{-# LANGUAGE TypeFamilies #-}
module Stochastic.Generator where

import Control.Monad.State.Lazy

class Generator g where
  type From g
  nextG   :: State g (From g)
  nextN  :: Int -> State g [(From g)]
  nextN 0 = state $ \g0 -> ([], g0)
  nextN n =
    do
      x  <- nextG
      xs <- nextN (n-1)
      return (x:xs)

instance Generator [a] where
  type From [a] = a
  nextG = state $ \ g0 -> (head g0, tail g0)

foldWhile :: Generator g
         => (From g -> a -> a)
         -> a
         -> (a -> Bool)
         -> State g [From g]
foldWhile f z p =
  do
   x  <- nextG
   let y = f x z
   xs <- if (p y)
         then foldWhile f y p 
         else return []
   return (x:xs)

while :: Generator g => ((From g) -> Bool) -> State g [From g]
while p =
  do
   x  <- nextG
   xs <- if (p x)
         then while p
         else return []
   return (x:xs)