PStream.hs

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE PatternSynonyms #-}

import Control.Applicative (Applicative(..))
import Control.Monad (ap)

import Data.Function (fix)


-- | Original AST, as described in "Functional Programming with Structured
-- Graphs". It's a PHOAS encoding of stream with cycle, with the abstract
-- type 'b' and the element type 'x'.
data PS0 x b = Var0 b
             | Mu0 (b -> PS0 x b) -- recursive binder
             | Cons0 x  (PS0 x b)

-- | Closed terms:
newtype Stream0 x = Stream0 { runS0 :: forall b. PS0 x b }

-- | Quasi-monadic join:
joinPS0 :: PS0 x (PS0 x b) -> PS0 x b
joinPS0 (Var0 x) = x
joinPS0 (Mu0 h) = Mu0 $ joinPS0 . h . Var0
joinPS0 (Cons0 x xs) = Cons0 x $ joinPS0 xs

-- | Unroll one layer of recursive binding : Mu $ \x -> E[x] => E[Mu $ \x -> E[x]]
unrollPS0 :: Stream0 x -> Stream0 x
unrollPS0 s = Stream0 $ joinPS0 . go $ runS0 s
  where
    go (Mu0 g) = g . joinPS0 . Mu0 $ g
    go (Cons0 x xs) = Cons0 x $ go xs
    go e = e

-- | Simple acyclic fold (don't unroll to infinite lists)
toListPS0 :: Stream0 x -> [x]
toListPS0 = go . runS0
  where
    go (Var0 x) = x
    go (Mu0 h) = go $ h [] -- nil
    go (Cons0 x xs) =  x : go xs

-- | Cyclic fold (unroll to infinite lists)
toListRecPS0 :: Stream0 x -> [x]
toListRecPS0 = go . runS0
  where
    go (Var0 x) = x
    go (Mu0 h) = fix $ go . h -- fixpoint
    go (Cons0 x xs) = x : go xs

-- | Some examples
exPS0, exPS0' :: Stream0 Int
exPS0 = Stream0 $ Cons0 0 (Mu0 $ \x -> Cons0 1 (Cons0 2 (Var0 x)))
exPS0' = unrollPS0 exPS0


-- -----------------------------------------------------------------------------
--               Monadic PHOAS
-- -----------------------------------------------------------------------------

-- | Free Monad in Fegaras and Sheard style
data Rec p a b = Place b
               | Roll (p a (Rec p a b))

-- | It's a profunctor, and thus a Functor over the positive 'b'
instance Functor (p a) => Functor (Rec p a) where
  fmap f (Place b) = Place $ f b
  fmap f (Roll p) = Roll $ fmap (fmap f) p

instance Functor (p a) => Applicative (Rec p a) where
  pure = Place
  (<*>) = ap

instance Functor (p a) => Monad (Rec p a) where
  return = Place
  (Place b) >>= f = f b
  (Roll p)  >>= f = Roll $ fmap (>>= f) p



-- "linearized" PS0 functor :
data PSF x a b = MuF (a -> b)
               | ConsF x   b

-- | Functor over 'b'
instance Functor (PSF x a) where
  fmap f (ConsF x xs) = ConsF x $ f xs
  fmap f (MuF h) = MuF $ f . h

-- | Type and pattern synonyms
type PS1 x = Rec (PSF x)
pattern Var1 x = Place x
pattern Mu1 h = Roll (MuF h)
pattern Cons1 x xs = Roll (ConsF x xs)

-- | Closed terms
newtype Stream1 x = Stream1 { runS1 :: forall b. PS1 x b b }



-- Literate translation of 'joinPS0'. The type annotation is inferred.
joinPS1 :: PS1 x (PS1 x' a a') (PS1 x a' b) -> PS1 x a' b
joinPS1 (Var1 x) = x
joinPS1 (Mu1 h) = Mu1 $ joinPS1 . h . Var1 -- Var1 affects the negative occurrences
joinPS1 (Cons1 x xs) = Cons1 x $ joinPS1 xs
joinPS1 _ = undefined --  make the exhaustivity checker quiet with PSs

-- Some specializations that typecheck in 'unrollPS1':
-- joinPS1 :: PS1 x (PS1 x  a a') (PS1 x a' b) -> PS1 x a' b   -- x' ~ x
-- joinPS1 :: PS1 x (PS1 x' a b ) (PS1 x b  b) -> PS1 x b  b   -- a' ~ b
-- joinPS1 :: PS1 x (PS1 x' a a ) (PS1 x a  b) -> PS1 x a  b   -- a' ~ a
-- joinPS1 :: PS1 x (PS1 x  b b ) (PS1 x b  b) -> PS1 x b  b   -- a' ~ a ~ b, x' ~ x


-- Inlining '(>>=)' and 'fmap' in '(>>= id)' give us :
joinFreePSF :: PS1 x a (PS1 x a b) -> PS1 x a b -- inferred
joinFreePSF (Var1 b) = b
joinFreePSF (Mu1 h) = Mu1 $ joinFreePSF . h
joinFreePSF (Cons1 x xs) = Cons1 x $ joinFreePSF xs
joinFreePSF _ = undefined


-- | Unroll operation in term of joinPS1
unrollPS1 :: Stream1 x -> Stream1 x
unrollPS1 s = Stream1 $ joinPS1 . go $ runS1 s
  where
    go (Mu1 g) = g . joinPS1 . Mu1 $ g
    go (Cons1 x xs) = Cons1 x $ go xs
    go e = e


toListPS1 :: Stream1 x -> [x]
toListPS1 = go . runS1
  where
    go (Var1 x) = x
    go (Mu1 h) = go . h $ []
    go (Cons1 x xs) =  x : go xs
    go _ = undefined

exPS1, exPS1' :: Stream1 Int
exPS1 = Stream1 $ Cons1 0 (Mu1 $ \x -> Cons1 1 (Cons1 2 (Var1 x)))
exPS1' = unrollPS1 exPS1