@@ -0,0 +1,190 @@
1+{-# OPTIONS_GHC -fno-warn-orphans #-}
2+
3+{-# LANGUAGE CPP #-}
4+{-# LANGUAGE MagicHash #-}
5+{-# LANGUAGE UnboxedTuples #-}
6+
7+#include "inline.hs"
8+
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+--
18+module Streamly.Internal.Data.PrimArray
19+ ( PrimArray(..)
20+
21+ , foldl'
22+ , foldr
23+
24+ , length
25+
26+ , writeN
27+ , write
28+
29+ , toStreamD
30+ , toStreamDRev
31+
32+ , toStream
33+ , toStreamRev
34+ , read
35+
36+ , fromListN
37+ , fromList
38+ , fromStreamDN
39+ , fromStreamD
40+
41+ , fromStreamN
42+ , fromStream
43+
44+ , streamFold
45+ , fold
46+ )
47+where
48+
49+import Prelude hiding (foldr, length, read)
50+import Control.DeepSeq (NFData(..))
51+import Control.Monad (when)
52+import Control.Monad.IO.Class (liftIO, MonadIO)
53+import GHC.IO (unsafePerformIO)
54+import GHC.Base (Int(..))
55+import Data.Primitive.PrimArray
56+import Data.Primitive.Types (Prim(..))
57+
58+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
59+import Streamly.Internal.Data.Fold.Types (Fold(..))
60+import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
61+import Streamly.Internal.Data.Stream.Serial (SerialT)
62+
63+import qualified Streamly.Internal.Data.Stream.StreamD as D
64+
65+{-# INLINE_NORMAL toStreamD #-}
66+toStreamD :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
67+toStreamD 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)
73+
74+{-# INLINE length #-}
75+length :: Prim a => PrimArray a -> Int
76+length arr = sizeofPrimArray arr
77+
78+{-# INLINE_NORMAL toStreamDRev #-}
79+toStreamDRev :: (Prim a, Monad m) => PrimArray a -> D.Stream m a
80+toStreamDRev 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)
86+
87+{-# INLINE_NORMAL foldl' #-}
88+foldl' :: Prim a => (b -> a -> b) -> b -> PrimArray a -> b
89+foldl' = foldlPrimArray'
90+
91+{-# INLINE_NORMAL foldr #-}
92+foldr :: Prim a => (a -> b -> b) -> b -> PrimArray a -> b
93+foldr = foldrPrimArray
94+
95+-- writeN n = S.evertM (fromStreamDN n)
96+{-# INLINE_NORMAL writeN #-}
97+writeN :: (MonadIO m, Prim a) => Int -> Fold m a (PrimArray a)
98+writeN 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
109+
110+{-# INLINE_NORMAL write #-}
111+write :: (MonadIO m, Prim a) => Fold m a (PrimArray a)
112+write = 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
128+
129+{-# INLINE_NORMAL fromStreamDN #-}
130+fromStreamDN :: (MonadIO m, Prim a) => Int -> D.Stream m a -> m (PrimArray a)
131+fromStreamDN 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
139+
140+{-# INLINE fromStreamD #-}
141+fromStreamD :: (MonadIO m, Prim a) => D.Stream m a -> m (PrimArray a)
142+fromStreamD str = D.runFold write str
143+
144+{-# INLINABLE fromListN #-}
145+fromListN :: Prim a => Int -> [a] -> PrimArray a
146+fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
147+
148+{-# INLINABLE fromList #-}
149+fromList :: Prim a => [a] -> PrimArray a
150+fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
151+
152+instance NFData (PrimArray a) where
153+ {-# INLINE rnf #-}
154+ rnf = const ()
155+
156+{-# INLINE fromStreamN #-}
157+fromStreamN :: (MonadIO m, Prim a) => Int -> SerialT m a -> m (PrimArray a)
158+fromStreamN n m = do
159+ when (n < 0) $ error "fromStreamN: negative write count specified"
160+ fromStreamDN n $ D.toStreamD m
161+
162+{-# INLINE fromStream #-}
163+fromStream :: (MonadIO m, Prim a) => SerialT m a -> m (PrimArray a)
164+fromStream m = fromStreamD $ D.toStreamD m
165+
166+{-# INLINE_EARLY toStream #-}
167+toStream :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
168+toStream = D.fromStreamD . toStreamD
169+
170+{-# INLINE_EARLY toStreamRev #-}
171+toStreamRev :: (Prim a, Monad m, IsStream t) => PrimArray a -> t m a
172+toStreamRev = D.fromStreamD . toStreamDRev
173+
174+{-# INLINE fold #-}
175+fold :: (Prim a, Monad m) => Fold m a b -> PrimArray a -> m b
176+fold f arr = D.runFold f (toStreamD arr)
177+
178+{-# INLINE streamFold #-}
179+streamFold :: (Prim a, Monad m) => (SerialT m a -> m b) -> PrimArray a -> m b
180+streamFold f arr = f (toStream arr)
181+
182+{-# INLINE_NORMAL read #-}
183+read :: (Prim a, Monad m) => Unfold m (PrimArray a) a
184+read = 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)
190+