git.iko.soy/streamly
on fa8a2cf help sign in
streamly/src/Streamly/Internal/Data/Prim/Array.hs
205 lines 6.1 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.Prim.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.Prim.Array
19 (
21 -- XXX should it be just Array instead? We should be able to replace one
22 -- array type with another easily.
23 PrimArray(..)
25 -- XXX Prim should be exported from Data.Prim module?
26 , Prim(..)
28 , foldl'
29 , foldr
31 , length
33 , writeN
34 , write
36 , toStreamD
37 , toStreamDRev
39 , toStream
40 , toStreamRev
41 , read
42 , readSlice
44 , fromListN
45 , fromList
46 , fromStreamDN
47 , fromStreamD
49 , fromStreamN
50 , fromStream
52 , streamFold
53 , fold
54 )
55where
57import Prelude hiding (foldr, length, read)
58import Control.DeepSeq (NFData(..))
59import Control.Monad (when)
60import Control.Monad.IO.Class (liftIO, MonadIO)
61import GHC.IO (unsafePerformIO)
62import Data.Primitive.Types (Prim(..))
64import Streamly.Internal.Data.Prim.Array.Types
65import Streamly.Internal.Data.Unfold.Types (Unfold(..))
66import Streamly.Internal.Data.Fold.Types (Fold(..))
67import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
68import Streamly.Internal.Data.Stream.Serial (SerialT)
70import qualified Streamly.Internal.Data.Stream.StreamD as D
72{-# INLINE_NORMAL toStreamD #-}
73toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
74toStreamD arr = D.Stream step 0
75 where
76 {-# INLINE_LATE step #-}
77 step _ i
78 | i == sizeofPrimArray arr = return D.Stop
79 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
81{-# INLINE length #-}
82length :: Prim a => PrimArray a -> Int
83length arr = sizeofPrimArray arr
85{-# INLINE_NORMAL toStreamDRev #-}
86toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
87toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
88 where
89 {-# INLINE_LATE step #-}
90 step _ i
91 | i < 0 = return D.Stop
92 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
94{-# INLINE_NORMAL foldl' #-}
95foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
96foldl' = foldlPrimArray'
98{-# INLINE_NORMAL foldr #-}
99foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
100foldr = foldrPrimArray
102-- writeN n = S.evertM (fromStreamDN n)
103{-# INLINE_NORMAL writeN #-}
104writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
105writeN limit = Fold step initial extract
106 where
107 initial = do
108 marr <- liftIO $ newPrimArray limit
109 return (marr, 0)
110 step (marr, i) x
111 | i == limit = return (marr, i)
112 | otherwise = do
113 liftIO $ writePrimArray marr i x
114 return (marr, i + 1)
115 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
117{-# INLINE_NORMAL write #-}
118write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
119write = Fold step initial extract
120 where
121 initial = do
122 marr <- liftIO $ newPrimArray 0
123 return (marr, 0, 0)
124 step (marr, i, capacity) x
125 | i == capacity =
126 let newCapacity = max (capacity * 2) 1
127 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
128 liftIO $ writePrimArray newMarr i x
129 return (newMarr, i + 1, newCapacity)
130 | otherwise = do
131 liftIO $ writePrimArray marr i x
132 return (marr, i + 1, capacity)
133 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
134 liftIO $ unsafeFreezePrimArray marr
136{-# INLINE_NORMAL fromStreamDN #-}
137fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
138fromStreamDN limit str = do
139 marr <- liftIO $ newPrimArray (max limit 0)
140 _ <-
141 D.foldlM'
142 (\i x -> i `seq` (liftIO $ writePrimArray marr i x) >> return (i + 1))
143 0 $
144 D.take limit str
145 liftIO $ unsafeFreezePrimArray marr
147{-# INLINE fromStreamD #-}
148fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
149fromStreamD str = D.runFold write str
151{-# INLINABLE fromListN #-}
152fromListN :: Prim a => Int -> [a] -> PrimArray a
153fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
155{-# INLINABLE fromList #-}
156fromList :: Prim a => [a] -> PrimArray a
157fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
159instance Prim a => NFData (PrimArray a) where
160 {-# INLINE rnf #-}
161 rnf = foldl' (\_ _ -> ()) ()
163{-# INLINE fromStreamN #-}
164fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
165fromStreamN n m = do
166 when (n < 0) $ error "fromStreamN: negative write count specified"
167 fromStreamDN n $ D.toStreamD m
169{-# INLINE fromStream #-}
170fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
171fromStream m = fromStreamD $ D.toStreamD m
173{-# INLINE_EARLY toStream #-}
174toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
175toStream = D.fromStreamD . toStreamD
177{-# INLINE_EARLY toStreamRev #-}
178toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
179toStreamRev = D.fromStreamD . toStreamDRev
181{-# INLINE fold #-}
182fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
183fold f arr = D.runFold f (toStreamD arr)
185{-# INLINE streamFold #-}
186streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
187streamFold f arr = f (toStream arr)
189{-# INLINE_NORMAL read #-}
190read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
191read = Unfold step inject
192 where
193 inject arr = return (arr, 0)
194 step (arr, i)
195 | i == length arr = return D.Stop
196 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
198{-# INLINE_NORMAL readSlice #-}
199readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
200readSlice off len = Unfold step inject
201 where
202 inject arr = return (arr, off)
203 step (arr, i)
204 | i == min (off + len) (length arr) = return D.Stop
205 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)