git.iko.soy/streamly
on fab76f2 help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
194 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.Serial (SerialT(..))
58import qualified Streamly.Internal.Data.Stream.StreamD as D
59import qualified Streamly.Internal.Data.Fold.Type as FL
61{-# NOINLINE bottomElement #-}
62bottomElement :: a
63bottomElement = undefined
65{-# INLINE length #-}
66length :: SmallArray a -> Int
67length = sizeofSmallArray
69{-# INLINE_NORMAL toStreamD #-}
70toStreamD :: Monad m => SmallArray a -> D.Stream m a
71toStreamD arr = D.Stream step 0
72 where
73 {-# INLINE_LATE step #-}
74 step _ i
75 | i == length arr = return D.Stop
76 | otherwise =
77 return $
78 case indexSmallArray## arr i of
79 (# x #) -> D.Yield x (i + 1)
81{-# INLINE_NORMAL toStreamDRev #-}
82toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
83toStreamDRev arr = D.Stream step (length arr - 1)
84 where
85 {-# INLINE_LATE step #-}
86 step _ i
87 | i < 0 = return D.Stop
88 | otherwise =
89 return $
90 case indexSmallArray## arr i of
91 (# x #) -> D.Yield x (i - 1)
93{-# INLINE_NORMAL foldl' #-}
94foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
95foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
97{-# INLINE_NORMAL foldr #-}
98foldr :: (a -> b -> b) -> b -> SmallArray a -> b
99foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
101-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
102-- 'SmallArray'.
103--
104-- Since we are folding to a 'SmallArray' @n@ should be <= 128, for larger number
105-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
106{-# INLINE_NORMAL writeN #-}
107writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
108writeN len = Fold step initial extract
110 where
112 {-# INLINE next #-}
113 next marr i = do
114 let i1 = i + 1
115 st = Tuple' marr i1
116 if len > i1
117 then return $ FL.Partial st
118 else fmap FL.Done $ extract st
120 initial = do
121 marr <- liftIO $ newSmallArray len bottomElement
122 next marr (-1)
124 step (Tuple' marr i) x = do
125 liftIO $ writeSmallArray marr i x
126 next marr i
128 extract (Tuple' marr l) = liftIO $ freezeSmallArray marr 0 l
131{-# INLINE_NORMAL fromStreamDN #-}
132fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
133fromStreamDN limit str = do
134 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
135 i <-
136 D.foldlM'
137 (\i x -> i `seq` liftIO (writeSmallArray marr i x) >> return (i + 1))
138 (return 0) $
139 D.take limit str
140 liftIO $ freezeSmallArray marr 0 i
142-- | Create a 'SmallArray' from the first @n@ elements of a list. The
143-- array may hold less than @n@ elements if the length of the list <=
144-- @n@.
145--
146-- It is recommended to use a value of @n@ <= 128. For larger sized
147-- arrays, use an 'Array' from "Streamly.Data.Array" or
148-- "Streamly.Data.Array.Foreign"
149{-# INLINABLE fromListN #-}
150fromListN :: Int -> [a] -> SmallArray a
151fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
153instance NFData a => NFData (SmallArray a) where
154 {-# INLINE rnf #-}
155 rnf = foldl' (\_ x -> rnf x) ()
157-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
158-- array is allocated to size @n@, if the stream terminates before @n@
159-- elements then the array may hold less than @n@ elements.
160--
161-- For optimal performance use this with @n@ <= 128.
162{-# INLINE fromStreamN #-}
163fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
164fromStreamN n (SerialT m) = do
165 when (n < 0) $ error "fromStreamN: negative write count specified"
166 fromStreamDN n $ D.fromStreamK m
168{-# INLINE_EARLY toStream #-}
169toStream :: Monad m => SmallArray a -> SerialT m a
170toStream = SerialT . D.toStreamK . toStreamD
172{-# INLINE_EARLY toStreamRev #-}
173toStreamRev :: Monad m => SmallArray a -> SerialT m a
174toStreamRev = SerialT . D.toStreamK . toStreamDRev
176{-# INLINE fold #-}
177fold :: Monad m => Fold m a b -> SmallArray a -> m b
178fold f arr = D.fold f (toStreamD arr)
180{-# INLINE streamFold #-}
181streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
182streamFold f arr = f (toStream arr)
184{-# INLINE_NORMAL read #-}
185read :: Monad m => Unfold m (SmallArray a) a
186read = Unfold step inject
187 where
188 inject arr = return (arr, 0)
189 step (arr, i)
190 | i == length arr = return D.Stop
191 | otherwise =
192 return $
193 case indexSmallArray## arr i of
194 (# x #) -> D.Yield x (arr, i + 1)