git.iko.soy/streamly
on 602011f help sign in
streamly/src/Streamly/Internal/Data/PrimArray.hs
191 lines 5.6 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(..)
20 , Prim(..)
22 , foldl'
23 , foldr
25 , length
27 , writeN
28 , write
30 , toStreamD
31 , toStreamDRev
33 , toStream
34 , toStreamRev
35 , read
37 , fromListN
38 , fromList
39 , fromStreamDN
40 , fromStreamD
42 , fromStreamN
43 , fromStream
45 , streamFold
46 , fold
47 )
48where
50import Prelude hiding (foldr, length, read)
51import Control.DeepSeq (NFData(..))
52import Control.Monad (when)
53import Control.Monad.IO.Class (liftIO, MonadIO)
54import GHC.IO (unsafePerformIO)
55import GHC.Base (Int(..))
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)