Rand.bosatsu

package Bosatsu/Rand

from Bosatsu/Nat import (Nat, Zero, Succ, exp as exp_Nat,
  to_Int as nat_to_Int, to_Nat as int_to_Nat)

from Bosatsu/BinNat import (BinNat, Zero as BZero, Odd, Even,
  next as next_BinNat, prev as prev_BinNat, toInt as binNat_to_Int)

export (Rand, run_Rand, prod_Rand, map_Rand, flat_map_Rand, const_Rand,
  int_range, sequence_Rand, bool_Rand, geometric_Int, from_pair, one_of)

struct State(s0: Int, s1: Int, s2: Int, s3: Int)
struct UInt64(toInt: Int)

def operator -(a, b): a.sub(b)
def operator +(a, b): a.add(b)
def operator *(a, b): a.times(b)
def operator &(a, b): a.and_Int(b)
def operator <<(a, b): a.shift_left_Int(b)
def operator >>(a, b): a.shift_right_Int(b)
def operator |(a, b): a.or_Int(b)
def operator ^(a, b): a.xor_Int(b)

bitmask_64 = (1 << 64) - 1

def uint64(i): i & bitmask_64

def rotl(x, k):
  lx = uint64(x << k)
  rx = uint64(x >> (64 - k))
  lx | rx

# this is the Xoshiro256** algorithm
def next(state: State) -> (State, UInt64):
  State { s0, s1, s2, s3 } = state
  result = uint64(rotl(uint64(s1 * 5), 7) * 9)
  t = uint64(s1 << 17)
  s2 = s2 ^ s0
  s3 = s3 ^ s1
  s1 = s1 ^ s2
  s0 = s0 ^ s3

  s2 = s2 ^ t
  s3 = uint64(rotl(s3, 45))

  (State { s0, s1, s2, s3}, UInt64(result))

# this is taken from the git-shas of the two previous commits to main
default_state = 0x78a2_951d_9698_3f3f_8ff9_e45e_c217_8773_4e55_075f_bd57_0ae8_c3d4_d351_f3bd_9bfd

def state_from_Int(i: Int) -> State:
  not_zero = match (i ^ default_state):
      case 0: default_state
      case nz: nz

  State(uint64(not_zero), uint64(not_zero >> 64), uint64(not_zero >> 128), uint64(not_zero >> 192))

struct Rand[a: +*](fn: State -> (State, a))

def map_Rand[a, b](r: Rand[a], fn: a -> b) -> Rand[b]:
  Rand(fna) = r
  Rand(s -> (
    (s1, a) = fna(s)
    (s1, fn(a))
  ))

def flat_map_Rand[a, b](r: Rand[a], fn: a -> Rand[b]) -> Rand[b]:
  Rand(fna) = r
  Rand(s -> (
    (s1, a) = fna(s)
    Rand(fnb) = fn(a)
    fnb(s1)
  ))

def prod_Rand[a, b](ra: Rand[a], rb: Rand[b]) -> Rand[(a, b)]:
  Rand(fna) = ra
  Rand(fnb) = rb
  Rand(s0 -> (
    (s1, a) = fna(s0)
    (s2, b) = fnb(s1)
    (s2, (a, b))
  ))

def const_Rand[a](a: a) -> Rand[a]: Rand(s -> (s, a))

def xor(a, b):
    match a:
        case True: False if b else True
        case False: b

nat_2 = Succ(Succ(Zero))

def parity(n: Nat, i: Int) -> Bool:
    recur n:
        case Zero: (i & 1) matches 1
        case Succ(p):
            half_bits = exp_Nat(nat_2, n).nat_to_Int()
            left_half = i >> half_bits
            bl = parity(p, left_half)
            br = parity(p, i)
            xor(bl, br)

# 2^6 = 64
zero = Zero
succ = Succ
six = zero.succ().succ().succ().succ().succ().succ()

bool_Rand: Rand[Bool] = Rand(s -> (
  (s, UInt64(i)) = next(s)
  (s, parity(six, i))
))

def run_Rand[a](rand: Rand[a], seed: Int) -> a:
  Rand(fn) = rand
  (_, a) = fn(state_from_Int(seed))
  a

def sequence_Rand[a](rands: List[Rand[a]]) -> Rand[List[a]]:
  def sample(rands: List[Rand[a]], s: State, acc: List[a]) -> (State, List[a]):
    recur rands:
        case []: (s, reverse(acc))
        case [Rand(hfn), *rt]: 
          (s1, h) = hfn(s)
          sample(rt, s1, [h, *acc])


  Rand(s -> sample(rands, s, []))

uint64_Rand = Rand(next)

