@@ -0,0 +1,165 @@
2+-- Module : Streamly.Internal.Data.SmallArray
3+-- Copyright : (c) 2019 Composewell Technologies
5+-- License : BSD-3-Clause
6+-- Maintainer : streamly@composewell.com
7+-- Stability : experimental
10+{-# OPTIONS_GHC -fno-warn-orphans #-}
13+{-# LANGUAGE MagicHash #-}
14+{-# LANGUAGE UnboxedTuples #-}
18+module Streamly.Internal.Data.SmallArray
44+import Prelude hiding (foldr, length, read)
45+import Control.DeepSeq (NFData(..))
46+import Control.Monad (when)
47+import Control.Monad.IO.Class (MonadIO, liftIO)
48+import GHC.IO (unsafePerformIO)
49+import Data.Functor.Identity (runIdentity)
50+import Data.Primitive.SmallArray
52+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
53+import Streamly.Internal.Data.Fold.Types (Fold(..))
54+import Streamly.Streams.StreamK.Type (IsStream)
55+import Streamly.Streams.Serial (SerialT)
57+import qualified Streamly.Streams.StreamD as D
59+{-# NOINLINE bottomElement #-}
61+bottomElement = undefined
64+length :: SmallArray a -> Int
65+length arr = sizeofSmallArray arr
67+{-# INLINE_NORMAL toStreamD #-}
68+toStreamD :: Monad m => SmallArray a -> D.Stream m a
69+toStreamD arr = D.Stream step 0
71+ {-# INLINE_LATE step #-}
73+ | i == length arr = return D.Stop
76+ case indexSmallArray## arr i of
77+ (# x #) -> D.Yield x (i + 1)
79+{-# INLINE_NORMAL toStreamDRev #-}
80+toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
81+toStreamDRev arr = D.Stream step (length arr - 1)
83+ {-# INLINE_LATE step #-}
85+ | i < 0 = return D.Stop
88+ case indexSmallArray## arr i of
89+ (# x #) -> D.Yield x (i - 1)
91+{-# INLINE_NORMAL foldl' #-}
92+foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
93+foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
95+{-# INLINE_NORMAL foldr #-}
96+foldr :: (a -> b -> b) -> b -> SmallArray a -> b
97+foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
99+-- writeN n = S.evertM (fromStreamDN n)
100+{-# INLINE_NORMAL writeN #-}
101+writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
102+writeN limit = Fold step initial extract
105+ marr <- liftIO $ newSmallArray limit bottomElement
108+ | i == limit = return (marr, i)
110+ liftIO $ writeSmallArray marr i x
111+ return (marr, i + 1)
112+ extract (marr, len) = liftIO $ freezeSmallArray marr 0 len
114+{-# INLINE_NORMAL fromStreamDN #-}
115+fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
116+fromStreamDN limit str = do
117+ marr <- liftIO $ newSmallArray (max limit 0) bottomElement
120+ (\i x -> i `seq` (liftIO $ writeSmallArray marr i x) >> return (i + 1))
123+ liftIO $ freezeSmallArray marr 0 i
125+{-# INLINABLE fromListN #-}
126+fromListN :: Int -> [a] -> SmallArray a
127+fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
129+instance NFData a => NFData (SmallArray a) where
131+ rnf = foldl' (\_ x -> rnf x) ()
133+{-# INLINE fromStreamN #-}
134+fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
136+ when (n < 0) $ error "fromStreamN: negative write count specified"
137+ fromStreamDN n $ D.toStreamD m
139+{-# INLINE_EARLY toStream #-}
140+toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
141+toStream = D.fromStreamD . toStreamD
143+{-# INLINE_EARLY toStreamRev #-}
144+toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
145+toStreamRev = D.fromStreamD . toStreamDRev
148+fold :: Monad m => Fold m a b -> SmallArray a -> m b
149+fold f arr = D.runFold f (toStreamD arr)
151+{-# INLINE streamFold #-}
152+streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
153+streamFold f arr = f (toStream arr)
155+{-# INLINE_NORMAL read #-}
156+read :: Monad m => Unfold m (SmallArray a) a
157+read = Unfold step inject
159+ inject arr = return (arr, 0)
161+ | i == length arr = return D.Stop
164+ case indexSmallArray## arr i of
165+ (# x #) -> D.Yield x (arr, i + 1)