BinNat.bosatsu

package Bosatsu/BinNat

from Bosatsu/Nat import Nat, Zero as NatZero, Succ as NatSucc, times2 as times2_Nat

export (BinNat(), toInt, toNat, toBinNat, next, add_BinNat, times2, div2,
  prev, times_BinNat, exp, cmp_BinNat, is_even, sub_BinNat, sub_Option, eq_BinNat)
# a natural number with three variants:
# Zero = 0
# Odd(n) = 2n + 1
# Even(n) = 2(n + 1)
# e.g:
# Zero, Odd(Zero), Even(Zero), Odd(Odd(Zero)), Even(Odd(Zero))
enum BinNat: Zero, Odd(half: BinNat), Even(half1: BinNat)

def is_even(b: BinNat) -> Bool:
  b matches Zero | Even(_)

# Convert a BinNat into the equivalent Int
# this is O(log(b)) operation
def toInt(b: BinNat) -> Int:
  recur b:
    Zero: 0
    Odd(n): toInt(n).times(2).add(1)
    Even(n): toInt(n).times(2).add(2)

# This is an O(b) operation
def toNat(b: BinNat) -> Nat:
  recur b:
    Zero: NatZero
    Odd(n): NatSucc(toNat(n).times2_Nat())
    Even(n): NatSucc(NatSucc(toNat(n).times2_Nat()))

# Convert a built in integer to a BinNat. <= 0 is converted to 0
def toBinNat(n: Int) -> BinNat:
  # build up a list in reverse of transformations
  fns = int_loop(n, [], \n, fns ->
    is_even = mod_Int(n, 2).eq_Int(0)
    (hfn, dec) = (n -> Even(n), n -> n.sub(1)) if is_even else (n -> Odd(n), n -> n)
    fns = [hfn, *fns]
    n = n.div(2)
    (dec(n), fns)
  )
  # Now apply all the transformations
  fns.foldLeft(Zero, (n, fn) -> fn(n))

def cmp_BinNat(a: BinNat, b: BinNat) -> Comparison:
  recur a:
    case Odd(a1):
      match b:
        case Odd(b1): cmp_BinNat(a1, b1)
        case Even(b1):
          # 2n + 1 <> 2m + 2
          # if n <= m, LT
          # if n > m GT
          match cmp_BinNat(a1, b1):
            case LT | EQ: LT
            case GT: GT
        case Zero: GT
    case Even(a1):
      match b:
        case Even(b1): cmp_BinNat(a1, b1)
        case Odd(b1):
          # 2n + 2 <> 2m + 1
          # if n >= m, GT
          # if n < m LT
          match cmp_BinNat(a1, b1):
            case GT | EQ: GT
            case LT: LT
        case Zero: GT
    case Zero:
      match b:
        case Odd(_) | Even(_): LT
        case Zero: EQ

# this is more efficient potentially than cmp_BinNat
# because at the first difference we can stop. In the worst
# case of equality, the cost is the same.
def eq_BinNat(a: BinNat, b: BinNat) -> Bool:
  recur a:
    case Zero: b matches Zero
    case Odd(n):
      match b:
        case Odd(m): eq_BinNat(n, m)
        case _: False
    case Even(n):
      match b:
        case Even(m): eq_BinNat(n, m)
        case _: False

# Return the next number
def next(b: BinNat) -> BinNat:
  recur b:
    Odd(half):
      # (2n + 1) + 1 = 2(n + 1)
      Even(half)
    Even(half1):
      # 2(n + 1) + 1
      Odd(next(half1))
    Zero: Odd(Zero)

# Return the previous number if the number is > 0, else return 0
def prev(b: BinNat) -> BinNat:
  recur b:
    case Zero | Odd(Zero): Zero
    case Odd(half):
      # (2n + 1) - 1 = 2n = 2(n-1 + 1)
      Even(prev(half))
    case Even(half1):
      # 2(n + 1) - 1 = 2n + 1
      Odd(half1)

