git.iko.soy/streamly
on 604f430 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
206 lines 5.9 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 , 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 Data.Functor.Identity (runIdentity)
53import Data.Primitive.Array hiding (fromList, fromListN)
54import GHC.Base (Int(..))
55import GHC.IO (unsafePerformIO)
56import qualified GHC.Exts as Exts
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 Streamly.Internal.Data.Fold.Type as FL
65import qualified Streamly.Internal.Data.Stream.StreamD as D
67{-# NOINLINE bottomElement #-}
68bottomElement :: a
69bottomElement = undefined
71{-# INLINE_NORMAL toStreamD #-}
72toStreamD :: Monad m => Array a -> D.Stream m a
73toStreamD arr = D.Stream step 0
74 where
75 {-# INLINE_LATE step #-}
76 step _ i
77 | i == length arr = return D.Stop
78 step _ (I# i) =
79 return $
80 case Exts.indexArray# (array# arr) i of
81 (# x #) -> D.Yield x (I# i + 1)
83{-# INLINE length #-}
84length :: Array a -> Int
85length = sizeofArray
87{-# INLINE_NORMAL toStreamDRev #-}
88toStreamDRev :: Monad m => Array a -> D.Stream m a
89toStreamDRev arr = D.Stream step (length arr - 1)
90 where
91 {-# INLINE_LATE step #-}
92 step _ i
93 | i < 0 = return D.Stop
94 step _ (I# i) =
95 return $
96 case Exts.indexArray# (array# arr) i of
97 (# x #) -> D.Yield x (I# i - 1)
99{-# INLINE_NORMAL foldl' #-}
100foldl' :: (b -> a -> b) -> b -> Array a -> b
101foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
103{-# INLINE_NORMAL foldr #-}
104foldr :: (a -> b -> b) -> b -> Array a -> b
105foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
107-- writeN n = S.evertM (fromStreamDN n)
108{-# INLINE_NORMAL writeN #-}
109writeN :: MonadIO m => Int -> Fold m a (Array a)
110writeN limit = Fold step initial extract
111 where
112 initial = do
113 marr <- liftIO $ newArray limit bottomElement
114 return $ FL.Partial (Tuple' marr 0)
115 step st@(Tuple' marr i) x
116 | i == limit = fmap FL.Done $ extract st
117 | otherwise = do
118 liftIO $ writeArray marr i x
119 return $ FL.Partial $ Tuple' marr (i + 1)
120 extract (Tuple' marr len) = liftIO $ freezeArray marr 0 len
122{-# INLINE_NORMAL write #-}
123write :: MonadIO m => Fold m a (Array a)
124write = Fold step initial extract
125 where
126 initial = do
127 marr <- liftIO $ newArray 0 bottomElement
128 return $ FL.Partial (Tuple3' marr 0 0)
129 step (Tuple3' marr i capacity) x
130 | i == capacity =
131 let newCapacity = max (capacity * 2) 1
132 in do newMarr <- liftIO $ newArray newCapacity bottomElement
133 liftIO $ copyMutableArray newMarr 0 marr 0 i
134 liftIO $ writeArray newMarr i x
135 return $ FL.Partial $ Tuple3' newMarr (i + 1) newCapacity
136 | otherwise = do
137 liftIO $ writeArray marr i x
138 return $ FL.Partial $ Tuple3' marr (i + 1) capacity
139 extract (Tuple3' marr len _) = liftIO $ freezeArray marr 0 len
141{-# INLINE_NORMAL fromStreamDN #-}
142fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
143fromStreamDN limit str = do
144 marr <- liftIO $ newArray (max limit 0) bottomElement
145 i <-
146 D.foldlM'
147 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
148 (return 0) $
149 D.take limit str
150 liftIO $ freezeArray marr 0 i
152{-# INLINE fromStreamD #-}
153fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
154fromStreamD = D.fold write
156{-# INLINABLE fromListN #-}
157fromListN :: Int -> [a] -> Array a
158fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
160{-# INLINABLE fromList #-}
161fromList :: [a] -> Array a
162fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
164#if !MIN_VERSION_primitive(0,7,1)
165instance NFData a => NFData (Array a) where
166 {-# INLINE rnf #-}
167 rnf = foldl' (\_ x -> rnf x) ()
168#endif
170{-# INLINE fromStreamN #-}
171fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
172fromStreamN n m = do
173 when (n < 0) $ error "fromStreamN: negative write count specified"
174 fromStreamDN n $ D.toStreamD m
176{-# INLINE fromStream #-}
177fromStream :: MonadIO m => SerialT m a -> m (Array a)
178fromStream m = fromStreamD $ D.toStreamD m
180{-# INLINE_EARLY toStream #-}
181toStream :: (Monad m, IsStream t) => Array a -> t m a
182toStream = D.fromStreamD . toStreamD
184{-# INLINE_EARLY toStreamRev #-}
185toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
186toStreamRev = D.fromStreamD . toStreamDRev
188{-# INLINE fold #-}
189fold :: Monad m => Fold m a b -> Array a -> m b
190fold f arr = D.fold f (toStreamD arr)
192{-# INLINE streamFold #-}
193streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
194streamFold f arr = f (toStream arr)
196{-# INLINE_NORMAL read #-}
197read :: Monad m => Unfold m (Array a) a
198read = Unfold step inject
199 where
200 inject arr = return (arr, 0)
201 step (arr, i)
202 | i == length arr = return D.Stop
203 step (arr, I# i) =
204 return $
205 case Exts.indexArray# (array# arr) i of
206 (# x #) -> D.Yield x (arr, I# i + 1)