Canos UwU

pipes.idr

module Pipes 

import Control.Monad.Trans
import Data.IORef

-- It's just https://github.com/QuentinDuval/IdrisPipes with one small change.

||| PipeT is a composable data type to describe streams of data that 
||| can be generated and can flow upstream and downstream just like 
||| javascript generators 
||| 
|||   Ex:   runEffect $ stdinC .| mapC (cast) .| mapC (+2) .| stdoutC
||| 
|||   * `u` : Data that is flowing upwards
|||   * `d` : Data that is flowing downwards
|||   * `i` : Type from the last pipe
|||   * `m` : Internal monad 
|||   * `r` : Return type of the transformer. Usually it's unit because
|||           the data will flow inside the monad and will not generate
|||           any value in the end, just computing some effects.

public export
data PipeT : (u, d, i : Type) -> (m : Type -> Type) -> (r : Type) -> Type where
    ||| Lift a value to the pipe transformer
    Pure   : r                                       -> PipeT u d i m r
    
    ||| Executes an action (it's the reason it's inside a monad) and 
    ||| then returns a continuation of the pipe. It's `Inf` because it's lazy.
    Action : Inf (m (PipeT u d i m r))               -> PipeT u d i m r 
    
    ||| Yield stops the execution and sends a value. It works by running
    ||| two pipes that are together until the left pipe reaches a Yield
    ||| and the right pipe reachs the Await so they change values and continue
    ||| the execution.
    Yield  : Inf (PipeT u d i m r)     -> m () -> d  -> PipeT u d i m r
    --                                    └ A finalizer effect! that can be useful
    --                                      to remove file descriptors and things like that


    ||| Await stops the execution to get the previous value return type or 
    ||| the yield value and returns the continuation of the pipe.
    Await  : ((Either i u) -> Inf (PipeT u d i m r)) -> PipeT u d i m r  

||| Source is a Sink that cannot use `await` and will not receive any values from previous sinks. 
||| It's possible because due to the type that is set to Void (there are no constructors).
public export
Source : (Type -> Type) -> (Type) -> Type
Source m b = PipeT Void b Void m ()
--                  |      └ Cannot receive values from previous sinks
--               Cannot send values upstream (cannot await) 

||| A pipe can `yield` and `await`
public export
Pipe : {r: Type} -> Type -> (Type -> Type) -> Type -> Type
Pipe {r} a m b = PipeT a b r m ()

||| A sink cannot send values upstream so it cannot yield
public export
Sink : {r1: Type} -> Type -> (Type -> Type) -> Type -> Type
Sink {r1} a m r = PipeT a Void r1 m r
--                         └ Cannot send values upstream

||| An Effect is just the composition of a Source and a Sink
||| So it cannot await or yield values by it's oww, dont have a 
||| previous value and just sends the result value.
public export
Effect : (m: Type -> Type) -> (r: Type) -> Type
Effect m r = PipeT Void Void Void m r

||| Yields a value
public export
yield : Monad m => d -> PipeT u d i m ()
yield = Yield (Pure ()) (pure ())

||| Yields a continuation
public export
await : Monad m => PipeT u d i m (Maybe u)
await = Await $ \case Right x => Pure (Just x)
                      _       => Pure Nothing

-- Some cool implementations 

public export
Monad m => Functor (PipeT u d i m) where
    map f (Pure r)           = Pure (f r)
    map f (Action r)         = Action $ r >>= \x => pure (map f x)
    map f (Yield next fin b) = Yield (map f next) fin b
    map f (Await cont)       = Await $ \x => map f (cont x) 

public export
Monad m => Applicative (PipeT u d i m) where 
    pure = Pure 
    (Pure f)           <*> pA = map f pA
    (Action r)         <*> pA = Action $ r >>= \x => pure (x <*> pA)
    (Yield next fin b) <*> pA = Yield (next <*> pA) fin b
    (Await cont)       <*> pA = Await $ \x => (cont x) <*> pA

public export
Monad m => Monad (PipeT u d i m) where 
    (Pure f)           >>= fb = fb f 
    (Action r)         >>= fb = assert_total $ Action $ (>>= fb) <$> r 
    (Yield next fin b) >>= fb = Yield (next >>= fb) fin b
    (Await cont)       >>= fb = Await $ \x => (cont x) >>= fb

