git.iko.soy/streamly
on 294aa21 help sign in
streamly/src/Streamly/Internal/Data/Array.hs
203 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)
50import Control.DeepSeq (NFData(..))
51import Control.Monad (when)
52import Control.Monad.IO.Class (liftIO, MonadIO)
53import GHC.IO (unsafePerformIO)
54import GHC.Base (Int(..))
55import Data.Functor.Identity (runIdentity)
56import Data.Primitive.Array hiding (fromList, fromListN)
57import qualified GHC.Exts as Exts
59import Streamly.Internal.Data.Unfold.Types (Unfold(..))
60import Streamly.Internal.Data.Fold.Types (Fold(..))
61import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
62import Streamly.Internal.Data.Stream.Serial (SerialT)
64import qualified Streamly.Internal.Data.Stream.StreamD as D
66{-# NOINLINE bottomElement #-}
67bottomElement :: a
68bottomElement = undefined
70{-# INLINE_NORMAL toStreamD #-}
71toStreamD :: Monad m => Array a -> D.Stream m a
72toStreamD arr = D.Stream step 0
73 where
74 {-# INLINE_LATE step #-}
75 step _ i
76 | i == length arr = return D.Stop
77 step _ (I# i) =
78 return $
79 case Exts.indexArray# (array# arr) i of
80 (# x #) -> D.Yield x (I# i + 1)
82{-# INLINE length #-}
83length :: Array a -> Int
84length = sizeofArray
86{-# INLINE_NORMAL toStreamDRev #-}
87toStreamDRev :: Monad m => Array a -> D.Stream m a
88toStreamDRev arr = D.Stream step (length arr - 1)
89 where
90 {-# INLINE_LATE step #-}
91 step _ i
92 | i < 0 = return D.Stop
93 step _ (I# i) =
94 return $
95 case Exts.indexArray# (array# arr) i of
96 (# x #) -> D.Yield x (I# i - 1)
98{-# INLINE_NORMAL foldl' #-}
99foldl' :: (b -> a -> b) -> b -> Array a -> b
100foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
102{-# INLINE_NORMAL foldr #-}
103foldr :: (a -> b -> b) -> b -> Array a -> b
104foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
106-- writeN n = S.evertM (fromStreamDN n)
107{-# INLINE_NORMAL writeN #-}
108writeN :: MonadIO m => Int -> Fold m a (Array a)
109writeN limit = Fold step initial extract
110 where
111 initial = do
112 marr <- liftIO $ newArray limit bottomElement
113 return (marr, 0)
114 step (marr, i) x
115 | i == limit = return (marr, i)
116 | otherwise = do
117 liftIO $ writeArray marr i x
118 return (marr, i + 1)
119 extract (marr, len) = liftIO $ freezeArray marr 0 len
121{-# INLINE_NORMAL write #-}
122write :: MonadIO m => Fold m a (Array a)
123write = Fold step initial extract
124 where
125 initial = do
126 marr <- liftIO $ newArray 0 bottomElement
127 return (marr, 0, 0)
128 step (marr, i, capacity) x
129 | i == capacity =
130 let newCapacity = max (capacity * 2) 1
131 in do newMarr <- liftIO $ newArray newCapacity bottomElement
132 liftIO $ copyMutableArray newMarr 0 marr 0 i
133 liftIO $ writeArray newMarr i x
134 return (newMarr, i + 1, newCapacity)
135 | otherwise = do
136 liftIO $ writeArray marr i x
137 return (marr, i + 1, capacity)
138 extract (marr, len, _) = liftIO $ freezeArray marr 0 len
140{-# INLINE_NORMAL fromStreamDN #-}
141fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
142fromStreamDN limit str = do
143 marr <- liftIO $ newArray (max limit 0) bottomElement
144 i <-
145 D.foldlM'
146 (\i x -> i `seq` liftIO $ writeArray marr i x >> return (i + 1))
147 0 $
148 D.take limit str
149 liftIO $ freezeArray marr 0 i
151{-# INLINE fromStreamD #-}
152fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
153fromStreamD = D.runFold write
155{-# INLINABLE fromListN #-}
156fromListN :: Int -> [a] -> Array a
157fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
159{-# INLINABLE fromList #-}
160fromList :: [a] -> Array a
161fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
163instance NFData a => NFData (Array a) where
164 {-# INLINE rnf #-}
165 rnf = foldl' (\_ x -> rnf x) ()
167{-# INLINE fromStreamN #-}
168fromStreamN :: MonadIO m => Int -> SerialT m a -> m (Array a)
169fromStreamN n m = do
170 when (n < 0) $ error "fromStreamN: negative write count specified"
171 fromStreamDN n $ D.toStreamD m
173{-# INLINE fromStream #-}
174fromStream :: MonadIO m => SerialT m a -> m (Array a)
175fromStream m = fromStreamD $ D.toStreamD m
177{-# INLINE_EARLY toStream #-}
178toStream :: (Monad m, IsStream t) => Array a -> t m a
179toStream = D.fromStreamD . toStreamD
181{-# INLINE_EARLY toStreamRev #-}
182toStreamRev :: (Monad m, IsStream t) => Array a -> t m a
183toStreamRev = D.fromStreamD . toStreamDRev
185{-# INLINE fold #-}
186fold :: Monad m => Fold m a b -> Array a -> m b
187fold f arr = D.runFold f (toStreamD arr)
189{-# INLINE streamFold #-}
190streamFold :: Monad m => (SerialT m a -> m b) -> Array a -> m b
191streamFold f arr = f (toStream arr)
193{-# INLINE_NORMAL read #-}
194read :: Monad m => Unfold m (Array a) a
195read = Unfold step inject
196 where
197 inject arr = return (arr, 0)
198 step (arr, i)
199 | i == length arr = return D.Stop
200 step (arr, I# i) =
201 return $
202 case Exts.indexArray# (array# arr) i of
203 (# x #) -> D.Yield x (arr, I# i + 1)