git.iko.soy/streamly
on 4008a26 help sign in
streamly/src/Streamly/Internal/Data/PrimArray.hs
190 lines 5.5 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.PrimArray
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.PrimArray
19 ( PrimArray(..)
21 , foldl'
22 , foldr
24 , length
26 , writeN
27 , write
29 , toStreamD
30 , toStreamDRev
32 , toStream
33 , toStreamRev
34 , read
36 , fromListN
37 , fromList
38 , fromStreamDN
39 , fromStreamD
41 , fromStreamN
42 , fromStream
44 , streamFold
45 , fold
46 )
47where
49import Prelude hiding (foldr, length, read)
50import Control.DeepSeq (NFData(..))
51import Control.Monad (when)
52import Control.Monad.IO.Class (liftIO, MonadIO)
53import GHC.IO (unsafePerformIO)
54import GHC.Base (Int(..))
55import Data.Primitive.PrimArray
56import Data.Primitive.Types (Prim(..))
58import Streamly.Internal.Data.Unfold.Types (Unfold(..))
59import Streamly.Internal.Data.Fold.Types (Fold(..))
60import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
61import Streamly.Internal.Data.Stream.Serial (SerialT)
63import qualified Streamly.Internal.Data.Stream.StreamD as D
65{-# INLINE_NORMAL toStreamD #-}
66toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
67toStreamD arr = D.Stream step 0
68 where
69 {-# INLINE_LATE step #-}
70 step _ i
71 | i == sizeofPrimArray arr = return D.Stop
72 step _ i = return $ D.Yield (indexPrimArray arr i) (i + 1)
74{-# INLINE length #-}
75length :: Prim a => PrimArray a -> Int
76length arr = sizeofPrimArray arr
78{-# INLINE_NORMAL toStreamDRev #-}
79toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
80toStreamDRev arr = D.Stream step (sizeofPrimArray arr - 1)
81 where
82 {-# INLINE_LATE step #-}
83 step _ i
84 | i < 0 = return D.Stop
85 step _ i = return $ D.Yield (indexPrimArray arr i) (i - 1)
87{-# INLINE_NORMAL foldl' #-}
88foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
89foldl' = foldlPrimArray'
91{-# INLINE_NORMAL foldr #-}
92foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
93foldr = foldrPrimArray
95-- writeN n = S.evertM (fromStreamDN n)
96{-# INLINE_NORMAL writeN #-}
97writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
98writeN limit = Fold step initial extract
99 where
100 initial = do
101 marr <- liftIO $ newPrimArray limit
102 return (marr, 0)
103 step (marr, i) x
104 | i == limit = return (marr, i)
105 | otherwise = do
106 liftIO $ writePrimArray marr i x
107 return (marr, i + 1)
108 extract (marr, _) = liftIO $ unsafeFreezePrimArray marr
110{-# INLINE_NORMAL write #-}
111write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
112write = Fold step initial extract
113 where
114 initial = do
115 marr <- liftIO $ newPrimArray 0
116 return (marr, 0, 0)
117 step (marr, i, capacity) x
118 | i == capacity =
119 let newCapacity = max (capacity * 2) 1
120 in do newMarr <- liftIO $ resizeMutablePrimArray marr newCapacity
121 liftIO $ writePrimArray newMarr i x
122 return (newMarr, i + 1, newCapacity)
123 | otherwise = do
124 liftIO $ writePrimArray marr i x
125 return (marr, i + 1, capacity)
126 extract (marr, len, _) = do liftIO $ shrinkMutablePrimArray marr len
127 liftIO $ unsafeFreezePrimArray marr
129{-# INLINE_NORMAL fromStreamDN #-}
130fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
131fromStreamDN limit str = do
132 marr <- liftIO $ newPrimArray (max limit 0)
133 _ <-
134 D.foldlM'
135 (\i x -> i `seq` (liftIO $ writePrimArray marr i x) >> return (i + 1))
136 0 $
137 D.take limit str
138 liftIO $ unsafeFreezePrimArray marr
140{-# INLINE fromStreamD #-}
141fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
142fromStreamD str = D.runFold write str
144{-# INLINABLE fromListN #-}
145fromListN :: Prim a => Int -> [a] -> PrimArray a
146fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
148{-# INLINABLE fromList #-}
149fromList :: Prim a => [a] -> PrimArray a
150fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
152instance NFData (PrimArray a) where
153 {-# INLINE rnf #-}
154 rnf = const ()
156{-# INLINE fromStreamN #-}
157fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
158fromStreamN n m = do
159 when (n < 0) $ error "fromStreamN: negative write count specified"
160 fromStreamDN n $ D.toStreamD m
162{-# INLINE fromStream #-}
163fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
164fromStream m = fromStreamD $ D.toStreamD m
166{-# INLINE_EARLY toStream #-}
167toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
168toStream = D.fromStreamD . toStreamD
170{-# INLINE_EARLY toStreamRev #-}
171toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
172toStreamRev = D.fromStreamD . toStreamDRev
174{-# INLINE fold #-}
175fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
176fold f arr = D.runFold f (toStreamD arr)
178{-# INLINE streamFold #-}
179streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
180streamFold f arr = f (toStream arr)
182{-# INLINE_NORMAL read #-}
183read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
184read = Unfold step inject
185 where
186 inject arr = return (arr, 0)
187 step (arr, i)
188 | i == length arr = return D.Stop
189 step (arr, i) = return $ D.Yield (indexPrimArray arr i) (arr, i + 1)