StreamTest.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# OPTIONS_GHC -ddump-simpl #-}
module StreamTest where

import Language.Haskell.TH
import Data.Functor.Compose
import Control.Monad.Trans.Identity
import Control.Monad.Trans
import Control.Monad.Trans.Free.Church

import Language.Haskell.TH
import Control.Monad.Trans.Writer (WriterT(..), runWriterT, tell)
import Unsafe.Coerce

-- Minimal instance
instance Quote m => Quote (WriterT [String] m) where
  newName s = lift (newName s)

-- Example Code
example :: Code (WriterT [String] Q) Int
example = [|| 1 + 2 + 3 ||]

--example2 :: Code (CT Q) In/
--example2 = [|| 0 ||]

example3 :: Code (CT CodeQ) Int
example3 = [|| 0 ||]

lifted = tell ["writer"] `bindCode_` example

-- Runner
runExample :: Q (Exp, [String])
runExample = runWriterT (unTypeCode lifted)


data Stream a r = Yield a (Stream a r)
                | Done r

s1 = Yield 0 (Yield 1 (Done 10))

data Stream2 a r = Stream2 { runStream :: forall z . (a -> z -> z) -> (r -> z) -> z }

s2 = Stream2 $ \y d -> y 0 $ y 1 $ d 10

y1 a (Stream2 k) = Stream2 $ \y d -> y a (k y d)
d1 r = Stream2 $ \_ d -> d r

s3 0 = d1 0
s3 n = y1 n (s3 (n - 1))

sink (Stream2 k) = k (:) (const [])

res = sink (s3 10)

data Stream3 a r = Stream3 { runStreamC :: forall z . (CodeQ a -> CodeQ z -> CodeQ z) -> (CodeQ r -> CodeQ z) -> CodeQ z }

--s3' = Stream3 $ \y d -> y 0 $ y 1 $ d 10

y3 a (Stream3 k) = Stream3 $ \y d -> y a (k y d)
d3 r = Stream3 $ \_ d -> d r

s4 n = Stream3 $ \y d ->
  [|| let go 0 = $$(d [|| 0 ||])
          go n = $$(y [|| n ||] [|| (go (n-1)) ||])
      in go $$n
  ||]


sink3 (Stream3 k) = k (\x y -> [|| $$x : $$y ||]) (\_ -> [|| [] ||])

sink4 (Stream3 k) = k (\x y -> [|| $$x + $$y ||]) (\r -> r)

res2 :: CodeQ [Int]
res2 = sink3 (s4 [|| 10 ||])

newtype CT m a = C { runC :: forall r . (a -> m r) -> m r } deriving (Functor)

instance Applicative (CT m) where
    pure a = C $ \k -> k a
    (<*>) gf ga = C $ \k -> runC gf $ \f -> runC ga $ \a -> k (f a)

instance Monad (CT m) where
    return = pure
    (>>=) ga f = C $ \k -> runC ga $ \a -> runC (f a) k


type C = CT CodeQ

type CM m = CT (Compose CodeQ m)




r :: C (CodeQ a) -> CodeQ a
r (C k) = k id

r2 :: CT m (m a) -> m a
r2 (C k) = k id

r3 :: Applicative m => CM m (CodeQ a) -> CodeQ (m a)
r3 = getCompose . ($ Compose) . runC . fmap p

r3' :: Applicative m => CM m (CodeQ (m a)) -> CodeQ (m a)
r3' = getCompose . ($ Compose) . runC

r4 :: Applicative m => CM m (TExp (m a)) -> Q (TExp (m a))
r4 (C k) = examineCode $ getCompose $ k (Compose . Code . pure)

p :: Applicative m => CodeQ a -> CodeQ (m a)
p a = [|| pure $$a ||]


data St a r = Y (CodeQ a) (St a r) | D (CodeQ r)

letrec :: (CodeQ a -> CodeQ a) -> C (CodeQ a)
letrec a = C $ \k -> [|| let foo = $$(a [|| foo ||]) in $$(k [|| foo ||]) ||]

