git.iko.soy/streamly
on a46c7bf help sign in
streamly/src/Streamly/Internal/Data/Array.hs
250 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.Serial (SerialT(..))
60import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
61import Streamly.Internal.Data.Unfold.Type (Unfold(..))
63import qualified GHC.Exts as Exts
64import qualified Streamly.Internal.Data.Fold.Type as FL
65import qualified Streamly.Internal.Data.Stream.StreamD as D
67import Data.Primitive.Array hiding (fromList, fromListN)
68import Prelude hiding (foldr, length, read)
70{-# NOINLINE bottomElement #-}
71bottomElement :: a
72bottomElement = undefined
74-------------------------------------------------------------------------------
75-- Construction - Folds
76-------------------------------------------------------------------------------
78{-# INLINE_NORMAL writeN #-}
79writeN :: MonadIO m => Int -> Fold m a (Array a)
80writeN len = Fold step initial extract
82 where
84 {-# INLINE next #-}
85 next marr i = do
86 let i1 = i + 1
87 st = Tuple' marr i1
88 if len > i1
89 then return $ FL.Partial st
90 else fmap FL.Done $ extract st
92 initial = do
93 marr <- liftIO $ newArray len bottomElement
94 next marr (-1)
96 step (Tuple' marr i) x = do
97 liftIO $ writeArray marr i x
98 next marr i
100 extract (Tuple' marr l) = liftIO $ freezeArray marr 0 l
102{-# INLINE_NORMAL write #-}
103write :: MonadIO m => Fold m a (Array a)
104write = Fold step initial extract
105 where
106 initial = do
107 marr <- liftIO $ newArray 0 bottomElement
108 return $ FL.Partial (Tuple3' marr 0 0)
109 step (Tuple3' marr i capacity) x
110 | i == capacity =
111 let newCapacity = max (capacity * 2) 1
112 in do newMarr <- liftIO $ newArray newCapacity bottomElement
113 liftIO $ copyMutableArray newMarr 0 marr 0 i
114 liftIO $ writeArray newMarr i x
115 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
116 | otherwise = do
117 liftIO $ writeArray marr i x
118 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
119 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
121-------------------------------------------------------------------------------
122-- Construction - from streams
123-------------------------------------------------------------------------------
125{-# INLINE_NORMAL fromStreamDN #-}
126fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
127fromStreamDN limit str = do
128 marr <- liftIO $ newArray (max limit 0) bottomElement
129 i <-
130 D.foldlM'
131 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
132 (return 0) $
133 D.take limit str
134 liftIO $ freezeArray marr 0 i
136{-# INLINE fromStreamD #-}
137fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
138fromStreamD = D.fold write
140{-# INLINE fromStreamN #-}
141fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
142fromStreamN n (SerialT m) = do
143 when (n < 0) $ error "fromStreamN: negative write count specified"
144 fromStreamDN n $ D.fromStreamK m
146{-# INLINE fromStream #-}
147fromStream :: MonadIO m => SerialT m a -> m (Array a)
148fromStream (SerialT m) = fromStreamD $ D.fromStreamK m
150{-# INLINABLE fromListN #-}
151fromListN :: Int -> [a] -> Array a
152fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
154{-# INLINABLE fromList #-}
155fromList :: [a] -> Array a
156fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
158-------------------------------------------------------------------------------
159-- Elimination - Unfolds
160-------------------------------------------------------------------------------
162{-# INLINE length #-}
163length :: Array a -> Int
164length = sizeofArray
166{-# INLINE_NORMAL read #-}
167read :: Monad m => Unfold m (Array a) a
168read = Unfold step inject
169 where
170 inject arr = return (arr, 0)
171 step (arr, i)
172 | i == length arr = return D.Stop
173 step (arr, I# i) =
174 return $
175 case Exts.indexArray# (array# arr) i of
176 (# x #) -> D.Yield x (arr, I# i + 1)
178-------------------------------------------------------------------------------
179-- Elimination - to streams
180-------------------------------------------------------------------------------
182{-# INLINE_NORMAL toStreamD #-}
183toStreamD :: Monad m => Array a -> D.Stream m a
184toStreamD arr = D.Stream step 0
185 where
186 {-# INLINE_LATE step #-}
187 step _ i
188 | i == length arr = return D.Stop
189 step _ (I# i) =
190 return $
191 case Exts.indexArray# (array# arr) i of
192 (# x #) -> D.Yield x (I# i + 1)
194{-# INLINE_NORMAL toStreamDRev #-}
195toStreamDRev :: Monad m => Array a -> D.Stream m a
196toStreamDRev arr = D.Stream step (length arr - 1)
197 where
198 {-# INLINE_LATE step #-}
199 step _ i
200 | i < 0 = return D.Stop
201 step _ (I# i) =
202 return $
203 case Exts.indexArray# (array# arr) i of
204 (# x #) -> D.Yield x (I# i - 1)
206{-# INLINE_EARLY toStream #-}
207toStream :: Monad m => Array a -> SerialT m a
208toStream = SerialT . D.toStreamK . toStreamD
210{-# INLINE_EARLY toStreamRev #-}
211toStreamRev :: Monad m => Array a -> SerialT m a
212toStreamRev = SerialT . D.toStreamK . toStreamDRev
214-------------------------------------------------------------------------------
215-- Elimination - using Folds
216-------------------------------------------------------------------------------
218{-# INLINE_NORMAL foldl' #-}
219foldl' :: (b -> a -> b) -> b -> Array a -> b
220foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
222{-# INLINE_NORMAL foldr #-}
223foldr :: (a -> b -> b) -> b -> Array a -> b
224foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
226#if !MIN_VERSION_primitive(0,7,1)
227instance NFData a => NFData (Array a) where
228 {-# INLINE rnf #-}
229 rnf = foldl' (\_ x -> rnf x) ()
230#endif
232{-# INLINE fold #-}
233fold :: Monad m => Fold m a b -> Array a -> m b
234fold f arr = D.fold f (toStreamD arr)
236{-# INLINE streamFold #-}
237streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
238streamFold f arr = f (toStream arr)
240-------------------------------------------------------------------------------
241-- Random reads and writes
242-------------------------------------------------------------------------------
244-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does
245-- not check the bounds.
246--
247-- @since 0.8.0
248{-# INLINE getIndexUnsafe #-}
249getIndexUnsafe :: Array a -> Int -> a
250getIndexUnsafe = indexArray