git.iko.soy/streamly
on 3d033ae help sign in
streamly/src/Streamly/Internal/Data/Prim/Array.hs
201 lines 6.0 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Prim.Array
7-- Copyright : (c) 2019 Composewell Technologies
8--
9-- License : BSD-3-Clause
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
12-- Portability : GHC
13--
14module Streamly.Internal.Data.Prim.Array
15 (
17 -- XXX should it be just Array instead? We should be able to replace one
18 -- array type with another easily.
19 PrimArray(..)
21 -- XXX Prim should be exported from Data.Prim module?
22 , Prim(..)
24 , foldl'
25 , foldr
27 , length
29 , writeN
30 , write
32 , toStreamD
33 , toStreamDRev
35 , toStream
36 , toStreamRev
37 , read
38 , readSlice
40 , fromListN
41 , fromList
42 , fromStreamDN
43 , fromStreamD
45 , fromStreamN
46 , fromStream
48 , streamFold
49 , fold
50 )
51where
53import Prelude hiding (foldr, length, read)
54import Control.DeepSeq (NFData(..))
55import Control.Monad (when)
56import Control.Monad.IO.Class (liftIO, MonadIO)
57import GHC.IO (unsafePerformIO)
58import Data.Primitive.Types (Prim(..))
60import Streamly.Internal.Data.Prim.Array.Types
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{-# INLINE_NORMAL toStreamD #-}
69toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
70toStreamD arr = D.Stream step 0
71 where
72 {-# INLINE_LATE step #-}
73 step _ i
74 | i == sizeofPrimArray arr = return D.Stop
75 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
77{-# INLINE length #-}
78length :: Prim a => PrimArray a -> Int
79length = sizeofPrimArray
81{-# INLINE_NORMAL toStreamDRev #-}
82toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
83toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
84 where
85 {-# INLINE_LATE step #-}
86 step _ i
87 | i < 0 = return D.Stop
88 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
90{-# INLINE_NORMAL foldl' #-}
91foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
92foldl' = foldlPrimArray'
94{-# INLINE_NORMAL foldr #-}
95foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
96foldr = foldrPrimArray
98-- writeN n = S.evertM (fromStreamDN n)
99{-# INLINE_NORMAL writeN #-}
100writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
101writeN limit = Fold step initial extract
102 where
103 initial = do
104 marr <- liftIO $ newPrimArray limit
105 return (marr, 0)
106 step (marr, i) x
107 | i == limit = return (marr, i)
108 | otherwise = do
109 liftIO $ writePrimArray marr i x
110 return (marr, i + 1)
111 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
113{-# INLINE_NORMAL write #-}
114write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
115write = Fold step initial extract
116 where
117 initial = do
118 marr <- liftIO $ newPrimArray 0
119 return (marr, 0, 0)
120 step (marr, i, capacity) x
121 | i == capacity =
122 let newCapacity = max (capacity * 2) 1
123 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
124 liftIO $ writePrimArray newMarr i x
125 return (newMarr, i + 1, newCapacity)
126 | otherwise = do
127 liftIO $ writePrimArray marr i x
128 return (marr, i + 1, capacity)
129 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
130 liftIO $ unsafeFreezePrimArray marr
132{-# INLINE_NORMAL fromStreamDN #-}
133fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
134fromStreamDN limit str = do
135 marr <- liftIO $ newPrimArray (max limit 0)
136 _ <-
137 D.foldlM'
138 (\i x -> i `seq` liftIO (writePrimArray marr i x) >> return (i + 1))
139 (return 0) $
140 D.take limit str
141 liftIO $ unsafeFreezePrimArray marr
143{-# INLINE fromStreamD #-}
144fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
145fromStreamD = D.runFold write
147{-# INLINABLE fromListN #-}
148fromListN :: Prim a => Int -> [a] -> PrimArray a
149fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
151{-# INLINABLE fromList #-}
152fromList :: Prim a => [a] -> PrimArray a
153fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
155instance Prim a => NFData (PrimArray a) where
156 {-# INLINE rnf #-}
157 rnf = foldl' (\_ _ -> ()) ()
159{-# INLINE fromStreamN #-}
160fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
161fromStreamN n m = do
162 when (n < 0) $ error "fromStreamN: negative write count specified"
163 fromStreamDN n $ D.toStreamD m
165{-# INLINE fromStream #-}
166fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
167fromStream m = fromStreamD $ D.toStreamD m
169{-# INLINE_EARLY toStream #-}
170toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
171toStream = D.fromStreamD . toStreamD
173{-# INLINE_EARLY toStreamRev #-}
174toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
175toStreamRev = D.fromStreamD . toStreamDRev
177{-# INLINE fold #-}
178fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
179fold f arr = D.runFold f (toStreamD arr)
181{-# INLINE streamFold #-}
182streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
183streamFold f arr = f (toStream arr)
185{-# INLINE_NORMAL read #-}
186read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
187read = Unfold step inject
188 where
189 inject arr = return (arr, 0)
190 step (arr, i)
191 | i == length arr = return D.Stop
192 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
194{-# INLINE_NORMAL readSlice #-}
195readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
196readSlice off len = Unfold step inject
197 where
198 inject arr = return (arr, off)
199 step (arr, i)
200 | i == min (off + len) (length arr) = return D.Stop
201 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)