Fancier `Gen`

gen.u

unique type test.Result = Fail Failure | Ok Text

unique ability Weighted a where
  bump : Nat -> ()
  give : a -> ()

unique ability Gen where
  generate : '{Weighted a} () -> a

Weighted.map : (a ->{g} b) -> '{Weighted a, g} r -> '{Weighted b,g} r 
Weighted.map f w = do
  go = cases
    { Weighted.bump n -> k } ->
       Weighted.bump n 
       handle !k with go
    { Weighted.give a -> k } -> 
       Weighted.give (f a)
       handle !k with go
    { r } -> r
  handle !w with go

Weighted.join : '{Weighted ('{Weighted a,g} ()),g} () -> '{Weighted a,g} ()
Weighted.join outer = do 
  go : Map Nat (Either ('{Weighted a,g} ()) ('{Weighted ('{Weighted a,g} ()), g} ()))
     ->{Weighted a, g} ()
  go heap = match Map.breakOffMin heap with
    None -> ()
    Some ((cost, Left wa), heap) -> 
      handle !wa with go.inner heap cost (Optional.fold 'maxNat at1 (Map.getMin heap))
    Some ((cost, Right ww), heap) -> 
      handle !ww with go.outer heap cost (Optional.fold 'maxNat at1 (Map.getMin heap))

  go.outer : Map Nat (Either ('{Weighted a,g} ()) ('{Weighted ('{Weighted a,g} ()), g} ()))
        -> Nat
        -> Nat 
        -> Request {Weighted ('{Weighted a,g} ())} ()
        ->{Weighted a, g} ()
  go.outer heap curCost suspendMark = cases
    { _ } -> go heap
    { Weighted.give w -> k } -> 
      handle !k with go.outer (Map.insert curCost (Left w) heap) curCost suspendMark
    { Weighted.bump by -> k } -> 
      cost' = curCost + by
      if cost' > suspendMark then 
        go (Map.insert cost' (Right k) heap) 
      else
        handle !k with go.outer heap cost' suspendMark

  go.inner : Map Nat (Either ('{Weighted a,g} ()) ('{Weighted ('{Weighted a,g} ()), g} ()))
        -> Nat
        -> Nat 
        -> Request {Weighted a} ()
        ->{Weighted a, g} ()
  go.inner heap curCost suspendMark = cases
    { _ } -> go heap
    { Weighted.give a -> k } -> 
      Weighted.give a
      handle !k with go.inner heap curCost suspendMark
    { Weighted.bump by -> k } -> 
      cost' = curCost + by
      if cost' > suspendMark then 
        go (Map.insert cost' (Left k) heap) 
      else
        handle !k with go.inner heap cost' suspendMark
  
  heap0 = Map.singleton 0 (Right outer)
  go heap0 
  
Weighted.flatMap : (a ->{g, Weighted b} ()) -> '{g, Weighted a} () -> '{g, Weighted b} ()
Weighted.flatMap f w = 
  mapped = Weighted.map (a -> '(f a)) w
  Weighted.join mapped

Weighted.toStream : '{g, Weighted a} r -> {g, Stream a} r
Weighted.toStream w = let
  go = cases
    { r } -> r
    { Weighted.give a -> k } -> 
      Stream.emit a
      handle !k with go
    { Weighted.bump by -> k } -> 
      handle !k with go
  handle !w with go

Gen.take : Nat -> '{Gen, g} r ->{g} [r]
Gen.take n g = Stream.toList! '(Stream.take! n '(Weighted.toStream '(toWeighted g)))

-- Question: how will this get used alongside other testing abilities?
-- Consideration: Generate pairs of lists of nats in range 0 to 100
--   Probably don't want to generate them in order or whatever
--   ; switch to random generation?

Gen.toWeighted : '{Gen, g} r ->{g, Weighted r} ()
Gen.toWeighted g =
  go : Request {Gen} x ->{Weighted x} ()
  go = cases
    { a } -> give a
    { Gen.generate xs -> k } -> 
      Weighted.flatMap (a -> handle k a with go) xs () 
  handle !g with go

---
unique ability Labeled where
  labels : [Text]

unique ability Gen where
  generate: '{Weighted a}'{Random} a -> {Gen} a

unique ability Generator a where
  next : {Generator} a

Weighted.nats = do
  go n = 
    Weighted.give n
    Weighted.bump 1
    go (lfsr 1)
  Weighted.give 0
  Weighted.bump 1
  Weighted.give maxNat
  Weighted.bump 1
  go 1