git.iko.soy/streamly
on 88c2aed help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
184 lines 5.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 (
20 -- XXX should it be just Array instead? We should be able to replace one
21 -- array type with another easily.
22 SmallArray(..)
24 , foldl'
25 , foldr
27 , length
29 , writeN
31 , toStreamD
32 , toStreamDRev
34 , toStream
35 , toStreamRev
36 , read
38 , fromListN
39 , fromStreamDN
40 , fromStreamN
42 , streamFold
43 , fold
44 )
45where
47import Prelude hiding (foldr, length, read)
48import Control.DeepSeq (NFData(..))
49import Control.Monad (when)
50import Control.Monad.IO.Class (MonadIO, liftIO)
51import GHC.IO (unsafePerformIO)
52import Data.Functor.Identity (runIdentity)
54import Streamly.Internal.Data.SmallArray.Types
55import Streamly.Internal.Data.Unfold.Types (Unfold(..))
56import Streamly.Internal.Data.Fold.Types (Fold(..))
57import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
58import Streamly.Internal.Data.Stream.Serial (SerialT)
60import qualified Streamly.Internal.Data.Stream.StreamD as D
62{-# NOINLINE bottomElement #-}
63bottomElement :: a
64bottomElement = undefined
66{-# INLINE length #-}
67length :: SmallArray a -> Int
68length = sizeofSmallArray
70{-# INLINE_NORMAL toStreamD #-}
71toStreamD :: Monad m => SmallArray a -> D.Stream m a
72toStreamD arr = D.Stream step 0
73 where
74 {-# INLINE_LATE step #-}
75 step _ i
76 | i == length arr = return D.Stop
77 | otherwise =
78 return $
79 case indexSmallArray## arr i of
80 (# x #) -> D.Yield x (i + 1)
82{-# INLINE_NORMAL toStreamDRev #-}
83toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
84toStreamDRev arr = D.Stream step (length arr - 1)
85 where
86 {-# INLINE_LATE step #-}
87 step _ i
88 | i < 0 = return D.Stop
89 | otherwise =
90 return $
91 case indexSmallArray## arr i of
92 (# x #) -> D.Yield x (i - 1)
94{-# INLINE_NORMAL foldl' #-}
95foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
96foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
98{-# INLINE_NORMAL foldr #-}
99foldr :: (a -> b -> b) -> b -> SmallArray a -> b
100foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
102-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
103-- 'SmallArray'.
104--
105-- Since we are folding to a 'SmallArray' @n@ should be <= 128, for larger number
106-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Memory.Array".
107{-# INLINE_NORMAL writeN #-}
108writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
109writeN limit = Fold step initial extract
110 where
111 initial = do
112 marr <- liftIO $ newSmallArray limit bottomElement
113 return (marr, 0)
114 step (marr, i) x
115 | i == limit = return (marr, i)
116 | otherwise = do
117 liftIO $ writeSmallArray marr i x
118 return (marr, i + 1)
119 extract (marr, len) = liftIO $ freezeSmallArray marr 0 len
121{-# INLINE_NORMAL fromStreamDN #-}
122fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
123fromStreamDN limit str = do
124 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
125 i <-
126 D.foldlM'
127 (\i x -> i `seq` liftIO (writeSmallArray marr i x) >> return (i + 1))
128 0 $
129 D.take limit str
130 liftIO $ freezeSmallArray marr 0 i
132-- | Create a 'SmallArray' from the first @n@ elements of a list. The
133-- array may hold less than @n@ elements if the length of the list <=
134-- @n@.
135--
136-- It is recommended to use a value of @n@ <= 128. For larger sized
137-- arrays, use an 'Array' from "Streamly.Data.Array" or
138-- "Streamly.Memory.Array"
139{-# INLINABLE fromListN #-}
140fromListN :: Int -> [a] -> SmallArray a
141fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
143instance NFData a => NFData (SmallArray a) where
144 {-# INLINE rnf #-}
145 rnf = foldl' (\_ x -> rnf x) ()
147-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
148-- array is allocated to size @n@, if the stream terminates before @n@
149-- elements then the array may hold less than @n@ elements.
150--
151-- For optimal performance use this with @n@ <= 128.
152{-# INLINE fromStreamN #-}
153fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
154fromStreamN n m = do
155 when (n < 0) $ error "fromStreamN: negative write count specified"
156 fromStreamDN n $ D.toStreamD m
158{-# INLINE_EARLY toStream #-}
159toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
160toStream = D.fromStreamD . toStreamD
162{-# INLINE_EARLY toStreamRev #-}
163toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
164toStreamRev = D.fromStreamD . toStreamDRev
166{-# INLINE fold #-}
167fold :: Monad m => Fold m a b -> SmallArray a -> m b
168fold f arr = D.runFold f (toStreamD arr)
170{-# INLINE streamFold #-}
171streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
172streamFold f arr = f (toStream arr)
174{-# INLINE_NORMAL read #-}
175read :: Monad m => Unfold m (SmallArray a) a
176read = Unfold step inject
177 where
178 inject arr = return (arr, 0)
179 step (arr, i)
180 | i == length arr = return D.Stop
181 | otherwise =
182 return $
183 case indexSmallArray## arr i of
184 (# x #) -> D.Yield x (arr, i + 1)