{-# OPTIONS -fplugin=Rattus.Plugin #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE CPP #-} -- | Programming with streams. module Rattus.Stream ( map , hd , tl , const , constBox , shift , shiftMany , scan , scanMap , scanMap2 , Str(..) , zipWith , zip , unfold , filter , integral ) where import Rattus import Prelude hiding (map, const, zipWith, zip, filter) import Data.VectorSpace -- | @Str a@ is a stream of values of type @a@. data Str a = ! a ::: (O (Str a)) -- all functions in this module are in Rattus {-# ANN module Rattus #-} -- | Get the first element (= head) of a stream. hd :: Str a -> a hd (x ::: _) = x -- | Get the tail of a stream, i.e. the remainder after removing the -- first element. tl :: Str a -> O (Str a) tl (_ ::: xs) = xs -- | Apply a function to each element of a stream. map :: Box (a -> b) -> Str a -> Str b {-# NOINLINE [1] map #-} map f (x ::: xs) = unbox f x ::: delay (map f (adv xs)) -- | Construct a stream that has the same given value at each step. {-# NOINLINE [1] const #-} const :: Stable a => a -> Str a const a = a ::: delay (const a) -- | Variant of 'const' that allows any type @a@ as argument as long -- as it is boxed. {-# NOINLINE [1] constBox #-} constBox :: Box a -> Str a constBox a = unbox a ::: delay (constBox a) -- | Construct a stream by repeatedly applying a function to a given -- start element. That is, @unfold (box f) x@ will produce the stream -- @x ::: f x ::: f (f x) ::: ...@ unfold :: Stable a => Box (a -> a) -> a -> Str a unfold f x = x ::: delay (unfold f (unbox f x)) -- | Similar to Haskell's 'scanl'. -- -- > scan (box f) x (v1 ::: v2 ::: v3 ::: ... ) == (x `f` v1) ::: ((x `f` v1) `f` v2) ::: ... -- -- Note: Unlike 'scanl', 'scan' starts with @x `f` v1@, not @x@. {-# NOINLINE [1] scan #-} scan :: (Stable b) => Box(b -> a -> b) -> b -> Str a -> Str b scan f acc (a ::: as) = acc' ::: delay (scan f acc' (adv as)) where acc' = unbox f acc a -- | 'scanMap' is a composition of 'map' and 'scan': -- -- > scanMap f g x === map g . scan f x {-# NOINLINE [1] scanMap #-} scanMap :: (Stable b) => Box(b -> a -> b) -> Box (b -> c) -> b -> Str a -> Str c scanMap f p acc (a ::: as) = unbox p acc' ::: delay (scanMap f p acc' (adv as)) where acc' = unbox f acc a -- | 'scanMap2' is similar to 'scanMap' but takes two input streams. scanMap2 :: (Stable b) => Box(b -> a1 -> a2 -> b) -> Box (b -> c) -> b -> Str a1 -> Str a2 -> Str c scanMap2 f p acc (a1 ::: as1) (a2 ::: as2) = unbox p acc' ::: delay (scanMap2 f p acc' (adv as1) (adv as2)) where acc' = unbox f acc a1 a2 -- | Similar to 'Prelude.zipWith' on Haskell lists. zipWith :: Box(a -> b -> c) -> Str a -> Str b -> Str c zipWith f (a ::: as) (b ::: bs) = unbox f a b ::: delay (zipWith f (adv as) (adv bs)) -- | Similar to 'Prelude.zip' on Haskell lists. {-# NOINLINE [1] zip #-} zip :: Str a -> Str b -> Str (a:*b) zip (a ::: as) (b ::: bs) = (a :* b) ::: delay (zip (adv as) (adv bs)) -- | Filter out elements from a stream according to a predicate. filter :: Box(a -> Bool) -> Str a -> Str(Maybe' a) filter p = map (box (\a -> if unbox p a then Just' a else Nothing')) {-| Given a value a and a stream as, this function produces a stream that behaves like -} shift :: Stable a => a -> Str a -> Str a shift a (x ::: xs) = a ::: delay (shift x (adv xs)) {-| Given a list @[a1, ..., an]@ of elements and a stream @xs@ this function constructs a stream that starts with the elements @a1, ..., an@, and then proceeds as @xs@. In particular, this means that the ith element of the original stream @xs@ is the (i+n)th element of the new stream. In other words @shiftMany@ behaves like repeatedly applying @shift@ for each element in the list. -} shiftMany :: Stable a => List a -> Str a -> Str a shiftMany l xs = run l Nil xs where run :: Stable a => List a -> List a -> Str a -> Str a run (b :! bs) buf (x ::: xs) = b ::: delay (run bs (x :! buf) (adv xs)) run Nil buf (x ::: xs) = case reverse' buf of b :! bs -> b ::: delay (run bs (x :! Nil) (adv xs)) Nil -> x ::: xs -- | Calculates an approximation of an integral of the stream of type -- @Str a@ (the y-axis), where the stream of type @Str s@ provides the -- distance between measurements (i.e. the distance along the y axis). integral :: (Stable a, VectorSpace a s) => a -> Str s -> Str a -> Str a integral acc (t ::: ts) (a ::: as) = acc' ::: delay (integral acc' (adv ts) (adv as)) where acc' = acc ^+^ (t *^ a) {-# RULES "const/map" forall (f :: Stable b => Box (a -> b)) x. map f (const x) = let x' = unbox f x in const x' ; "map/map" forall f g xs. map f (map g xs) = map (box (unbox f . unbox g)) xs ; "map/scan" forall f p acc as. map p (scan f acc as) = scanMap f p acc as ; "zip/map" forall xs ys f. map f (zip xs ys) = let f' = unbox f in zipWith (box (\ x y -> f' (x :* y))) xs ys #-} #if __GLASGOW_HASKELL__ >= 808 {-# RULES "scan/scan" forall f g b c as. scan g c (scan f b as) = let f' = unbox f; g' = unbox g in scanMap (box (\ (b:*c) a -> let b' = f' b a in (b':* g' c b'))) (box snd') (b:*c) as ; "scan/scanMap" forall f g p b c as. scan g c (scanMap f p b as) = let f' = unbox f; g' = unbox g; p' = unbox p in scanMap (box (\ (b:*c) a -> let b' = f' (p' b) a in (b':* g' c b'))) (box snd') (b:*c) as ; #-} #endif