@@ -0,0 +1,165 @@
1+-- |
2+-- Module : Streamly.Internal.Data.SmallArray
3+-- Copyright : (c) 2019 Composewell Technologies
4+--
5+-- License : BSD-3-Clause
6+-- Maintainer : streamly@composewell.com
7+-- Stability : experimental
8+-- Portability : GHC
9+
10+{-# OPTIONS_GHC -fno-warn-orphans #-}
11+
12+{-# LANGUAGE CPP #-}
13+{-# LANGUAGE MagicHash #-}
14+{-# LANGUAGE UnboxedTuples #-}
15+
16+#include "inline.hs"
17+
18+module Streamly.Internal.Data.SmallArray
19+ ( SmallArray(..)
20+
21+ , foldl'
22+ , foldr
23+
24+ , length
25+
26+ , writeN
27+
28+ , toStreamD
29+ , toStreamDRev
30+
31+ , toStream
32+ , toStreamRev
33+ , read
34+
35+ , fromListN
36+ , fromStreamDN
37+ , fromStreamN
38+
39+ , streamFold
40+ , fold
41+ )
42+where
43+
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
51+
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)
56+
57+import qualified Streamly.Streams.StreamD as D
58+
59+{-# NOINLINE bottomElement #-}
60+bottomElement :: a
61+bottomElement = undefined
62+
63+{-# INLINE length #-}
64+length :: SmallArray a -> Int
65+length arr = sizeofSmallArray arr
66+
67+{-# INLINE_NORMAL toStreamD #-}
68+toStreamD :: Monad m => SmallArray a -> D.Stream m a
69+toStreamD arr = D.Stream step 0
70+ where
71+ {-# INLINE_LATE step #-}
72+ step _ i
73+ | i == length arr = return D.Stop
74+ | otherwise =
75+ return $
76+ case indexSmallArray## arr i of
77+ (# x #) -> D.Yield x (i + 1)
78+
79+{-# INLINE_NORMAL toStreamDRev #-}
80+toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
81+toStreamDRev arr = D.Stream step (length arr - 1)
82+ where
83+ {-# INLINE_LATE step #-}
84+ step _ i
85+ | i < 0 = return D.Stop
86+ | otherwise =
87+ return $
88+ case indexSmallArray## arr i of
89+ (# x #) -> D.Yield x (i - 1)
90+
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
94+
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
98+
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
103+ where
104+ initial = do
105+ marr <- liftIO $ newSmallArray limit bottomElement
106+ return (marr, 0)
107+ step (marr, i) x
108+ | i == limit = return (marr, i)
109+ | otherwise = do
110+ liftIO $ writeSmallArray marr i x
111+ return (marr, i + 1)
112+ extract (marr, len) = liftIO $ freezeSmallArray marr 0 len
113+
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
118+ i <-
119+ D.foldlM'
120+ (\i x -> i `seq` (liftIO $ writeSmallArray marr i x) >> return (i + 1))
121+ 0 $
122+ D.take limit str
123+ liftIO $ freezeSmallArray marr 0 i
124+
125+{-# INLINABLE fromListN #-}
126+fromListN :: Int -> [a] -> SmallArray a
127+fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
128+
129+instance NFData a => NFData (SmallArray a) where
130+ {-# INLINE rnf #-}
131+ rnf = foldl' (\_ x -> rnf x) ()
132+
133+{-# INLINE fromStreamN #-}
134+fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
135+fromStreamN n m = do
136+ when (n < 0) $ error "fromStreamN: negative write count specified"
137+ fromStreamDN n $ D.toStreamD m
138+
139+{-# INLINE_EARLY toStream #-}
140+toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
141+toStream = D.fromStreamD . toStreamD
142+
143+{-# INLINE_EARLY toStreamRev #-}
144+toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
145+toStreamRev = D.fromStreamD . toStreamDRev
146+
147+{-# INLINE fold #-}
148+fold :: Monad m => Fold m a b -> SmallArray a -> m b
149+fold f arr = D.runFold f (toStreamD arr)
150+
151+{-# INLINE streamFold #-}
152+streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
153+streamFold f arr = f (toStream arr)
154+
155+{-# INLINE_NORMAL read #-}
156+read :: Monad m => Unfold m (SmallArray a) a
157+read = Unfold step inject
158+ where
159+ inject arr = return (arr, 0)
160+ step (arr, i)
161+ | i == length arr = return D.Stop
162+ | otherwise =
163+ return $
164+ case indexSmallArray## arr i of
165+ (# x #) -> D.Yield x (arr, i + 1)