def bit_count(i: Int) -> Int:
  int_loop(i, 0, (i, bits) -> (i >> 1, bits + 1))

def to_big_Int(us: List[UInt64], acc: Int) -> Int:
  recur us:
    case []: acc
    case [UInt64(h), *t]: to_big_Int(t, (acc << 64) | h)

nat30 = int_to_Nat(30)

def resample(rand_Int: Rand[Int], high: Int, uints: Int) -> Rand[Int]:
  Rand(fn) = rand_Int
  boundary = (1 << (uints * 64)).div(high) * high
  def next(s: State, fuel: Nat) -> (State, Int):
    recur fuel:
        case Zero: (s, (high - 1).div(2))
        case Succ(n):
          (s1, i) = fn(s)
          if i.cmp_Int(boundary) matches LT:
            # this sample worked
            (s1, i.mod_Int(high))
          else: next(s1, n)
  
  # each sample has a change > 1/2 of working, so if we try 30 times
  # we have a 1 in a billion chance to just choose the mean value
  Rand(s -> next(s, nat30))

const0 = const_Rand(0)

# if you pass high <= 0, you get const_Rand(0)
def int_range(high: Int) -> Rand[Int]:
  if high.cmp_Int(1) matches GT:
      # high >= 2
      # bits > 1 since high > 0
      bits = bit_count(high)
      uint_count = bits.div(64) + 1
      rand_Int = replicate_List(uint64_Rand, uint_count) \
        .sequence_Rand() \
        .map_Rand(us -> to_big_Int(us, 0))
      # now we know the integer we get out is >= high
      resample(rand_Int, high, uint_count)
  else: const0

def geometric(depth: Nat, acc: Int) -> Rand[Int]:
    recur depth:
        case Zero: const_Rand(acc)
        case Succ(prev):
          # prob 1/2 0, prob 1/2 geometric + 1
          bool_Rand.flat_map_Rand(b -> (
              match b:
                  case True: const_Rand(acc)
                  case False: geometric(prev, acc + 1)
          ))

geometric_Int: Rand[Int] = geometric(nat30, 0)

def len[a](list: List[a]) -> BinNat:
    def loop(list, acc):
        recur list:
            case []: acc
            case [_, *t]: loop(t, acc.next_BinNat())
    loop(list, BZero)

def split_at[a](list: List[a], idx: BinNat) -> (List[a], List[a]):
  recur list:
    case []: ([], [])
    case [h, *t]:
      match idx:
        case BZero: ([], list)
        case _:
          (left, right) = split_at(t, idx.prev_BinNat())
          ([h, *left], right)

def from_pair(left: Rand[a], right: Rand[a]) -> Rand[a]:
    bool_Rand.flat_map_Rand(b -> (
      match b:
        case True: left
        case False: right
    ))

def one_of[a](head: Rand[a], tail: List[Rand[a]]) -> Rand[a]:
    unreachable_case = head
    # invariant: items_len == len(items), and items_len > 0
    def loop(items_len: BinNat, items: List[Rand[a]]) -> Rand[a]:
        recur items_len:
            case BZero:
                # unreachable, just return head to typecheck
                unreachable_case
            case Odd(BZero):
                # 2 * 0 + 1, just choose the head of the list
                match items:
                    case [h, *_]: h
                    case _:
                      # unreacheable len > 0, so can't match
                      unreachable_case
            case Odd(n):
                # llen = 2n + 1
                # split into 1 + n + n 
                # with weight (n + 1) choose left, and with weight n choose right
                match items:
                    case [front, *tail]:
                        (left, right) = split_at(tail, n)
                        lrand = loop(n, left)
                        rrand = loop(n, right)
                        # equal chance left and right
                        back = from_pair(lrand, rrand)
                        int_range(binNat_to_Int(items_len)) \
                          .flat_map_Rand(idx -> (
                            match idx:
                              case 0: front
                              case _: back
                          ))
                    case []: unreachable_case
            case Even(BZero):
              # exactly two, choose each equally likely
              match items:
                  case [left, right, *_]:
                    # equal chance left and right
                    from_pair(left, right)
                  case _:
                    unreachable_case
            case Even(n):
              # 2n + 2, pull two off, choose them equally likely,
              match items:
                  case [f1, f2, *tail]:
                    front = from_pair(f1, f2)
                    back = loop(n, tail)
                    # with prob 1/(n + 1) choose front
                    int_range(binNat_to_Int(items_len).div(2)) \
                      .flat_map_Rand(idx -> (
                        match idx:
                          case 0: front
                          case _: back
                      ))
                  case _:
                    unreachable_case

    items = [head, *tail]
    loop(len(items), items)