def add_BinNat(left: BinNat, right: BinNat) -> BinNat:
  recur left:
    Odd(left) as odd:
      match right:
        Odd(right):
          # 2left + 1 + 2right + 1 = 2((left + right) + 1)
          Even(add_BinNat(left, right))
        Even(right):
          # 2left + 1 + 2(right + 1) = 2((left + right) + 1) + 1
          Odd(add_BinNat(left, right.next()))
        Zero: odd
    Even(left) as even:
      match right:
        Odd(right):
          # 2(left + 1) + 2right + 1 = 2((left + right) + 1) + 1
          Odd(add_BinNat(left, right.next()))
        Even(right):
          # 2(left + 1) + 2(right + 1) = 2((left + right + 1) + 1)
          Even(add_BinNat(left, right.next()))
        Zero: even
    Zero: right

# multiply by 2
def times2(b: BinNat) -> BinNat:
  recur b:
    Odd(n):
      #2(2n + 1) = Even(2n)
      Even(times2(n))
    Even(n):
      #2(2(n + 1)) = 2((2n + 1) + 1)
      Even(Odd(n))
    Zero: Zero

# 2n - 1 if it is defined
def doub_prev(b: BinNat) -> Option[BinNat]:
  match b:
    case Odd(n):
      # 2(2n + 1) - 1 = 4n + 1 = Odd(2n)
      Some(Odd(times2(n)))
    case Even(n):
      # 2(2n + 2) - 1 = 4n + 3 = 2(2n + 1) + 1
      Some(Odd(Odd(n)))
    case Zero: None

def sub_Option(left: BinNat, right: BinNat) -> Option[BinNat]:
  recur left:
    case Zero:
      match right:
        case Zero: Some(Zero)
        case _: None
    case Odd(left) as odd:
      match right:
        case Zero: Some(odd)
        case Odd(right):
          # (2n + 1) - (2m + 1) = 2(n - m)
          match sub_Option(left, right):
            case Some(n_m): Some(times2(n_m))
            case None: None
        case Even(right):
          # (2n + 1) - (2m + 2) = 2(n - m) - 1
          # note if (2n + 1) > (2m + 2), then n > m
          match sub_Option(left, right):
            case Some(n_m): doub_prev(n_m)
            case None: None
    case Even(left) as even:
      match right:
        case Zero: Some(even)
        case Odd(right):
          # Even can't equal odd, so we never return
          # zero. Next an even - odd is odd.
          # (2n + 2) - (2m + 1) = 2(n - m) + 1
          match sub_Option(left, right):
            case Some(n_m): Some(Odd(n_m))
            case None: None
        case Even(right):
          # (2n + 2) - (2m + 2) = 2(n - m)
          match sub_Option(left, right):
            case Some(n_m): Some(times2(n_m))
            case None: None

def sub_BinNat(left: BinNat, right: BinNat) -> BinNat:
  match sub_Option(left, right):
    case Some(v): v
    case None: Zero

def div2(b: BinNat) -> BinNat:
    match b:
        case Zero: Zero
        case Odd(n):
          # (2n + 1)/2 = n
          n
        case Even(n):
          # (2n + 2)/2 = n + 1
          next(n)

# multiply two BinNat together
def times_BinNat(left: BinNat, right: BinNat) -> BinNat:
  recur left:
    Zero: Zero
    Odd(left):
      match right:
        Zero: Zero
        right:
          # (2l + 1) * r = 2lr + r
          prod = times_BinNat(left, right)
          times2(prod).add_BinNat(right)
    Even(left):
      match right:
        Zero: Zero
        right:
          # 2(l + 1) * n = 2(l*n + n)
          prod = times_BinNat(left, right)
          times2(prod.add_BinNat(right))

one = Odd(Zero)

def exp(base: BinNat, power: BinNat) -> BinNat:
  recur power:
    case Zero: one
    case Odd(n):
      # b^(2n + 1) == (b^n) * (b^n) * b
      bn = exp(base, n)
      bn.times_BinNat(bn).times_BinNat(base)
    case Even(n):
      # b^(2n + 2) = (b^n * b)^2
      bn = exp(base, n)
      bn1 = bn.times_BinNat(base)
      bn1.times_BinNat(bn1)

# fold(fn, a, Zero) = a
# fold(fn, a, n) = fold(fn, fn(a, n - 1), n - 1)
def fold_left_BinNat(fn: (a, BinNat) -> a, init: a, cnt: BinNat) -> a:
  # use the Nat as a trick to prove termination
  def loop(init: a, cnt: BinNat, cnt_Nat: Nat) -> a:
    recur cnt_Nat:
      NatZero: init
      NatSucc(prevNat):
        cnt = prev(cnt)
        init = fn(init, cnt)
        loop(init, cnt, prevNat)

  # this is O(cnt) to build the Nat, but the fold is already
  # O(cnt), so this isn't a complexity change
  loop(init, cnt, toNat(cnt))

