Another effectful stream representation

Streams.hs

{-# Language ExistentialQuantification #-}
{-# Language GADTs #-}
{-# Language RankNTypes #-}
{-# Language ScopedTypeVariables #-}

import Control.Monad
import Control.Applicative
import Data.List hiding (uncons)
import Data.Time

data Free f r
  = Pure r
  | Eval (f r)
  | Fail Err
  | forall x . Bind (Free f x) (x -> Free f r)
  | OnError (Free f r) (Err -> Free f r)

type Err = String

type Pull f o r = Free (StreamF f o) r
type Stream f o = Pull f o ()

data ViewL f r where
  Bound :: f x -> (x -> Free f r) -> Maybe (Err -> Free f r) -> ViewL f r
  Done :: r -> ViewL f r
  Failed :: Err -> ViewL f r

viewL :: Free f r -> ViewL f r
viewL f = go f K0 Nothing where
  bind :: Applicative f => K f a b -> a -> f b
  bind K0 a = pure a
  bind (K f) a = f a
  go :: Free f x -> K (Free f) x r -> Maybe (Err -> Free f r) -> ViewL f r
  go (Pure x) k _ = case k of K0 -> Done x; K k -> viewL (k x)
  go (Eval fx) k onErr = case k of
    K0 -> Bound fx pure onErr
    K k -> Bound fx k onErr
  go (Fail err) _ onErr = case onErr of
    Nothing -> Failed err
    Just onErr -> viewL (onErr err)
  go (OnError fx onErrInner) k onErr = case k of
    K0 -> case onErr of
      Nothing -> go fx K0 (Just onErrInner)
      Just onErr -> go fx K0 (Just $ \e -> OnError (onErrInner e) onErr)
    K k -> case onErr of
      Nothing -> go fx (K k) (Just $ \e -> onErrInner e >>= k)
      Just onErr -> go fx (K k) (Just $ \e -> OnError (onErrInner e >>= k) onErr)
  go (Bind x f) K0 onErr = go x (K f) onErr
  go (Bind x f) (K k) onErr = go x (K $ \x -> f x >>= k) onErr

data K f a b where K0 :: K f a a; K :: (a -> f b) -> K f a b

data StreamF f o r where
  Wrap :: f r -> StreamF f x r
  Output :: [o] -> StreamF f o () -- use more interesting pure stream type here
  Outputs :: Stream f o -> StreamF f o ()
  -- todo: add other instructions here, like UnconsAsync, Acquire, Release, Snapshot
  -- implement runFold' which handles resource map

uncons :: Stream f o -> Pull f x (Maybe ([o], Stream f o))
uncons s = case viewL s of
  Done () -> pure Nothing
  Bound e f onErr -> case e of
    Wrap fx -> eval (Wrap fx) >>= (uncons . f)
    Output os -> pure (Just (os, f ()))
    Outputs os -> uncons os >>= \o -> case o of
      Nothing -> uncons (f ())
      Just (hd, tl) -> pure (Just (hd, tl >> f ()))

flatMap :: Stream f o -> (o -> Stream f o2) -> Stream f o2
flatMap s f = uncons s >>= \o -> case o of
  Nothing -> pure ()
  -- todo - a bit of extra work on the result of `map f os` to collapse
  -- pure segments as much as possible
  Just (hd, tl) -> eval (Outputs (foldr (>>) (pure ()) (map f hd))) >> (tl `flatMap` f)

runFold :: Monad f => (b -> a -> b) -> b -> Stream f a -> f b
runFold f z s = go f z (viewL (uncons s)) where
  go :: Monad f => (b -> a -> b) -> b -> ViewL (StreamF f x) (Maybe ([a], Stream f a)) -> f b
  go _ z (Done Nothing) = pure z
  go f z (Done (Just (hd, tl))) = go f (foldl' f z hd) (viewL (uncons tl))
  go f z (Bound (Wrap fx) g onErr) = fx >>= \x -> go f z (viewL (g x))

transform :: (forall x . f x -> g x) -> Free f a -> Free g a
transform nt f = case viewL f of
  Done x -> Pure x
  Bound fx g onErr -> Bind (Eval (nt fx)) (\x -> transform nt (g x))

eval :: f a -> Free f a
eval = Eval

interpret :: Monad f => Free f a -> f a
interpret f = case viewL f of
  Done a -> pure a
  Bound fa k onErr -> fa >>= (interpret . k)

emit :: a -> Stream f a
emit a = eval (Output [a])

emits :: [a] -> Stream f a
emits as = eval (Output as)

timeit :: IO () -> IO ()
timeit e = do
  start <- getCurrentTime
  e
  end <- getCurrentTime
  print (diffUTCTime end start)

appendEx :: [(Int, Stream IO Int)]
appendEx = [ (i, foldl' (>>) (pure ()) (map emit [1..i])) | i <- [100000,200000,400000]]

constantAppend = appendEx `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

bindEx :: [(Int, Stream IO Int)]
bindEx = [ (i, void $ foldl' (>>) (pure 1) (map pure [1..i])) | i <- [100000,200000,400000]]

constantBind = bindEx `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

bindEx2 :: [(Int, Stream IO Int)]
bindEx2 = [ (i, void $ foldl' (>>) (pure 1) (map (eval . Wrap . pure) [1..i])) | i <- [100000,200000,400000]]

constantBind2 = bindEx2 `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

bindEx3 :: [(Int, Stream IO Int)]
bindEx3 = [ (i, void $ foldr (>>) (pure 1) (map (eval . Wrap . pure) [1..i])) | i <- [100000,200000,400000]]

constantBind3 = bindEx3 `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

flatMapEx :: [(Int, Stream IO Int)]
flatMapEx = [ (i, void $ foldl' go (emit 1) [1..i]) | i <- [100000,200000,400000]] where
  go s i = s `flatMap` \k -> emit i

flatMapEx2 :: [(Int, Stream IO Int)]
flatMapEx2 = [ (i, void $ foldl' go (emits [1..i]) [1..1000]) | i <- [1000,2000,4000]] where
  go s i = s `flatMap` \k -> emit i

constantFlatMap = flatMapEx `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

constantFlatMap2 = flatMapEx2 `forM_` \(i, s) -> timeit $ do
  n <- runFold (+) 0 s
  putStrLn (show i ++ ": " ++ show n)

tests = do
  putStrLn "- bind(3) ----"
  constantBind3 -- O(n)
  putStrLn "- bind(2) ----"
  constantBind2 -- O(n)
  putStrLn "- constantBind ----"
  constantBind -- O(n)
  putStrLn "- constantAppend ----"
  constantAppend -- O(n)
  putStrLn "- constantFlatMap ----"
  constantFlatMap -- O(n)
  putStrLn "- constantFlatMap (2) ----"
  constantFlatMap2 -- O(n)

instance Monad (Free f) where return = Pure; (>>=) = Bind
instance Applicative (Free f) where pure = return; (<*>) = ap
instance Functor (Free f) where fmap = liftM