git.iko.soy/streamly
on cee86ff help sign in
streamly/src/Streamly/Internal/Data/Prim/Array.hs
199 lines 5.9 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 ( PrimArray(..)
20 , Prim(..)
22 , foldl'
23 , foldr
25 , length
27 , writeN
28 , write
30 , toStreamD
31 , toStreamDRev
33 , toStream
34 , toStreamRev
35 , read
36 , readSlice
38 , fromListN
39 , fromList
40 , fromStreamDN
41 , fromStreamD
43 , fromStreamN
44 , fromStream
46 , streamFold
47 , fold
48 )
49where
51import Prelude hiding (foldr, length, read)
52import Control.DeepSeq (NFData(..))
53import Control.Monad (when)
54import Control.Monad.IO.Class (liftIO, MonadIO)
55import GHC.IO (unsafePerformIO)
56import Data.Primitive.PrimArray
57import Data.Primitive.Types (Prim(..))
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{-# INLINE_NORMAL toStreamD #-}
67toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
68toStreamD arr = D.Stream step 0
69 where
70 {-# INLINE_LATE step #-}
71 step _ i
72 | i == sizeofPrimArray arr = return D.Stop
73 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
75{-# INLINE length #-}
76length :: Prim a => PrimArray a -> Int
77length arr = sizeofPrimArray arr
79{-# INLINE_NORMAL toStreamDRev #-}
80toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
81toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
82 where
83 {-# INLINE_LATE step #-}
84 step _ i
85 | i < 0 = return D.Stop
86 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
88{-# INLINE_NORMAL foldl' #-}
89foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
90foldl' = foldlPrimArray'
92{-# INLINE_NORMAL foldr #-}
93foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
94foldr = foldrPrimArray
96-- writeN n = S.evertM (fromStreamDN n)
97{-# INLINE_NORMAL writeN #-}
98writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
99writeN limit = Fold step initial extract
100 where
101 initial = do
102 marr <- liftIO $ newPrimArray limit
103 return (marr, 0)
104 step (marr, i) x
105 | i == limit = return (marr, i)
106 | otherwise = do
107 liftIO $ writePrimArray marr i x
108 return (marr, i + 1)
109 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
111{-# INLINE_NORMAL write #-}
112write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
113write = Fold step initial extract
114 where
115 initial = do
116 marr <- liftIO $ newPrimArray 0
117 return (marr, 0, 0)
118 step (marr, i, capacity) x
119 | i == capacity =
120 let newCapacity = max (capacity * 2) 1
121 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
122 liftIO $ writePrimArray newMarr i x
123 return (newMarr, i + 1, newCapacity)
124 | otherwise = do
125 liftIO $ writePrimArray marr i x
126 return (marr, i + 1, capacity)
127 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
128 liftIO $ unsafeFreezePrimArray marr
130{-# INLINE_NORMAL fromStreamDN #-}
131fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
132fromStreamDN limit str = do
133 marr <- liftIO $ newPrimArray (max limit 0)
134 _ <-
135 D.foldlM'
136 (\i x -> i `seq` (liftIO $ writePrimArray marr i x) >> return (i + 1))
137 0 $
138 D.take limit str
139 liftIO $ unsafeFreezePrimArray marr
141{-# INLINE fromStreamD #-}
142fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
143fromStreamD str = D.runFold write str
145{-# INLINABLE fromListN #-}
146fromListN :: Prim a => Int -> [a] -> PrimArray a
147fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
149{-# INLINABLE fromList #-}
150fromList :: Prim a => [a] -> PrimArray a
151fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
153instance Prim a => NFData (PrimArray a) where
154 {-# INLINE rnf #-}
155 rnf = foldl' (\_ _ -> ()) ()
157{-# INLINE fromStreamN #-}
158fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
159fromStreamN n m = do
160 when (n < 0) $ error "fromStreamN: negative write count specified"
161 fromStreamDN n $ D.toStreamD m
163{-# INLINE fromStream #-}
164fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
165fromStream m = fromStreamD $ D.toStreamD m
167{-# INLINE_EARLY toStream #-}
168toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
169toStream = D.fromStreamD . toStreamD
171{-# INLINE_EARLY toStreamRev #-}
172toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
173toStreamRev = D.fromStreamD . toStreamDRev
175{-# INLINE fold #-}
176fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
177fold f arr = D.runFold f (toStreamD arr)
179{-# INLINE streamFold #-}
180streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
181streamFold f arr = f (toStream arr)
183{-# INLINE_NORMAL read #-}
184read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
185read = Unfold step inject
186 where
187 inject arr = return (arr, 0)
188 step (arr, i)
189 | i == length arr = return D.Stop
190 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)
192{-# INLINE_NORMAL readSlice #-}
193readSlice :: (Prim a, Monad m) => Int -> Int -> Unfold m (PrimArray a) a
194readSlice off len = Unfold step inject
195 where
196 inject arr = return (arr, off)
197 step (arr, i)
198 | i == min (off + len) (length arr) = return D.Stop
199 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)