down1 :: (CodeQ a -> CodeQ b) -> CodeQ (a -> b)
down1 f = [|| \a -> $$(f [|| a ||]) ||]


data Stream4 a m r =
      Stream4 { runStreamC3 :: forall z . (CodeQ a -> CodeQ (m z) -> CodeQ (m z))
                           -> (CodeQ r -> CodeQ (m z))
                           -> CodeQ (m z) }

data Stream5 a m base r = Stream5 { runStreamC4 :: forall z . (CodeQ a -> m (CM base) z -> m (CM base) z)
                                                     -> (CodeQ r -> m (CM base) z)
                                                     -> m (CM base) z }

y5 :: CodeQ a -> Stream5 a m base r -> Stream5 a m base r
y5 x r = Stream5 $ \y d -> y x (runStreamC4 r y d)

d5 :: CodeQ r -> Stream5 a m base r
d5 x = Stream5 $ \y d -> d x

iostream :: Stream5 Int IdentityT IO ()
iostream = y5 [|| 0 ||] $ y5 [|| 1 ||] $ d5 [|| () ||]

gen :: CodeQ a -> CM m (CodeQ a)
gen x = C $ \k -> Compose $ [|| do
                            let g = $$x
                            $$(getCompose $ k [|| g ||]) ||]

genM :: Monad m => CodeQ (m a) -> CM m (CodeQ a)
genM x = C $ \k -> Compose $ [|| do
                            g <- $$x
                            $$(getCompose $ k [|| g ||]) ||]


--liftBase :: CM base (CodeQ a) -> Stream5 z m base a
--liftBase b

sinkStream5 :: forall n a m r . (MonadTrans n) => Stream5 a n m r -> n (CM m) (CodeQ [a])
sinkStream5 (Stream5 k) = k ly ld
  where
    ld _ = lift $ pure [|| [] ||]

    ly :: CodeQ a -> n (CM m) (CodeQ [a]) -> n (CM m) (CodeQ [a])
    ly y n = do
      y' <- lift $ gen y
      n' <- n
      lift $ pure $ [|| $$y' : $$n' ||]

data YieldF a r = YieldF a r

type Stream6 a m base = FT (YieldF (CodeQ a)) (m (CM base))

