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