
   0 | -- Taken and modified from https://github.com/MarcelineVQ/idris2-streaming
   1 | 
   2 | ||| References:
   3 | ||| - https://github.com/MarcelineVQ/idris2-streaming
   4 | ||| - https://hackage.haskell.org/package/streaming
   5 | module Utils.Streaming
   6 | 
   7 | import Control.Monad.Trans
   8 | import Data.List
   9 | import Data.Nat
  10 | 
  11 | export infixl 0 :>
  12 | 
  13 | public export
  14 | data Of : (a : Type) -> (b : Type) -> Type where
  15 |   (:>) : a -> Lazy b -> Of a b
  16 | 
  17 | export
  18 | Bifunctor Of where
  19 |   mapFst f (a :> b) = f a :> b
  20 |   mapSnd f (a :> b) = a :> f b
  21 |   bimap f g (a :> b) = f a :> g b
  22 | 
  23 | export
  24 | Functor (Of a) where
  25 |   map = mapSnd
  26 | 
  27 | ||| The `Stream` type
  28 | public export
  29 | data Stream : (f : Type -> Type) -> (m : Type -> Type) -> (r : Type) -> Type where
  30 |   Step : Inf (f (Stream f m r)) -> Stream f m r
  31 |   Effect : Inf (m (Stream f m r)) -> Stream f m r
  32 |   Return : r -> Stream f m r
  33 |   Build : (forall b. (r -> b) -> (m b -> b) -> (f b -> b) -> b) -> Stream f m r
  34 | 
  35 | ||| Wrap a new layer of a `Stream`
  36 | export
  37 | wrap : f (Stream f m r) -> Stream f m r
  38 | wrap x = Step x
  39 | 
  40 | ||| Wrap a new effect layer of a `Stream`
  41 | export
  42 | effect : m (Stream f m r) -> Stream f m r
  43 | effect x = Effect x
  44 | 
  45 | export
  46 | build : (forall b. (r -> b) -> (m b -> b) -> (f b -> b) -> b) -> Stream f m r
  47 | build = \phi => phi Return (\x => Effect x) (\x => Step x)
  48 | 
  49 | export
  50 | fold : (Functor f, Monad m) => (r -> b) -> (m b -> b) -> (f b -> b) -> Stream f m r -> b
  51 | fold return effect step (Return x) = return x
  52 | fold return effect step (Effect x) = effect (fold return effect step <$> x)
  53 | fold return effect step (Step x) = step (fold return effect step <$> x)
  54 | fold return effect step (Build g) = g return effect step
  55 | 
  56 | export
  57 | destroy : (Functor f, Monad m) => (f a -> a) -> (m a -> a) -> (r -> a) -> Stream f m r -> a
  58 | destroy step effect return = fold return effect step
  59 | 
  60 | ||| Unfold a `Stream`
  61 | public export
  62 | unfold : (Functor f, Monad m) => (a -> m (Either r (f a))) -> a -> Stream f m r
  63 | unfold f a = Effect $ do
  64 |   Right a' <- f a
  65 |   | Left r => pure (Return r)
  66 |   pure (Step (unfold f <$> a'))
  67 | 
  68 | export
  69 | inspect : (Functor f, Monad m) => Stream f m r -> m (Either r (f (Stream f m r)))
  70 | inspect = destroy (pure . (Right . map (effect {f} {m} . map (either Return wrap)))) join (pure . Left)
  71 | 
  72 | export
  73 | hoist : (Functor f, Monad m) => (forall a. m a -> n a) -> Stream f m r -> Stream f n r
  74 | hoist f = fold Return (\x => Effect $ f x) (\x => Step x)
  75 | 
  76 | public export
  77 | (Functor f, Monad m) => Functor (Stream f m) where
  78 |   map f x = Build (\return, effect, step => fold (return . f) effect step x)
  79 | 
  80 | mutual
  81 |   public export
  82 |   covering
  83 |   (Functor f, Monad m) => Applicative (Stream f m) where
  84 |     pure = Return
  85 |     x <*> y = do
  86 |       f <- x
  87 |       v <- y
  88 |       pure (f v)
  89 | 
  90 |   public export
  91 |   covering
  92 |     (Functor f, Monad m) => Monad (Stream f m) where
  93 |       x >>= k = assert_total Build 
  94 |         (\return, effect, step => 
  95 |           fold (fold return effect step . k) effect step x)
  96 | 
  97 | public export
  98 | MonadTrans (Stream f) where
  99 |   lift x = Effect (map Return x)
 100 | 
 101 | public export
 102 | (HasIO m, Monad (Stream f m)) => HasIO (Stream f m) where
 103 |   liftIO x = lift (liftIO x)
 104 | 
 105 | export
 106 | yield : Monad m => a -> Stream (Of a) m ()
 107 | yield x = Step (x :> Return ())
 108 | 
 109 | export
 110 | run : Monad m => Stream m m r -> m r
 111 | run (Return x) = pure x
 112 | run (Effect x) = x >>= run
 113 | run (Step x) = x >>= run
 114 | run (Build g) = run (build g)
 115 | 
 116 | ||| Turns a `Stream` into a list
 117 | public export
 118 | toList : Monad m => Stream (Of a) m r -> m (List a, r)
 119 | toList = destroy (\(a :> b) => map (mapFst (a ::)) b) join (\x => pure (Nil, x))
 120 | 
 121 | ||| `toList` but discards the result
 122 | public export
 123 | toList_ : Monad m => Stream (Of a) m r -> m (List a)
 124 | toList_ = destroy (\(a :> b) => map (a ::) b) join (const (pure Nil))
 125 | 
 126 | ||| Construct a `Stream` from a `List` with a result type
 127 | public export
 128 | fromList : r -> List a -> Stream (Of a) m r
 129 | fromList r Nil = Return r
 130 | fromList r (a :: as) = Step (a :> fromList r as)
 131 | 
 132 | ||| `fromList` but discards the result
 133 | public export
 134 | fromList_ : List a -> Stream (Of a) m ()
 135 | fromList_ = fromList ()
 136 | 
 137 | ||| Concatenate an element into a `Stream`
 138 | public export
 139 | cons : a -> Stream (Of a) m r -> Stream (Of a) m r
 140 | cons x stream = Step (x :> stream)
 141 | 
 142 | export
 143 | next : Monad m => Stream (Of a) m r -> m (Either r (a, Stream (Of a) m r))
 144 | next stream = inspect stream >>= \case
 145 |   Left r => pure (Left r)
 146 |   Right (x :> xs) => pure (Right (x, xs))
 147 | 
 148 | export
 149 | mapf : (Functor f, Monad m) => (forall x. f x -> g x) -> Stream f m r -> Stream g m r
 150 | mapf f s = Build (\return, effect, step => fold return effect (step . f) s)
 151 | 
 152 | export
 153 | mapfM : (Monad m, Functor f) => (forall x. f x -> m (g x)) -> Stream f m r -> Stream g m r
 154 | mapfM f stream = Build (\return, effect, step => fold return effect (effect . map step . f) stream)
 155 | 
 156 | export
 157 | maps : Monad m => (a -> b) -> Stream (Of a) m r -> Stream (Of b) m r
 158 | maps f s = mapf (mapFst f) s
 159 | 
 160 | export
 161 | mapsM : Monad m => (a -> m b) -> Stream (Of a) m r -> Stream (Of b) m r
 162 | mapsM f s = mapfM (\(c :> g) => (:> g) <$> f c) s
 163 | 
 164 | export
 165 | fors : Monad m => (a -> m x) -> Stream (Of a) m r -> m r
 166 | fors f = fold pure join (\(x :> act) => ignore (f x) >> act)
 167 | 
 168 | export
 169 | takeStream : Monad m => Nat -> Stream (Of a) m r -> m (List a, Either r (Stream (Of a) m r))
 170 | takeStream n stream = loop [] n stream where
 171 |   loop : List a -> Nat -> Stream (Of a) m r -> m (List a, Either r (Stream (Of a) m r))
 172 |   loop acc Z stream = pure (acc, Right stream)
 173 |   loop acc (S n) stream = do
 174 |     Right (elem, rest) <- next stream
 175 |     | Left r => pure (acc, Left r)
 176 |     loop (snoc acc elem) n rest
 177 | 
 178 | ||| Split the stream into stream of sublist of length at most n
 179 | export
 180 | chunksOf : Monad m => (n : Nat) -> {auto 0 ok : NonZero n} -> Stream (Of a) m r -> Stream (Of (List a)) m r
 181 | chunksOf n stream = do
 182 |   (chunk, Right rest) <- lift $ takeStream n stream
 183 |   | (chunk, Left r) => yield chunk >> pure r
 184 |   yield chunk *> chunksOf n rest
 185 | 
 186 | ||| Consume all the content in the stream
 187 | export
 188 | consume : Monad m => Stream (Of a) m r -> m ()
 189 | consume = fold (const $ pure ()) join (\(a :> b) => b)
 190 |