sinkStream6 :: MonadTrans n => Stream6 a n m r -> n (CM m) (CodeQ [a])
sinkStream6 (FT k) = k ld ly
  where
    ld _ = lift $ pure [|| [] ||]

    ly k1 (YieldF a r) = do
      r' <- k1 r
      lift $ pure $ [|| $$a : $$r' ||]

sinkStream7 :: MonadTrans n => Stream6 a n m r -> n (CM m) r
sinkStream7 (FT k) = k ld ly
  where
    ld x = pure x

    ly k1 (YieldF a r) = k1 r

sinkStream8 :: (Show a, MonadTrans n) => Stream6 a n IO (CodeQ (IO r)) -> n (CM IO) (CodeQ (IO r))
sinkStream8 (FT k) = k ld ly
  where
    ld x = pure x

    ly k1 (YieldF a r) = do
      r' <- k1 r
      lift $ pure [|| print $$a >> $$r' ||]



-- MonadTrans lift
l :: Monad (n (CM m)) => n (CM m) r -> Stream6 a n m r
l m = FT $ \r _ -> m >>= r

lc :: (MonadTrans n, Monad m) => CodeQ (m r) -> Stream6 a n m (CodeQ r)
lc v = l $ lift (genM v)

y6 :: CodeQ a -> Stream6 a m base r -> Stream6 a m base r
y6 a r = wrap (YieldF a r)

stream6 :: Stream6 Int IdentityT IO ()
stream6 = y6 [|| 0 ||] (y6 [|| 1 ||] (pure ()))
--
stream6IO :: Stream6 Int IdentityT IO ()
stream6IO = lc [|| print "jimney" ||] >> stream6

--toCode :: Stream6 Int IdentityT IO (CodeQ Int) -> Code (Stream

  {-
countdown :: CodeQ Int -> Stream6 Int IdentityT IO (CodeQ Int)
countdown n = FT $ \r eff -> dfo
   _ $  [|| let go 0 = $$(r [|| 0 ||])
                go n = $$(eff _)
            in go $$n ||]
            -}


--test :: IdentityT C (CodeQ a) -> Code _ a
--test (IdentityT (C r)) = Code _

convert :: C a -> CT (Compose Q TExp) a
convert = hoist (liftCode . getCompose) (Compose . examineCode)

hoist :: (forall a . n a -> m a) -> (forall a . m a -> n a) -> CT m a -> CT n a
hoist f t (C m) = C $ \k -> t (m (f . k))

instance MonadTrans CT where
  lift m = C $ \k -> m >>= k

instance Quote (CT CodeQ) where
  newName s = C$ \k -> newName s `bindCode` k

instance Quote (CT (Compose CodeQ base)) where
  newName s = C$ \k -> Compose $ newName s `bindCode` (getCompose .  k)

instance Quote m => Quote (FT f m) where
  newName s = lift (newName s)

instance Quote m => Quote (IdentityT m) where
  newName s = lift (newName s)

--instance Quote (Compose CodeQ n) where
--  newName s = Compose $ _ (newName s)

qt :: Code (Stream6 a IdentityT c) Int
qt = [|| 0 ||]

filip :: Code (Stream6 a IdentityT c) v -> Stream6 a IdentityT c (CodeQ v)
filip x = do
  r <- examineCode x
  pure $ Code (pure r)

flop :: Stream6 a IdentityT c (CodeQ v) -> Code (Stream6 a IdentityT c) v
flop s = Code $ do
  Code c' <- s
  lift $ lift $ C $ \k -> Compose $ Code $ c' >>= examineCode . getCompose . k

testing :: Code (Stream6 Int IdentityT c) Int
testing = [|| $$(flop $ y6 [|| 0 ||] (pure [|| 1 ||])) + $$(flop $ y6 [|| 1 ||] (pure [|| 2 ||])) ||]

countdown :: Code (Stream6 Int IdentityT c) Int -> Code (Stream6 Int IdentityT c) (IO Int)
countdown x = [|| let go :: Int -> IO Int
                      go 0 = $$(flop $ pure [|| pure 0 ||])
                      go n = $$(flop $ y6 [|| n ||] (filip $ [|| go (n - 1) ||]))
                  in go $$x ||]

sink6 :: Stream6 Int IdentityT IO (CodeQ (IO r)) -> CodeQ (IO r)
sink6 = r3' . runIdentityT . sinkStream8

countdown2:: Code Q Int -> Code Q (IO Int)
countdown2 x = [|| let go :: Int -> IO Int
                       go 0 = $$(sink6 $ pure [|| pure 0 ||])
                       go n = $$(sink6 $ y6 [|| n ||] (filip $ [|| go (n - 1) ||]))
                  in   go $$x ||]

r10 :: Applicative c => FT (YieldF (CodeQ Int)) (IdentityT (CM c)) (TExp (c a))
  -> Q (TExp (c a))
r10 (FT k) = r4 $ runIdentityT $ k ly ld
  where
    ly :: TExp a -> IdentityT (CM c) (TExp a)
    ly x = lift $ pure x

    ld :: (x -> IdentityT (CM c) (TExp a)) -> YieldF (CodeQ Int) x -> IdentityT (CM c) (TExp a)
    ld k1 (YieldF a r) = fmap texp [| print 0 >> $(fmap unType $ k1 r)  |]

texp :: Exp -> TExp a
texp = unsafeCoerce

runCS :: Applicative c => Code (Stream6 Int IdentityT c) (c Int) -> Code Q (c Int)
runCS (Code n) = Code (r10 n)

hc :: Monad m => (forall a . m (TExp a) -> n (TExp a)) -> Code m a -> Code n a
hc f (Code n) = Code (f n)

--Q (TExp a)