git.iko.soy/streamly
on 88c2aed help sign in
streamly/src/Streamly/Internal/Data/Prim/Array.hs
203 lines 6.0 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE CPP #-}
5#include "inline.hs"
7-- |
8-- Module : Streamly.Internal.Data.Prim.Array
9-- Copyright : (c) 2019 Composewell Technologies
10--
11-- License : BSD-3-Clause
12-- Maintainer : streamly@composewell.com
13-- Stability : experimental
14-- Portability : GHC
15--
16module Streamly.Internal.Data.Prim.Array
17 (
19 -- XXX should it be just Array instead? We should be able to replace one
20 -- array type with another easily.
21 PrimArray(..)
23 -- XXX Prim should be exported from Data.Prim module?
24 , Prim(..)
26 , foldl'
27 , foldr
29 , length
31 , writeN
32 , write
34 , toStreamD
35 , toStreamDRev
37 , toStream
38 , toStreamRev
39 , read
40 , readSlice
42 , fromListN
43 , fromList
44 , fromStreamDN
45 , fromStreamD
47 , fromStreamN
48 , fromStream
50 , streamFold
51 , fold
52 )
53where
55import Prelude hiding (foldr, length, read)
56import Control.DeepSeq (NFData(..))
57import Control.Monad (when)
58import Control.Monad.IO.Class (liftIO, MonadIO)
59import GHC.IO (unsafePerformIO)
60import Data.Primitive.Types (Prim(..))
62import Streamly.Internal.Data.Prim.Array.Types
63import Streamly.Internal.Data.Unfold.Types (Unfold(..))
64import Streamly.Internal.Data.Fold.Types (Fold(..))
65import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
66import Streamly.Internal.Data.Stream.Serial (SerialT)
68import qualified Streamly.Internal.Data.Stream.StreamD as D
70{-# INLINE_NORMAL toStreamD #-}
71toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
72toStreamD arr = D.Stream step 0
73 where
74 {-# INLINE_LATE step #-}
75 step _ i
76 | i == sizeofPrimArray arr = return D.Stop
77 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
79{-# INLINE length #-}
80length :: Prim a => PrimArray a -> Int
81length = sizeofPrimArray
83{-# INLINE_NORMAL toStreamDRev #-}
84toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
85toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
86 where
87 {-# INLINE_LATE step #-}
88 step _ i
89 | i < 0 = return D.Stop
90 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
92{-# INLINE_NORMAL foldl' #-}
93foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
94foldl' = foldlPrimArray'
96{-# INLINE_NORMAL foldr #-}
97foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
98foldr = foldrPrimArray
100-- writeN n = S.evertM (fromStreamDN n)
101{-# INLINE_NORMAL writeN #-}
102writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
103writeN limit = Fold step initial extract
104 where
105 initial = do
106 marr <- liftIO $ newPrimArray limit
107 return (marr, 0)
108 step (marr, i) x
109 | i == limit = return (marr, i)
110 | otherwise = do
111 liftIO $ writePrimArray marr i x
112 return (marr, i + 1)
113 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
115{-# INLINE_NORMAL write #-}
116write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
117write = Fold step initial extract
118 where
119 initial = do
120 marr <- liftIO $ newPrimArray 0
121 return (marr, 0, 0)
122 step (marr, i, capacity) x
123 | i == capacity =
124 let newCapacity = max (capacity * 2) 1
125 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
126 liftIO $ writePrimArray newMarr i x
127 return (newMarr, i + 1, newCapacity)
128 | otherwise = do
129 liftIO $ writePrimArray marr i x
130 return (marr, i + 1, capacity)
131 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
132 liftIO $ unsafeFreezePrimArray marr
134{-# INLINE_NORMAL fromStreamDN #-}
135fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
136fromStreamDN limit str = do
137 marr <- liftIO $ newPrimArray (max limit 0)
138 _ <-
139 D.foldlM'
140 (\i x -> i `seq` liftIO (writePrimArray marr i x) >> return (i + 1))
141 0 $
142 D.take limit str
143 liftIO $ unsafeFreezePrimArray marr
145{-# INLINE fromStreamD #-}
146fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
147fromStreamD = D.runFold write
149{-# INLINABLE fromListN #-}
150fromListN :: Prim a => Int -> [a] -> PrimArray a
151fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
153{-# INLINABLE fromList #-}
154fromList :: Prim a => [a] -> PrimArray a
155fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
157instance Prim a => NFData (PrimArray a) where
158 {-# INLINE rnf #-}
159 rnf = foldl' (\_ _ -> ()) ()
161{-# INLINE fromStreamN #-}
162fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
163fromStreamN n m = do
164 when (n < 0) $ error "fromStreamN: negative write count specified"
165 fromStreamDN n $ D.toStreamD m
167{-# INLINE fromStream #-}
168fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
169fromStream m = fromStreamD $ D.toStreamD m
171{-# INLINE_EARLY toStream #-}
172toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
173toStream = D.fromStreamD . toStreamD
175{-# INLINE_EARLY toStreamRev #-}
176toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
177toStreamRev = D.fromStreamD . toStreamDRev
179{-# INLINE fold #-}
180fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
181fold f arr = D.runFold f (toStreamD arr)
183{-# INLINE streamFold #-}
184streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
185streamFold f arr = f (toStream arr)
187{-# INLINE_NORMAL read #-}
188read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
189read = Unfold step inject
190 where
191 inject arr = return (arr, 0)
192 step (arr, i)
193 | i == length arr = return D.Stop
194 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
196{-# INLINE_NORMAL readSlice #-}
197readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
198readSlice off len = Unfold step inject
199 where
200 inject arr = return (arr, off)
201 step (arr, i)
202 | i == min (off + len) (length arr) = return D.Stop
203 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)