# fibonacci using the fuel pattern
def fib(b: BinNat) -> BinNat:
    def loop(n: Nat, cur: BinNat, next: BinNat) -> BinNat:
        recur n:
            NatZero: cur
            NatSucc(n):
                sum = add_BinNat(cur, next)
                loop(n, next, sum)
    one = Odd(Zero)
    loop(toNat(b), one, one)

def round_trip_law(i, msg):
  Assertion(i.toBinNat().toInt().eq_Int(i), msg)

def next_law(i, msg):
  Assertion(i.toBinNat().next().toInt().eq_Int(i.add(1)), msg)

def times2_law(i, msg):
  Assertion(i.toBinNat().times2().toInt().eq_Int(i.times(2)), msg)

two = one.next()
three = two.next()
four = three.next()


test = TestSuite(
  "BinNat tests", [
    Assertion(Zero.toInt().eq_Int(0), "0.toBinNat"),
    Assertion(Odd(Zero).toInt().eq_Int(1), "1.toBinNat"),
    Assertion(Even(Zero).toInt().eq_Int(2), "2.toBinNat"),
    Assertion(Odd(Odd(Zero)).toInt().eq_Int(3), "3.toBinNat"),
    Assertion(Even(Odd(Zero)).toInt().eq_Int(4), "4.toBinNat"),
    TestSuite("round trip laws", [ round_trip_law(i, m) for (i, m) in [
      (0, "roundtrip 0"),
      (1, "roundtrip 1"),
      (2, "roundtrip 2"),
      (3, "roundtrip 3"),
      (4, "roundtrip 4"),
      (5, "roundtrip 5"),
      (6, "roundtrip 6"),
      (7, "roundtrip 7"),
      (50, "roundtrip 50"),
      (61, "roundtrip 61"),
      (72, "roundtrip 72"),
    ]]),
    TestSuite("next law", [ next_law(i, msg) for (i, msg) in [
      (0, "0.next"),
      (5, "5.next"),
      (10, "10.next"),
      (113, "113.next"),
    ]]),
    Assertion(0.toBinNat().next().prev().toInt().eq_Int(0), "0.next().prev == 0"),
    Assertion(5.toBinNat().next().prev().toInt().eq_Int(5), "5.next().prev == 5"),
    Assertion(10.toBinNat().next().prev().toInt().eq_Int(10), "10.next().prev == 10"),
    Assertion(10.toBinNat().add_BinNat(11.toBinNat()).toInt().eq_Int(21), "add_BinNat(10, 11) == 21"),
    TestSuite("times2 law", [times2_law(i, msg) for (i, msg) in [
      (0, "0 * 2"),
      (1, "1 * 2"),
      (2, "2 * 2"),
      (5, "5 * 2"),
      (10, "10 * 2"),
    ]]),
    Assertion(10.toBinNat().times_BinNat(11.toBinNat()).toInt().eq_Int(110), "10*11 = 110"),
    Assertion(0.toBinNat().times_BinNat(11.toBinNat()).toInt().eq_Int(0), "0*11 = 0"),
    Assertion(fold_left_BinNat(\n, _ -> n.next(), Zero, 10.toBinNat()).toInt().eq_Int(10), "1 + ... + 1 = 10"),
    Assertion(fold_left_BinNat(\n1, n2 -> n1.add_BinNat(n2), Zero, 4.toBinNat()).toInt().eq_Int(6), "1+2+3=6"),
    Assertion(fib(Zero).toInt().eq_Int(1), "fib(0) == 1"),
    Assertion(fib(one).toInt().eq_Int(1), "fib(1) == 1"),
    Assertion(fib(two).toInt().eq_Int(2), "fib(2) == 2"),
    Assertion(fib(three).toInt().eq_Int(3), "fib(3) == 3"),
    Assertion(fib(four).toInt().eq_Int(5), "fib(4) == 5"),
    Assertion(cmp_BinNat(54.toBinNat(), 54.toBinNat()) matches EQ, "54 == 54"),
  ])