{-# LANGUAGE GADTs, LambdaCase #-}

module Test.Credit.Queue.Streams (Stream(..), SThunk, SLazyCon(..), smatch, credit, evalone, toList, ifIndirect, test) where

import Control.Monad
import Control.Monad.Credit

data Stream m a
  = SCons a (Stream m a)
  | SNil
  | SIndirect (SThunk m (Stream m a))

type SThunk m = Thunk m (SLazyCon m)

data SLazyCon m a where
  SAppend :: Stream m a -> Stream m a -> SLazyCon m (Stream m a)
  SReverse :: Stream m a -> Stream m a -> SLazyCon m (Stream m a)

instance MonadInherit m => HasStep (SLazyCon m) m where
  step :: forall a. SLazyCon m a -> m a
step (SAppend Stream m a
xs Stream m a
ys) = Stream m a -> Stream m a -> m (Stream m a)
forall (m :: * -> *) a.
MonadInherit m =>
Stream m a -> Stream m a -> m (Stream m a)
sappend Stream m a
xs Stream m a
ys
  step (SReverse Stream m a
xs Stream m a
ys) = Stream m a -> Stream m a -> m (Stream m a)
forall (m :: * -> *) a.
MonadInherit m =>
Stream m a -> Stream m a -> m (Stream m a)
sreverse Stream m a
xs Stream m a
ys

-- | Smart destructor for streams, consuming one credit
smatch :: MonadInherit m => Stream m a -- ^ Scrutinee
       -> (a -> Stream m a -> m b) -- ^ Cons case
       -> m b -- ^ Nil case
       -> m b
smatch :: forall (m :: * -> *) a b.
MonadInherit m =>
Stream m a -> (a -> Stream m a -> m b) -> m b -> m b
smatch Stream m a
x a -> Stream m a -> m b
cons m b
nil = m ()
forall (m :: * -> *). MonadCount m => m ()
tick m () -> m b -> m b
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Stream m a -> m b
eval Stream m a
x
  where
    eval :: Stream m a -> m b
eval Stream m a
x = case Stream m a
x of
      SCons a
a Stream m a
as -> a -> Stream m a -> m b
cons a
a Stream m a
as
      Stream m a
SNil -> m b
nil
      SIndirect SThunk m (Stream m a)
i -> SThunk m (Stream m a) -> m (Stream m a)
forall (m :: * -> *) (t :: * -> *) a.
(MonadLazy m, HasStep t m) =>
Thunk m t a -> m a
forall (t :: * -> *) a. HasStep t m => Thunk m t a -> m a
force SThunk m (Stream m a)
i m (Stream m a) -> (Stream m a -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Stream m a -> m b
eval

-- | delay a computation, consuming all credits
taildelay :: MonadInherit m => SLazyCon m (Stream m a) -> m (Stream m a)
taildelay :: forall (m :: * -> *) a.
MonadInherit m =>
SLazyCon m (Stream m a) -> m (Stream m a)
taildelay SLazyCon m (Stream m a)
t = do
  Thunk m (SLazyCon m) (Stream m a)
x <- SLazyCon m (Stream m a) -> m (Thunk m (SLazyCon m) (Stream m a))
forall (m :: * -> *) (t :: * -> *) a.
MonadLazy m =>
t a -> m (Thunk m t a)
forall (t :: * -> *) a. t a -> m (Thunk m t a)
delay SLazyCon m (Stream m a)
t
  Thunk m (SLazyCon m) (Stream m a) -> m ()
forall (m :: * -> *) (t :: * -> *) a.
MonadInherit m =>
Thunk m t a -> m ()
forall (t :: * -> *) a. Thunk m t a -> m ()
creditAllTo Thunk m (SLazyCon m) (Stream m a)
x
  Stream m a -> m (Stream m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Thunk m (SLazyCon m) (Stream m a) -> Stream m a
forall (m :: * -> *) a. SThunk m (Stream m a) -> Stream m a
SIndirect Thunk m (SLazyCon m) (Stream m a)
x)

sreverse :: MonadInherit m => Stream m a -> Stream m a -> m (Stream m a)
sreverse :: forall (m :: * -> *) a.
MonadInherit m =>
Stream m a -> Stream m a -> m (Stream m a)
sreverse Stream m a
xs Stream m a
ys = Stream m a
-> (a -> Stream m a -> m (Stream m a))
-> m (Stream m a)
-> m (Stream m a)
forall (m :: * -> *) a b.
MonadInherit m =>
Stream m a -> (a -> Stream m a -> m b) -> m b -> m b
smatch Stream m a
xs
  (\a
x Stream m a
xs -> SLazyCon m (Stream m a) -> m (Stream m a)
forall (m :: * -> *) a.
MonadInherit m =>
SLazyCon m (Stream m a) -> m (Stream m a)
taildelay (Stream m a -> Stream m a -> SLazyCon m (Stream m a)
forall (m :: * -> *) a.
Stream m a -> Stream m a -> SLazyCon m (Stream m a)
SReverse Stream m a
xs (a -> Stream m a -> Stream m a
forall (m :: * -> *) a. a -> Stream m a -> Stream m a
SCons a
x Stream m a
ys)))
  (Stream m a -> m (Stream m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stream m a
ys)

ifIndirect :: Monad m => Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
ifIndirect :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
ifIndirect (SIndirect SThunk m (Stream m a)
i) SThunk m (Stream m a) -> m ()
f = SThunk m (Stream m a) -> m ()
f SThunk m (Stream m a)
i
ifIndirect Stream m a
_ SThunk m (Stream m a) -> m ()
_ = () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

credit :: MonadInherit m => Stream m a -> m ()
credit :: forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
credit Stream m a
s = Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
ifIndirect Stream m a
s (SThunk m (Stream m a) -> Credit -> m ()
forall (m :: * -> *) (t :: * -> *) a.
MonadCredit m =>
Thunk m t a -> Credit -> m ()
forall (t :: * -> *) a. Thunk m t a -> Credit -> m ()
`creditWith` Credit
1)

evalone :: MonadInherit m => Stream m a -> m ()
evalone :: forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
evalone Stream m a
s = Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
forall (m :: * -> *) a.
Monad m =>
Stream m a -> (SThunk m (Stream m a) -> m ()) -> m ()
ifIndirect Stream m a
s (m (Stream m a) -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m (Stream m a) -> m ())
-> (SThunk m (Stream m a) -> m (Stream m a))
-> SThunk m (Stream m a)
-> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SThunk m (Stream m a) -> m (Stream m a)
forall (m :: * -> *) (t :: * -> *) a.
(MonadLazy m, HasStep t m) =>
Thunk m t a -> m a
forall (t :: * -> *) a. HasStep t m => Thunk m t a -> m a
force)

sappend :: MonadInherit m => Stream m a -> Stream m a -> m (Stream m a)
sappend :: forall (m :: * -> *) a.
MonadInherit m =>
Stream m a -> Stream m a -> m (Stream m a)
sappend Stream m a
xs Stream m a
ys = Stream m a -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
credit Stream m a
ys m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Stream m a -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
evalone Stream m a
ys m () -> m (Stream m a) -> m (Stream m a)
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Stream m a
-> (a -> Stream m a -> m (Stream m a))
-> m (Stream m a)
-> m (Stream m a)
forall (m :: * -> *) a b.
MonadInherit m =>
Stream m a -> (a -> Stream m a -> m b) -> m b -> m b
smatch Stream m a
xs
  (\a
x Stream m a
xs -> a -> Stream m a -> Stream m a
forall (m :: * -> *) a. a -> Stream m a -> Stream m a
SCons a
x (Stream m a -> Stream m a) -> m (Stream m a) -> m (Stream m a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SLazyCon m (Stream m a) -> m (Stream m a)
forall (m :: * -> *) a.
MonadInherit m =>
SLazyCon m (Stream m a) -> m (Stream m a)
taildelay (Stream m a -> Stream m a -> SLazyCon m (Stream m a)
forall (m :: * -> *) a.
Stream m a -> Stream m a -> SLazyCon m (Stream m a)
SAppend Stream m a
xs Stream m a
ys))
  (Stream m a -> m (Stream m a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Stream m a
ys)

walk :: Stream m a -> m ()
walk Stream m a
s = Stream m a -> (a -> Stream m a -> m ()) -> m () -> m ()
forall (m :: * -> *) a b.
MonadInherit m =>
Stream m a -> (a -> Stream m a -> m b) -> m b -> m b
smatch Stream m a
s (\a
_ Stream m a
xs -> Stream m a -> m ()
walk Stream m a
xs) (() -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

foo :: MonadInherit m => Stream m a -> m ()
foo :: forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
foo Stream m a
s = Stream m a -> (a -> Stream m a -> m ()) -> m () -> m ()
forall (m :: * -> *) a b.
MonadInherit m =>
Stream m a -> (a -> Stream m a -> m b) -> m b -> m b
smatch Stream m a
s (\a
_ Stream m a
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()) (() -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

test :: MonadInherit m => m ()
test :: forall (m :: * -> *). MonadInherit m => m ()
test = do
  Stream m Integer
s <- Stream m Integer -> Stream m Integer -> m (Stream m Integer)
forall (m :: * -> *) a.
MonadInherit m =>
Stream m a -> Stream m a -> m (Stream m a)
sappend (Integer -> Stream m Integer -> Stream m Integer
forall (m :: * -> *) a. a -> Stream m a -> Stream m a
SCons Integer
1 Stream m Integer
forall (m :: * -> *) a. Stream m a
SNil) (Integer -> Stream m Integer -> Stream m Integer
forall (m :: * -> *) a. a -> Stream m a -> Stream m a
SCons Integer
2 Stream m Integer
forall (m :: * -> *) a. Stream m a
SNil)
  Stream m Integer -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
credit Stream m Integer
s m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Stream m Integer -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
credit Stream m Integer
s
  Stream m Integer -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
foo Stream m Integer
s
  Stream m Integer -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
credit Stream m Integer
s
  Stream m Integer -> m ()
forall (m :: * -> *) a. MonadInherit m => Stream m a -> m ()
walk Stream m Integer
s

toList :: MonadLazy m => Stream m a -> m [a]
toList :: forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
SNil = [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
toList (SCons a
x Stream m a
xs) = do
  [a]
xs' <- Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
xs
  [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
xs'
toList (SIndirect SThunk m (Stream m a)
t) = do
  SThunk m (Stream m a)
-> (Stream m a -> m [a])
-> (SLazyCon m (Stream m a) -> m [a])
-> m [a]
forall (m :: * -> *) (t :: * -> *) a b.
MonadLazy m =>
Thunk m t a -> (a -> m b) -> (t a -> m b) -> m b
forall (t :: * -> *) a b.
Thunk m t a -> (a -> m b) -> (t a -> m b) -> m b
lazymatch SThunk m (Stream m a)
t Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList ((SLazyCon m (Stream m a) -> m [a]) -> m [a])
-> (SLazyCon m (Stream m a) -> m [a]) -> m [a]
forall a b. (a -> b) -> a -> b
$ \case
      SAppend Stream m a
xs Stream m a
ys -> do
        [a]
xs' <- Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
xs
        [a]
ys' <- Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
ys
        [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ [a]
[a]
xs' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
[a]
ys'
      SReverse Stream m a
xs Stream m a
ys -> do
        [a]
xs' <- Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
xs
        [a]
ys' <- Stream m a -> m [a]
forall (m :: * -> *) a. MonadLazy m => Stream m a -> m [a]
toList Stream m a
ys
        [a] -> m [a]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([a] -> m [a]) -> [a] -> m [a]
forall a b. (a -> b) -> a -> b
$ [a] -> [a]
forall a. [a] -> [a]
reverse [a]
[a]
xs' [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
[a]
ys'

instance (MonadMemory m, MemoryCell m a) => MemoryCell m (SLazyCon m a) where
  prettyCell :: SLazyCon m a -> m Memory
prettyCell (SAppend Stream m a
xs Stream m a
ys) = do
    Memory
xs' <- Stream m a -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Stream m a
xs
    Memory
ys' <- Stream m a -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Stream m a
ys
    Memory -> m Memory
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Memory -> m Memory) -> Memory -> m Memory
forall a b. (a -> b) -> a -> b
$ String -> [Memory] -> Memory
mkMCell String
"SAppend" [Memory
xs', Memory
ys']
  prettyCell (SReverse Stream m a
xs Stream m a
ys) = do
    Memory
xs' <- Stream m a -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Stream m a
xs
    Memory
ys' <- Stream m a -> m Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell Stream m a
ys
    Memory -> m Memory
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Memory -> m Memory) -> Memory -> m Memory
forall a b. (a -> b) -> a -> b
$ String -> [Memory] -> Memory
mkMCell String
"SReverse" [Memory
xs', Memory
ys']

instance (MonadMemory m, MemoryCell m a) => MemoryCell m (Stream m a) where
  prettyCell :: Stream m a -> m Memory
prettyCell Stream m a
xs = [Memory] -> Maybe Memory -> Memory
mkMList ([Memory] -> Maybe Memory -> Memory)
-> m [Memory] -> m (Maybe Memory -> Memory)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Stream m a -> m [Memory]
forall {f :: * -> *} {a} {m :: * -> *}.
MemoryCell f a =>
Stream m a -> f [Memory]
toList Stream m a
xs m (Maybe Memory -> Memory) -> m (Maybe Memory) -> m Memory
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Stream m a -> m (Maybe Memory)
forall {f :: * -> *} {m :: * -> *} {a}.
MemoryCell f (SThunk m (Stream m a)) =>
Stream m a -> f (Maybe Memory)
toHole Stream m a
xs
    where
      toList :: Stream m a -> f [Memory]
toList Stream m a
SNil = [Memory] -> f [Memory]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Memory] -> f [Memory]) -> [Memory] -> f [Memory]
forall a b. (a -> b) -> a -> b
$ []
      toList (SCons a
x Stream m a
xs) = (:) (Memory -> [Memory] -> [Memory])
-> f Memory -> f ([Memory] -> [Memory])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell a
x f ([Memory] -> [Memory]) -> f [Memory] -> f [Memory]
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Stream m a -> f [Memory]
toList Stream m a
xs
      toList (SIndirect SThunk m (Stream m a)
t) = [Memory] -> f [Memory]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Memory] -> f [Memory]) -> [Memory] -> f [Memory]
forall a b. (a -> b) -> a -> b
$ []

      toHole :: Stream m a -> f (Maybe Memory)
toHole Stream m a
SNil = Maybe Memory -> f (Maybe Memory)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Memory -> f (Maybe Memory))
-> Maybe Memory -> f (Maybe Memory)
forall a b. (a -> b) -> a -> b
$ Maybe Memory
forall a. Maybe a
Nothing
      toHole (SCons a
x Stream m a
xs) = Stream m a -> f (Maybe Memory)
toHole Stream m a
xs
      toHole (SIndirect SThunk m (Stream m a)
t) = Memory -> Maybe Memory
forall a. a -> Maybe a
Just (Memory -> Maybe Memory) -> f Memory -> f (Maybe Memory)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SThunk m (Stream m a) -> f Memory
forall (m :: * -> *) a. MemoryCell m a => a -> m Memory
prettyCell SThunk m (Stream m a)
t