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