git.iko.soy/streamly
on 70fb992 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
207 lines 5.7 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE CPP #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE UnboxedTuples #-}
7#include "inline.hs"
9-- |
10-- Module : Streamly.Internal.Data.Array
11-- Copyright : (c) 2019 Composewell Technologies
12--
13-- License : BSD-3-Clause
14-- Maintainer : streamly@composewell.com
15-- Stability : experimental
16-- Portability : GHC
17--
18module Streamly.Internal.Data.Array
19 ( Array(..)
21 , foldl'
22 , foldr
24 , length
26 , writeN
27 , write
29 , toStreamD
30 , toStreamDRev
32 , toStream
33 , toStreamRev
34 , read
36 , fromListN
37 , fromList
38 , fromStreamDN
39 , fromStreamD
41 , fromStreamN
42 , fromStream
44 , streamFold
45 , fold
46 )
47where
49import Prelude hiding (foldr, length, read)
50#if !MIN_VERSION_primitive(0,7,1)
51import Control.DeepSeq (NFData(..))
52#endif
53import Control.Monad (when)
54import Control.Monad.IO.Class (liftIO, MonadIO)
55import GHC.IO (unsafePerformIO)
56import GHC.Base (Int(..))
57import Data.Functor.Identity (runIdentity)
58import Data.Primitive.Array hiding (fromList, fromListN)
59import qualified GHC.Exts as Exts
61import Streamly.Internal.Data.Unfold.Types (Unfold(..))
62import Streamly.Internal.Data.Fold.Types (Fold(..))
63import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
64import Streamly.Internal.Data.Stream.Serial (SerialT)
66import qualified Streamly.Internal.Data.Stream.StreamD as D
68{-# NOINLINE bottomElement #-}
69bottomElement :: a
70bottomElement = undefined
72{-# INLINE_NORMAL toStreamD #-}
73toStreamD :: Monad m => Array a -> D.Stream m a
74toStreamD arr = D.Stream step 0
75 where
76 {-# INLINE_LATE step #-}
77 step _ i
78 | i == length arr = return D.Stop
79 step _ (I# i) =
80 return $
81 case Exts.indexArray# (array# arr) i of
82 (# x #) -> D.Yield x (I# i + 1)
84{-# INLINE length #-}
85length :: Array a -> Int
86length = sizeofArray
88{-# INLINE_NORMAL toStreamDRev #-}
89toStreamDRev :: Monad m => Array a -> D.Stream m a
90toStreamDRev arr = D.Stream step (length arr - 1)
91 where
92 {-# INLINE_LATE step #-}
93 step _ i
94 | i < 0 = return D.Stop
95 step _ (I# i) =
96 return $
97 case Exts.indexArray# (array# arr) i of
98 (# x #) -> D.Yield x (I# i - 1)
100{-# INLINE_NORMAL foldl' #-}
101foldl' :: (b -> a -> b) -> b -> Array a -> b
102foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
104{-# INLINE_NORMAL foldr #-}
105foldr :: (a -> b -> b) -> b -> Array a -> b
106foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
108-- writeN n = S.evertM (fromStreamDN n)
109{-# INLINE_NORMAL writeN #-}
110writeN :: MonadIO m => Int -> Fold m a (Array a)
111writeN limit = Fold step initial extract
112 where
113 initial = do
114 marr <- liftIO $ newArray limit bottomElement
115 return (marr, 0)
116 step (marr, i) x
117 | i == limit = return (marr, i)
118 | otherwise = do
119 liftIO $ writeArray marr i x
120 return (marr, i + 1)
121 extract (marr, len) = liftIO $ freezeArray marr 0 len
123{-# INLINE_NORMAL write #-}
124write :: MonadIO m => Fold m a (Array a)
125write = Fold step initial extract
126 where
127 initial = do
128 marr <- liftIO $ newArray 0 bottomElement
129 return (marr, 0, 0)
130 step (marr, i, capacity) x
131 | i == capacity =
132 let newCapacity = max (capacity * 2) 1
133 in do newMarr <- liftIO $ newArray newCapacity bottomElement
134 liftIO $ copyMutableArray newMarr 0 marr 0 i
135 liftIO $ writeArray newMarr i x
136 return (newMarr, i + 1, newCapacity)
137 | otherwise = do
138 liftIO $ writeArray marr i x
139 return (marr, i + 1, capacity)
140 extract (marr, len, _) = liftIO $ freezeArray marr 0 len
142{-# INLINE_NORMAL fromStreamDN #-}
143fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
144fromStreamDN limit str = do
145 marr <- liftIO $ newArray (max limit 0) bottomElement
146 i <-
147 D.foldlM'
148 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
149 0 $
150 D.take limit str
151 liftIO $ freezeArray marr 0 i
153{-# INLINE fromStreamD #-}
154fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
155fromStreamD = D.runFold write
157{-# INLINABLE fromListN #-}
158fromListN :: Int -> [a] -> Array a
159fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
161{-# INLINABLE fromList #-}
162fromList :: [a] -> Array a
163fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
165#if !MIN_VERSION_primitive(0,7,1)
166instance NFData a => NFData (Array a) where
167 {-# INLINE rnf #-}
168 rnf = foldl' (\_ x -> rnf x) ()
169#endif
171{-# INLINE fromStreamN #-}
172fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
173fromStreamN n m = do
174 when (n < 0) $ error "fromStreamN: negative write count specified"
175 fromStreamDN n $ D.toStreamD m
177{-# INLINE fromStream #-}
178fromStream :: MonadIO m => SerialT m a -> m (Array a)
179fromStream m = fromStreamD $ D.toStreamD m
181{-# INLINE_EARLY toStream #-}
182toStream :: (Monad m, IsStream t) => Array a -> t m a
183toStream = D.fromStreamD . toStreamD
185{-# INLINE_EARLY toStreamRev #-}
186toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
187toStreamRev = D.fromStreamD . toStreamDRev
189{-# INLINE fold #-}
190fold :: Monad m => Fold m a b -> Array a -> m b
191fold f arr = D.runFold f (toStreamD arr)
193{-# INLINE streamFold #-}
194streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
195streamFold f arr = f (toStream arr)
197{-# INLINE_NORMAL read #-}
198read :: Monad m => Unfold m (Array a) a
199read = Unfold step inject
200 where
201 inject arr = return (arr, 0)
202 step (arr, i)
203 | i == length arr = return D.Stop
204 step (arr, I# i) =
205 return $
206 case Exts.indexArray# (array# arr) i of
207 (# x #) -> D.Yield x (arr, I# i + 1)