public export
MonadTrans (PipeT u d i) where 
    lift m = Action (m >>= \x => pure (Pure x))

infixr 9 .|


||| So we can assemble two pipes by getting the upstream and 
||| downstream values and moving them. 
|||
||| This function works by getting the first pipe and running all the Actions
||| Until it reaches the `yield` and the other pipe we run until it reaches `await`

export --       
(.|) : Monad m => PipeT u d i m o -> PipeT d c o m o2 -> PipeT u c i m o2
(.|) =  pull (pure ())
    where mutual
        pull : m () -> PipeT u d i m o -> PipeT d c o m o2 -> PipeT u c i m o2
        --                                       > Pull the other part | join the finalizer with the last finalizer
        pull final up (Yield next fin c) = Yield (pull final up next) (fin >> final) c
        --                                > Compose the pipe action with the pull of the rest of the pipe |
        pull final up (Action a)         = lift a >>= pull final up
        --                               > Push the up pipe and gets the value to continue the computation
        pull final up (Await cont)       = up `push` \x => cont x
        pull final up (Pure r)           = lift final >> Pure r

        push :  PipeT u d i m o -> (Either o d -> PipeT d c o m o2) -> PipeT u c i m o2
        --                             > Push to get the yield value
        push (Await t)          down = Await (\a => t a `push` down)
        --                             > Run the action and push the rest
        push (Action a)         down = lift a >>= (`push` down)
        --                             > pull the next but run the continuation
        push (Yield next fin b) down = pull fin next (down (Right b))
        push (Pure r)           down = pull (pure ()) (Pure r) (down (Left r))

||| Runs a pure pipe 
export
runPipe : Monad m => Effect m r -> m r
runPipe (Pure r)        = pure r 
runPipe (Action action) = action >>= runPipe
runPipe (Await cont)    = runPipe $ Await (either absurd absurd)
runPipe (Yield next finish b) impossible

export
awaitOr : PipeT d u i m (Either i d)
awaitOr = Await $ \v => Pure v

||| Function to await forever
export
awaitForever : Monad m => (d -> PipeT d u r m r1) -> PipeT d u r m r
awaitForever f = awaitOr >>= either Pure (\x => f x >>= \_ => awaitForever f)

||| Sink that ignores everything
export
discard : Monad m => Sink {r1=r} u m r
discard = awaitForever $ \_ => pure ()

||| It's not tail call optimized so if you run it a million times
||| you'll get in to some memory problems....
stdinLn : String -> Source IO String
stdinLn promptLine = recur where
  recur : Source IO String
  recur = do
    lift (putStr promptLine)
    res <- lift getLine -- From prelude 
    yield res
    recur

||| SInk to prnit
public export
stdoutLn : Show a => Sink {r1=r} a IO r
stdoutLn = awaitForever (\x => lift (printLn x))

-- Tail recursion on Monads just like Prescript-tailrec

-- data to represent if we should or not continue the recursion
public export
data Step a b = Loop a | Done b

public export
interface Monad m => MonadRec m where   
    -- It's partial and can only be total by expressing it 
    -- with WellFounded. https://github.com/stefan-hoeck/idris2-tailrec
    partial tailRecM : (a -> m (Step a b)) -> a -> m b

-- Tail call function to perform some effects with 
-- a recursive like structure
export
untilE : IO Bool -> IO ()
untilE n = if unsafePerformIO n 
            then untilE n
            else pure ()

-- Implementation for IO
export
MonadRec IO where 
    tailRecM f a = do 
        r <- f a >>= newIORef -- Stores the result of the monad
        untilE $ readIORef r >>= 
            \case Loop a => f a >>= writeIORef r >> pure True  
                  Done a => pure False
        fromDone <$> readIORef r
    where 
        partial 
        fromDone : Step _ b -> b
        fromDone (Done r) = r

-- Runs a monad forever with tail call optimization with constant space usage.
public export
forever : MonadRec m => m a -> m b
forever ma = tailRecM (\u => Loop u <$ ma) ()