git.iko.soy/streamly
on cf63b4e help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
188 lines 5.6 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 #-}
11{-# LANGUAGE UnboxedTuples #-}
13#include "inline.hs"
15module Streamly.Internal.Data.SmallArray
16 (
17 -- XXX should it be just Array instead? We should be able to replace one
18 -- array type with another easily.
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)
51import Streamly.Internal.Data.SmallArray.Type
53import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
54import Streamly.Internal.Data.Unfold.Type (Unfold(..))
55import Streamly.Internal.Data.Fold.Type (Fold(..))
56import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
57import Streamly.Internal.Data.Stream.Serial (SerialT)
59import qualified Streamly.Internal.Data.Stream.StreamD as D
60import qualified Streamly.Internal.Data.Fold.Type as FL
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.Data.Array.Foreign".
107{-# INLINE_NORMAL writeN #-}
108writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
109writeN limit = Fold step initial extract
111 where
113 initial = do
114 marr <- liftIO $ newSmallArray limit bottomElement
115 return $ FL.Partial (Tuple' marr 0)
117 step st@(Tuple' marr i) x
118 | i == limit = FL.Done <$> extract st
119 | otherwise = do
120 liftIO $ writeSmallArray marr i x
121 return $ FL.Partial (Tuple' marr (i + 1))
123 extract (Tuple' marr len) = liftIO $ freezeSmallArray marr 0 len
125{-# INLINE_NORMAL fromStreamDN #-}
126fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
127fromStreamDN limit str = do
128 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
129 i <-
130 D.foldlM'
131 (\i x -> i `seq` liftIO (writeSmallArray marr i x) >> return (i + 1))
132 (return 0) $
133 D.take limit str
134 liftIO $ freezeSmallArray marr 0 i
136-- | Create a 'SmallArray' from the first @n@ elements of a list. The
137-- array may hold less than @n@ elements if the length of the list <=
138-- @n@.
139--
140-- It is recommended to use a value of @n@ <= 128. For larger sized
141-- arrays, use an 'Array' from "Streamly.Data.Array" or
142-- "Streamly.Data.Array.Foreign"
143{-# INLINABLE fromListN #-}
144fromListN :: Int -> [a] -> SmallArray a
145fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
147instance NFData a => NFData (SmallArray a) where
148 {-# INLINE rnf #-}
149 rnf = foldl' (\_ x -> rnf x) ()
151-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
152-- array is allocated to size @n@, if the stream terminates before @n@
153-- elements then the array may hold less than @n@ elements.
154--
155-- For optimal performance use this with @n@ <= 128.
156{-# INLINE fromStreamN #-}
157fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
158fromStreamN n m = do
159 when (n < 0) $ error "fromStreamN: negative write count specified"
160 fromStreamDN n $ D.toStreamD m
162{-# INLINE_EARLY toStream #-}
163toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
164toStream = D.fromStreamD . toStreamD
166{-# INLINE_EARLY toStreamRev #-}
167toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
168toStreamRev = D.fromStreamD . toStreamDRev
170{-# INLINE fold #-}
171fold :: Monad m => Fold m a b -> SmallArray a -> m b
172fold f arr = D.fold f (toStreamD arr)
174{-# INLINE streamFold #-}
175streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
176streamFold f arr = f (toStream arr)
178{-# INLINE_NORMAL read #-}
179read :: Monad m => Unfold m (SmallArray a) a
180read = Unfold step inject
181 where
182 inject arr = return (arr, 0)
183 step (arr, i)
184 | i == length arr = return D.Stop
185 | otherwise =
186 return $
187 case indexSmallArray## arr i of
188 (# x #) -> D.Yield x (arr, i + 1)