git.iko.soy/streamly
on 80daa14 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
204 lines 5.7 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 : experimental
13-- Portability : GHC
14--
15module Streamly.Internal.Data.Array
16 ( Array(..)
18 , foldl'
19 , foldr
21 , length
23 , writeN
24 , write
26 , toStreamD
27 , toStreamDRev
29 , toStream
30 , toStreamRev
31 , read
33 , fromListN
34 , fromList
35 , fromStreamDN
36 , fromStreamD
38 , fromStreamN
39 , fromStream
41 , streamFold
42 , fold
43 )
44where
46import Prelude hiding (foldr, length, read)
47#if !MIN_VERSION_primitive(0,7,1)
48import Control.DeepSeq (NFData(..))
49#endif
50import Control.Monad (when)
51import Control.Monad.IO.Class (liftIO, MonadIO)
52import GHC.IO (unsafePerformIO)
53import GHC.Base (Int(..))
54import Data.Functor.Identity (runIdentity)
55import Data.Primitive.Array hiding (fromList, fromListN)
56import qualified GHC.Exts as Exts
58import Streamly.Internal.Data.Unfold.Types (Unfold(..))
59import Streamly.Internal.Data.Fold.Types (Fold(..))
60import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
61import Streamly.Internal.Data.Stream.Serial (SerialT)
63import qualified Streamly.Internal.Data.Stream.StreamD as D
65{-# NOINLINE bottomElement #-}
66bottomElement :: a
67bottomElement = undefined
69{-# INLINE_NORMAL toStreamD #-}
70toStreamD :: Monad m => Array a -> D.Stream m a
71toStreamD arr = D.Stream step 0
72 where
73 {-# INLINE_LATE step #-}
74 step _ i
75 | i == length arr = return D.Stop
76 step _ (I# i) =
77 return $
78 case Exts.indexArray# (array# arr) i of
79 (# x #) -> D.Yield x (I# i + 1)
81{-# INLINE length #-}
82length :: Array a -> Int
83length = sizeofArray
85{-# INLINE_NORMAL toStreamDRev #-}
86toStreamDRev :: Monad m => Array a -> D.Stream m a
87toStreamDRev arr = D.Stream step (length arr - 1)
88 where
89 {-# INLINE_LATE step #-}
90 step _ i
91 | i < 0 = return D.Stop
92 step _ (I# i) =
93 return $
94 case Exts.indexArray# (array# arr) i of
95 (# x #) -> D.Yield x (I# i - 1)
97{-# INLINE_NORMAL foldl' #-}
98foldl' :: (b -> a -> b) -> b -> Array a -> b
99foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
101{-# INLINE_NORMAL foldr #-}
102foldr :: (a -> b -> b) -> b -> Array a -> b
103foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
105-- writeN n = S.evertM (fromStreamDN n)
106{-# INLINE_NORMAL writeN #-}
107writeN :: MonadIO m => Int -> Fold m a (Array a)
108writeN limit = Fold step initial extract
109 where
110 initial = do
111 marr <- liftIO $ newArray limit bottomElement
112 return (marr, 0)
113 step (marr, i) x
114 | i == limit = return (marr, i)
115 | otherwise = do
116 liftIO $ writeArray marr i x
117 return (marr, i + 1)
118 extract (marr, len) = liftIO $ freezeArray marr 0 len
120{-# INLINE_NORMAL write #-}
121write :: MonadIO m => Fold m a (Array a)
122write = Fold step initial extract
123 where
124 initial = do
125 marr <- liftIO $ newArray 0 bottomElement
126 return (marr, 0, 0)
127 step (marr, i, capacity) x
128 | i == capacity =
129 let newCapacity = max (capacity * 2) 1
130 in do newMarr <- liftIO $ newArray newCapacity bottomElement
131 liftIO $ copyMutableArray newMarr 0 marr 0 i
132 liftIO $ writeArray newMarr i x
133 return (newMarr, i + 1, newCapacity)
134 | otherwise = do
135 liftIO $ writeArray marr i x
136 return (marr, i + 1, capacity)
137 extract (marr, len, _) = liftIO $ freezeArray marr 0 len
139{-# INLINE_NORMAL fromStreamDN #-}
140fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
141fromStreamDN limit str = do
142 marr <- liftIO $ newArray (max limit 0) bottomElement
143 i <-
144 D.foldlM'
145 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
146 (return 0) $
147 D.take limit str
148 liftIO $ freezeArray marr 0 i
150{-# INLINE fromStreamD #-}
151fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
152fromStreamD = D.runFold write
154{-# INLINABLE fromListN #-}
155fromListN :: Int -> [a] -> Array a
156fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
158{-# INLINABLE fromList #-}
159fromList :: [a] -> Array a
160fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
162#if !MIN_VERSION_primitive(0,7,1)
163instance NFData a => NFData (Array a) where
164 {-# INLINE rnf #-}
165 rnf = foldl' (\_ x -> rnf x) ()
166#endif
168{-# INLINE fromStreamN #-}
169fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
170fromStreamN n m = do
171 when (n < 0) $ error "fromStreamN: negative write count specified"
172 fromStreamDN n $ D.toStreamD m
174{-# INLINE fromStream #-}
175fromStream :: MonadIO m => SerialT m a -> m (Array a)
176fromStream m = fromStreamD $ D.toStreamD m
178{-# INLINE_EARLY toStream #-}
179toStream :: (Monad m, IsStream t) => Array a -> t m a
180toStream = D.fromStreamD . toStreamD
182{-# INLINE_EARLY toStreamRev #-}
183toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
184toStreamRev = D.fromStreamD . toStreamDRev
186{-# INLINE fold #-}
187fold :: Monad m => Fold m a b -> Array a -> m b
188fold f arr = D.runFold f (toStreamD arr)
190{-# INLINE streamFold #-}
191streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
192streamFold f arr = f (toStream arr)
194{-# INLINE_NORMAL read #-}
195read :: Monad m => Unfold m (Array a) a
196read = Unfold step inject
197 where
198 inject arr = return (arr, 0)
199 step (arr, i)
200 | i == length arr = return D.Stop
201 step (arr, I# i) =
202 return $
203 case Exts.indexArray# (array# arr) i of
204 (# x #) -> D.Yield x (arr, I# i + 1)