git.iko.soy/streamly
on 3be5560 help sign in
streamly/src/Streamly/Internal/Data/Prim/Array.hs
206 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 Data.Primitive.PrimArray
66import Streamly.Internal.Data.Unfold.Types (Unfold(..))
67import Streamly.Internal.Data.Fold.Types (Fold(..))
68import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
69import Streamly.Internal.Data.Stream.Serial (SerialT)
71import qualified Streamly.Internal.Data.Stream.StreamD as D
73{-# INLINE_NORMAL toStreamD #-}
74toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
75toStreamD arr = D.Stream step 0
76 where
77 {-# INLINE_LATE step #-}
78 step _ i
79 | i == sizeofPrimArray arr = return D.Stop
80 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
82{-# INLINE length #-}
83length :: Prim a => PrimArray a -> Int
84length arr = sizeofPrimArray arr
86{-# INLINE_NORMAL toStreamDRev #-}
87toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
88toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
89 where
90 {-# INLINE_LATE step #-}
91 step _ i
92 | i < 0 = return D.Stop
93 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
95{-# INLINE_NORMAL foldl' #-}
96foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
97foldl' = foldlPrimArray'
99{-# INLINE_NORMAL foldr #-}
100foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
101foldr = foldrPrimArray
103-- writeN n = S.evertM (fromStreamDN n)
104{-# INLINE_NORMAL writeN #-}
105writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
106writeN limit = Fold step initial extract
107 where
108 initial = do
109 marr <- liftIO $ newPrimArray limit
110 return (marr, 0)
111 step (marr, i) x
112 | i == limit = return (marr, i)
113 | otherwise = do
114 liftIO $ writePrimArray marr i x
115 return (marr, i + 1)
116 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
118{-# INLINE_NORMAL write #-}
119write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
120write = Fold step initial extract
121 where
122 initial = do
123 marr <- liftIO $ newPrimArray 0
124 return (marr, 0, 0)
125 step (marr, i, capacity) x
126 | i == capacity =
127 let newCapacity = max (capacity * 2) 1
128 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
129 liftIO $ writePrimArray newMarr i x
130 return (newMarr, i + 1, newCapacity)
131 | otherwise = do
132 liftIO $ writePrimArray marr i x
133 return (marr, i + 1, capacity)
134 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
135 liftIO $ unsafeFreezePrimArray marr
137{-# INLINE_NORMAL fromStreamDN #-}
138fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
139fromStreamDN limit str = do
140 marr <- liftIO $ newPrimArray (max limit 0)
141 _ <-
142 D.foldlM'
143 (\i x -> i `seq` (liftIO $ writePrimArray marr i x) >> return (i + 1))
144 0 $
145 D.take limit str
146 liftIO $ unsafeFreezePrimArray marr
148{-# INLINE fromStreamD #-}
149fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
150fromStreamD str = D.runFold write str
152{-# INLINABLE fromListN #-}
153fromListN :: Prim a => Int -> [a] -> PrimArray a
154fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
156{-# INLINABLE fromList #-}
157fromList :: Prim a => [a] -> PrimArray a
158fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
160instance Prim a => NFData (PrimArray a) where
161 {-# INLINE rnf #-}
162 rnf = foldl' (\_ _ -> ()) ()
164{-# INLINE fromStreamN #-}
165fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
166fromStreamN n m = do
167 when (n < 0) $ error "fromStreamN: negative write count specified"
168 fromStreamDN n $ D.toStreamD m
170{-# INLINE fromStream #-}
171fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
172fromStream m = fromStreamD $ D.toStreamD m
174{-# INLINE_EARLY toStream #-}
175toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
176toStream = D.fromStreamD . toStreamD
178{-# INLINE_EARLY toStreamRev #-}
179toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
180toStreamRev = D.fromStreamD . toStreamDRev
182{-# INLINE fold #-}
183fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
184fold f arr = D.runFold f (toStreamD arr)
186{-# INLINE streamFold #-}
187streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
188streamFold f arr = f (toStream arr)
190{-# INLINE_NORMAL read #-}
191read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
192read = Unfold step inject
193 where
194 inject arr = return (arr, 0)
195 step (arr, i)
196 | i == length arr = return D.Stop
197 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
199{-# INLINE_NORMAL readSlice #-}
200readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
201readSlice off len = Unfold step inject
202 where
203 inject arr = return (arr, off)
204 step (arr, i)
205 | i == min (off + len) (length arr) = return D.Stop
206 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)