git.iko.soy/streamly
on ee32239 help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
165 lines 4.5 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.SmallArray
3-- Copyright : (c) 2019 Composewell Technologies
4--
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
10{-# OPTIONS_GHC -fno-warn-orphans #-}
12{-# LANGUAGE CPP #-}
13{-# LANGUAGE MagicHash #-}
14{-# LANGUAGE UnboxedTuples #-}
16#include "inline.hs"
18module Streamly.Internal.Data.SmallArray
19 ( SmallArray(..)
21 , foldl'
22 , foldr
24 , length
26 , writeN
28 , toStreamD
29 , toStreamDRev
31 , toStream
32 , toStreamRev
33 , read
35 , fromListN
36 , fromStreamDN
37 , fromStreamN
39 , streamFold
40 , fold
41 )
42where
44import Prelude hiding (foldr, length, read)
45import Control.DeepSeq (NFData(..))
46import Control.Monad (when)
47import Control.Monad.IO.Class (MonadIO, liftIO)
48import GHC.IO (unsafePerformIO)
49import Data.Functor.Identity (runIdentity)
50import Data.Primitive.SmallArray
52import Streamly.Internal.Data.Unfold.Types (Unfold(..))
53import Streamly.Internal.Data.Fold.Types (Fold(..))
54import Streamly.Streams.StreamK.Type (IsStream)
55import Streamly.Streams.Serial (SerialT)
57import qualified Streamly.Streams.StreamD as D
59{-# NOINLINE bottomElement #-}
60bottomElement :: a
61bottomElement = undefined
63{-# INLINE length #-}
64length :: SmallArray a -> Int
65length arr = sizeofSmallArray arr
67{-# INLINE_NORMAL toStreamD #-}
68toStreamD :: Monad m => SmallArray a -> D.Stream m a
69toStreamD arr = D.Stream step 0
70 where
71 {-# INLINE_LATE step #-}
72 step _ i
73 | i == length arr = return D.Stop
74 | otherwise =
75 return $
76 case indexSmallArray## arr i of
77 (# x #) -> D.Yield x (i + 1)
79{-# INLINE_NORMAL toStreamDRev #-}
80toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
81toStreamDRev arr = D.Stream step (length arr - 1)
82 where
83 {-# INLINE_LATE step #-}
84 step _ i
85 | i < 0 = return D.Stop
86 | otherwise =
87 return $
88 case indexSmallArray## arr i of
89 (# x #) -> D.Yield x (i - 1)
91{-# INLINE_NORMAL foldl' #-}
92foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
93foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
95{-# INLINE_NORMAL foldr #-}
96foldr :: (a -> b -> b) -> b -> SmallArray a -> b
97foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
99-- writeN n = S.evertM (fromStreamDN n)
100{-# INLINE_NORMAL writeN #-}
101writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
102writeN limit = Fold step initial extract
103 where
104 initial = do
105 marr <- liftIO $ newSmallArray limit bottomElement
106 return (marr, 0)
107 step (marr, i) x
108 | i == limit = return (marr, i)
109 | otherwise = do
110 liftIO $ writeSmallArray marr i x
111 return (marr, i + 1)
112 extract (marr, len) = liftIO $ freezeSmallArray marr 0 len
114{-# INLINE_NORMAL fromStreamDN #-}
115fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
116fromStreamDN limit str = do
117 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
118 i <-
119 D.foldlM'
120 (\i x -> i `seq` (liftIO $ writeSmallArray marr i x) >> return (i + 1))
121 0 $
122 D.take limit str
123 liftIO $ freezeSmallArray marr 0 i
125{-# INLINABLE fromListN #-}
126fromListN :: Int -> [a] -> SmallArray a
127fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
129instance NFData a => NFData (SmallArray a) where
130 {-# INLINE rnf #-}
131 rnf = foldl' (\_ x -> rnf x) ()
133{-# INLINE fromStreamN #-}
134fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
135fromStreamN n m = do
136 when (n < 0) $ error "fromStreamN: negative write count specified"
137 fromStreamDN n $ D.toStreamD m
139{-# INLINE_EARLY toStream #-}
140toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
141toStream = D.fromStreamD . toStreamD
143{-# INLINE_EARLY toStreamRev #-}
144toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
145toStreamRev = D.fromStreamD . toStreamDRev
147{-# INLINE fold #-}
148fold :: Monad m => Fold m a b -> SmallArray a -> m b
149fold f arr = D.runFold f (toStreamD arr)
151{-# INLINE streamFold #-}
152streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
153streamFold f arr = f (toStream arr)
155{-# INLINE_NORMAL read #-}
156read :: Monad m => Unfold m (SmallArray a) a
157read = Unfold step inject
158 where
159 inject arr = return (arr, 0)
160 step (arr, i)
161 | i == length arr = return D.Stop
162 | otherwise =
163 return $
164 case indexSmallArray## arr i of
165 (# x #) -> D.Yield x (arr, i + 1)