{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE RankNTypes #-}
module GardGround.Utils.Fusion (
Step(..),
stepState,
runSteps,
StepFun(..),
Stream,
runStream,
embedTfs,
step
) where
import Control.Applicative (Alternative(..))
import Prelude hiding (filter)
import Data.Bifunctor
import qualified Data.Text.Internal.Fusion as TF
data Step s a = Done | Skip !s | Yield !a !s
instance Functor (Step s) where
fmap _ Done = Done
fmap _ (Skip s) = Skip s
fmap f (Yield a s) = Yield (f a) s
{-# INLINE fmap #-}
instance Bifunctor Step where
bimap _ _ Done = Done
bimap f _ (Skip s) = Skip (f s)
bimap f g (Yield a s) = Yield (g a) (f s)
{-# INLINE bimap #-}
stepState :: Step s a -> Maybe s
stepState Done = Nothing
stepState (Skip st) = Just st
stepState (Yield _ st) = Just st
data StepFun s a = StepFun (s -> Step s a)
instance Functor (StepFun s) where
fmap f inif = inif >>= (\x -> pure (f x))
{-# INLINE fmap #-}
instance Applicative (StepFun s) where
pure x = StepFun $ \st -> Yield x st
{-# INLINE pure #-}
f <*> v = f >>= (\fx -> v >>= (\vx -> StepFun $ Yield (fx vx)))
instance Monad (StepFun s) where
return = pure
(StepFun f1) >>= f2 = StepFun go
where
go st = case f1 st of
Done -> Done
Skip st2 -> Skip st2
Yield x st2 -> let (StepFun f3) = (f2 x) in f3 st2
-- this makes `guard` work
instance Alternative (StepFun s) where
empty = StepFun $ Skip
{-# INLINE empty #-}
StepFun f1 <|> StepFun f2 = StepFun $ \st ->
case f1 st of
Yield a st2 -> Yield a st2
-- backtracking
Done -> f2 st
Skip _ -> f2 st
runSteps :: StepFun s a -> s -> [a]
runSteps (StepFun sf) = go
where
go st = case sf st of
Done -> []
Skip st2 -> go st2
Yield a st2 -> a:(go st2)
data Stream a = forall s. Stream
(StepFun s a) -- ^ stepper function
!s -- ^ current state
instance Functor Stream where
fmap f (Stream step_ cst) = Stream (fmap f step_) cst
runStream :: Stream a -> [a]
runStream (Stream sf ctx) = runSteps sf ctx
{-# INLINE runStream #-}
smap :: TF.Step s a -> Step s a
smap TF.Done = Done
smap (TF.Skip st) = Skip st
smap (TF.Yield av st) = Yield av st
-- | embed a stream from Data.Text into our stream type
embedTfs :: TF.Stream x -> Stream x
embedTfs (TF.Stream step_ st _) = Stream (StepFun (smap . step_)) st
-- | step an inner stream;
-- forwards skips and turns Done into Nothing-yields
step :: StepFun (Stream a) (Maybe a)
step = StepFun $ \(Stream (StepFun f) st) -> case f st of
Done -> Yield Nothing (Stream (StepFun f) st)
Skip st2 -> Skip (Stream (StepFun f) st2)
Yield x st2 -> Yield (Just x) (Stream (StepFun f) st2)