git.iko.soy/streamly
on 705d2cf help sign in
streamly/src/Streamly/Internal/Data/Array.hs
242 lines 7.4 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE UnboxedTuples #-}
4#include "inline.hs"
6-- |
7-- Module : Streamly.Internal.Data.Array
8-- Copyright : (c) 2019 Composewell Technologies
9--
10-- License : BSD-3-Clause
11-- Maintainer : streamly@composewell.com
12-- Stability : pre-release
13-- Portability : GHC
14--
15module Streamly.Internal.Data.Array
16 ( Array(..)
18 -- * Construction
19 , writeN
20 , write
22 , fromStreamDN
23 , fromStreamD
24 , fromStreamN
25 , fromStream
27 , fromListN
28 , fromList
30 -- * Elimination
31 , length
32 , read
34 , toStreamD
35 , toStreamDRev
36 , toStream
37 , toStreamRev
39 , foldl'
40 , foldr
41 , streamFold
42 , fold
44 -- * Random Access
45 , getIndexUnsafe
46 )
47where
49#if !MIN_VERSION_primitive(0,7,1)
50import Control.DeepSeq (NFData(..))
51#endif
52import Control.Monad (when)
53import Control.Monad.IO.Class (liftIO, MonadIO)
54import Data.Functor.Identity (runIdentity)
55import GHC.Base (Int(..))
56import GHC.IO (unsafePerformIO)
58import Streamly.Internal.Data.Fold.Type (Fold(..))
59import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
60import Streamly.Internal.Data.Stream.Serial (SerialT)
61import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
62import Streamly.Internal.Data.Unfold.Type (Unfold(..))
64import qualified GHC.Exts as Exts
65import qualified Streamly.Internal.Data.Fold.Type as FL
66import qualified Streamly.Internal.Data.Stream.StreamD as D
68import Data.Primitive.Array hiding (fromList, fromListN)
69import Prelude hiding (foldr, length, read)
71{-# NOINLINE bottomElement #-}
72bottomElement :: a
73bottomElement = undefined
75-------------------------------------------------------------------------------
76-- Construction - Folds
77-------------------------------------------------------------------------------
79-- writeN n = S.evertM (fromStreamDN n)
80{-# INLINE_NORMAL writeN #-}
81writeN :: MonadIO m => Int -> Fold m a (Array a)
82writeN limit = Fold step initial extract
83 where
84 initial = do
85 marr <- liftIO $ newArray limit bottomElement
86 return $ FL.Partial (Tuple' marr 0)
87 step st@(Tuple' marr i) x
88 | i == limit = fmap FL.Done $ extract st
89 | otherwise = do
90 liftIO $ writeArray marr i x
91 return $ FL.Partial $ Tuple' marr (i + 1)
92 extract (Tuple' marr len) = liftIO $ freezeArray marr 0 len
94{-# INLINE_NORMAL write #-}
95write :: MonadIO m => Fold m a (Array a)
96write = Fold step initial extract
97 where
98 initial = do
99 marr <- liftIO $ newArray 0 bottomElement
100 return $ FL.Partial (Tuple3' marr 0 0)
101 step (Tuple3' marr i capacity) x
102 | i == capacity =
103 let newCapacity = max (capacity * 2) 1
104 in do newMarr <- liftIO $ newArray newCapacity bottomElement
105 liftIO $ copyMutableArray newMarr 0 marr 0 i
106 liftIO $ writeArray newMarr i x
107 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
108 | otherwise = do
109 liftIO $ writeArray marr i x
110 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
111 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
113-------------------------------------------------------------------------------
114-- Construction - from streams
115-------------------------------------------------------------------------------
117{-# INLINE_NORMAL fromStreamDN #-}
118fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
119fromStreamDN limit str = do
120 marr <- liftIO $ newArray (max limit 0) bottomElement
121 i <-
122 D.foldlM'
123 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
124 (return 0) $
125 D.take limit str
126 liftIO $ freezeArray marr 0 i
128{-# INLINE fromStreamD #-}
129fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
130fromStreamD = D.fold write
132{-# INLINE fromStreamN #-}
133fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
134fromStreamN n m = do
135 when (n < 0) $ error "fromStreamN: negative write count specified"
136 fromStreamDN n $ D.toStreamD m
138{-# INLINE fromStream #-}
139fromStream :: MonadIO m => SerialT m a -> m (Array a)
140fromStream m = fromStreamD $ D.toStreamD m
142{-# INLINABLE fromListN #-}
143fromListN :: Int -> [a] -> Array a
144fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
146{-# INLINABLE fromList #-}
147fromList :: [a] -> Array a
148fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
150-------------------------------------------------------------------------------
151-- Elimination - Unfolds
152-------------------------------------------------------------------------------
154{-# INLINE length #-}
155length :: Array a -> Int
156length = sizeofArray
158{-# INLINE_NORMAL read #-}
159read :: Monad m => Unfold m (Array a) a
160read = Unfold step inject
161 where
162 inject arr = return (arr, 0)
163 step (arr, i)
164 | i == length arr = return D.Stop
165 step (arr, I# i) =
166 return $
167 case Exts.indexArray# (array# arr) i of
168 (# x #) -> D.Yield x (arr, I# i + 1)
170-------------------------------------------------------------------------------
171-- Elimination - to streams
172-------------------------------------------------------------------------------
174{-# INLINE_NORMAL toStreamD #-}
175toStreamD :: Monad m => Array a -> D.Stream m a
176toStreamD arr = D.Stream step 0
177 where
178 {-# INLINE_LATE step #-}
179 step _ i
180 | i == length arr = return D.Stop
181 step _ (I# i) =
182 return $
183 case Exts.indexArray# (array# arr) i of
184 (# x #) -> D.Yield x (I# i + 1)
186{-# INLINE_NORMAL toStreamDRev #-}
187toStreamDRev :: Monad m => Array a -> D.Stream m a
188toStreamDRev arr = D.Stream step (length arr - 1)
189 where
190 {-# INLINE_LATE step #-}
191 step _ i
192 | i < 0 = return D.Stop
193 step _ (I# i) =
194 return $
195 case Exts.indexArray# (array# arr) i of
196 (# x #) -> D.Yield x (I# i - 1)
198{-# INLINE_EARLY toStream #-}
199toStream :: (Monad m, IsStream t) => Array a -> t m a
200toStream = D.fromStreamD . toStreamD
202{-# INLINE_EARLY toStreamRev #-}
203toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
204toStreamRev = D.fromStreamD . toStreamDRev
206-------------------------------------------------------------------------------
207-- Elimination - using Folds
208-------------------------------------------------------------------------------
210{-# INLINE_NORMAL foldl' #-}
211foldl' :: (b -> a -> b) -> b -> Array a -> b
212foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
214{-# INLINE_NORMAL foldr #-}
215foldr :: (a -> b -> b) -> b -> Array a -> b
216foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
218#if !MIN_VERSION_primitive(0,7,1)
219instance NFData a => NFData (Array a) where
220 {-# INLINE rnf #-}
221 rnf = foldl' (\_ x -> rnf x) ()
222#endif
224{-# INLINE fold #-}
225fold :: Monad m => Fold m a b -> Array a -> m b
226fold f arr = D.fold f (toStreamD arr)
228{-# INLINE streamFold #-}
229streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
230streamFold f arr = f (toStream arr)
232-------------------------------------------------------------------------------
233-- Random reads and writes
234-------------------------------------------------------------------------------
236-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
237-- not check the bounds.
238--
239-- @since 0.8.0
240{-# INLINE getIndexUnsafe #-}
241getIndexUnsafe :: Array a -> Int -> a
242getIndexUnsafe = indexArray