git.iko.soy/streamly
on bd5b6fc help sign in
streamly/src/Streamly/Internal/Data/Array/PrimInclude.hs
184 lines 6.4 kB history raw
1-- Copyright : (c) 2020 Composewell Technologies
2-- License : BSD-3-Clause
3-- Maintainer : streamly@composewell.com
4-- Stability : experimental
5-- Portability : GHC
7import Control.Monad (when)
8import Control.Monad.IO.Class (MonadIO(..))
9import Data.Primitive.Types (Prim(..))
10import Streamly.Internal.Data.Fold.Type (Fold(..))
11import Streamly.Internal.Data.Unfold.Type (Unfold(..))
12import Streamly.Internal.Data.Stream.Serial (SerialT(..))
14import qualified Streamly.Internal.Data.Stream.Prelude as P
15import qualified Streamly.Internal.Data.Stream.StreamD as D
17import Prelude hiding (length, null, last, map, (!!), read, concat)
19-------------------------------------------------------------------------------
20-- Construction
21-------------------------------------------------------------------------------
23-- | Create an 'Array' from the first N elements of a stream. The array is
24-- allocated to size N, if the stream terminates before N elements then the
25-- array may hold less than N elements.
26--
27-- /Pre-release/
28{-# INLINE fromStreamN #-}
29fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (Array a)
30fromStreamN n (SerialT m) = do
31 when (n < 0) $ error "writeN: negative write count specified"
32 A.fromStreamDN n $ D.fromStreamK m
34-- | Create an 'Array' from a stream. This is useful when we want to create a
35-- single array from a stream of unknown size. 'writeN' is at least twice
36-- as efficient when the size is already known.
37--
38-- Note that if the input stream is too large memory allocation for the array
39-- may fail. When the stream size is not known, `arraysOf` followed by
40-- processing of indvidual arrays in the resulting stream should be preferred.
41--
42-- /Pre-release/
43{-# INLINE fromStream #-}
44fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (Array a)
45fromStream (SerialT m) = P.fold A.write m
46-- write m = A.fromStreamD $ D.toStreamD m
48-------------------------------------------------------------------------------
49-- Elimination
50-------------------------------------------------------------------------------
52-- | Convert an 'Array' into a stream.
53--
54-- /Pre-release/
55{-# INLINE_EARLY toStream #-}
56toStream :: (MonadIO m, Prim a) => Array a -> SerialT m a
57toStream = SerialT . D.toStreamK . A.toStreamD
58-- XXX add fallback to StreamK rule
59-- {-# RULES "Streamly.Array.read fallback to StreamK" [1]
60-- forall a. S.readK (read a) = K.fromArray a #-}
62-- | Convert an 'Array' into a stream in reverse order.
63--
64-- /Pre-release/
65{-# INLINE_EARLY toStreamRev #-}
66toStreamRev :: (MonadIO m, Prim a) => Array a -> SerialT m a
67toStreamRev = SerialT . D.toStreamK . A.toStreamDRev
68-- XXX add fallback to StreamK rule
69-- {-# RULES "Streamly.Array.readRev fallback to StreamK" [1]
70-- forall a. S.toStreamK (readRev a) = K.revFromArray a #-}
72-- | Unfold an array into a stream.
73--
74-- @since 0.7.0
75{-# INLINE_NORMAL read #-}
76read :: (MonadIO m, Prim a) => Unfold m (Array a) a
77read = Unfold step inject
78 where
80 inject = return
82 {-# INLINE_LATE step #-}
83 step (Array _ _ len) | len == 0 = return D.Stop
84 step arr@(Array arr# off len) =
85 let !x = A.unsafeIndex arr 0
86 in return $ D.Yield x (Array arr# (off + 1) (len - 1))
88-- | Unfold an array into a stream, does not check the end of the array, the
89-- user is responsible for terminating the stream within the array bounds. For
90-- high performance application where the end condition can be determined by
91-- a terminating fold.
92--
93-- The following might not be true, not that the representation changed.
94-- Written in the hope that it may be faster than "read", however, in the case
95-- for which this was written, "read" proves to be faster even though the core
96-- generated with unsafeRead looks simpler.
97--
98-- /Pre-release/
99--
100{-# INLINE_NORMAL unsafeRead #-}
101unsafeRead :: (MonadIO m, Prim a) => Unfold m (Array a) a
102unsafeRead = Unfold step inject
103 where
105 inject = return
107 {-# INLINE_LATE step #-}
108 step arr@(Array arr# off len) =
109 let !x = A.unsafeIndex arr 0
110 in return $ D.Yield x (Array arr# (off + 1) (len - 1))
112-- | > null arr = length arr == 0
113--
114-- /Pre-release/
115{-# INLINE null #-}
116null :: Array a -> Bool
117null arr = length arr == 0
119-------------------------------------------------------------------------------
120-- Folds
121-------------------------------------------------------------------------------
123-- | Fold an array using a 'Fold'.
124--
125-- /Pre-release/
126{-# INLINE fold #-}
127fold :: forall m a b. (MonadIO m, Prim a) => Fold m a b -> Array a -> m b
128fold f arr = P.fold f (getSerialT (toStream arr))
130-- | Fold an array using a stream fold operation.
131--
132-- /Pre-release/
133{-# INLINE streamFold #-}
134streamFold :: (MonadIO m, Prim a) => (SerialT m a -> m b) -> Array a -> m b
135streamFold f arr = f (toStream arr)
137-------------------------------------------------------------------------------
138-- Random reads
139-------------------------------------------------------------------------------
141-- | /O(1)/ Lookup the element at the given index, starting from 0.
142--
143-- /Pre-release/
144{-# INLINE readIndex #-}
145readIndex :: Prim a => Array a -> Int -> Maybe a
146readIndex arr i =
147 if i < 0 || i > length arr - 1
148 then Nothing
149 else Just $ A.unsafeIndex arr i
151-- | > last arr = readIndex arr (length arr - 1)
152--
153-- /Pre-release/
154{-# INLINE last #-}
155last :: Prim a => Array a -> Maybe a
156last arr = readIndex arr (length arr - 1)
158-------------------------------------------------------------------------------
159-- Array stream operations
160-------------------------------------------------------------------------------
162-- | Convert a stream of arrays into a stream of their elements.
163--
164-- Same as the following but more efficient:
165--
166-- > concat = S.concatMap A.read
167--
168-- /Pre-release/
169{-# INLINE concat #-}
170concat :: (MonadIO m, Prim a) => SerialT m (Array a) -> SerialT m a
171-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
172-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
173concat (SerialT m) =
174 SerialT $ D.toStreamK $ D.unfoldMany read (D.fromStreamK m)
176-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
177-- maximum specified size in bytes.
178--
179-- /Pre-release/
180{-# INLINE compact #-}
181compact ::
182 (MonadIO m, Prim a) => Int -> SerialT m (Array a) -> SerialT m (Array a)
183compact n (SerialT xs) =
184 SerialT $ D.toStreamK $ A.packArraysChunksOf n (D.fromStreamK xs)