git.iko.soy/streamly
on master help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
195 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 #-}
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 , read
32 , readRev
33 , reader
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 (Unfold(..))
56import qualified Streamly.Internal.Data.Stream as D
57import qualified Streamly.Internal.Data.Fold as FL
58 (Fold (..), Step (Done, Partial))
60import Streamly.Internal.Data.SmallArray.Type
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.Generic" or "Streamly.Data.Array.Foreign".
107{-# INLINE_NORMAL writeN #-}
108writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
109writeN len = FL.Fold step initial extract extract
111 where
113 {-# INLINE next #-}
114 next marr i = do
115 let i1 = i + 1
116 st = Tuple' marr i1
117 if len > i1
118 then return $ FL.Partial st
119 else fmap FL.Done $ extract st
121 initial = do
122 marr <- liftIO $ newSmallArray len bottomElement
123 next marr (-1)
125 step (Tuple' marr i) x = do
126 liftIO $ writeSmallArray marr i x
127 next marr i
129 extract (Tuple' marr l) = liftIO $ freezeSmallArray marr 0 l
132{-# INLINE_NORMAL fromStreamDN #-}
133fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
134fromStreamDN limit str = do
135 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
136 i <-
137 D.foldlM'
138 (\i x -> i `seq` liftIO (writeSmallArray marr i x) >> return (i + 1))
139 (return 0) $
140 D.take limit str
141 liftIO $ freezeSmallArray marr 0 i
143-- | Create a 'SmallArray' from the first @n@ elements of a list. The
144-- array may hold less than @n@ elements if the length of the list <=
145-- @n@.
146--
147-- It is recommended to use a value of @n@ <= 128. For larger sized
148-- arrays, use an 'Array' from "Streamly.Data.Array" or
149-- "Streamly.Data.Array.Foreign"
150{-# INLINABLE fromListN #-}
151fromListN :: Int -> [a] -> SmallArray a
152fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
154instance NFData a => NFData (SmallArray a) where
155 {-# INLINE rnf #-}
156 rnf = foldl' (\_ x -> rnf x) ()
158-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
159-- array is allocated to size @n@, if the stream terminates before @n@
160-- elements then the array may hold less than @n@ elements.
161--
162-- For optimal performance use this with @n@ <= 128.
163{-# INLINE fromStreamN #-}
164fromStreamN :: MonadIO m => Int -> Stream m a -> m (SmallArray a)
165fromStreamN n m = do
166 when (n < 0) $ error "fromStreamN: negative write count specified"
167 fromStreamDN n m
169{-# INLINE_EARLY read #-}
170read :: Monad m => SmallArray a -> Stream m a
171read = toStreamD
173{-# INLINE_EARLY readRev #-}
174readRev :: Monad m => SmallArray a -> Stream m a
175readRev = toStreamDRev
177{-# INLINE fold #-}
178fold :: Monad m => Fold m a b -> SmallArray a -> m b
179fold f arr = D.fold f (toStreamD arr)
181{-# INLINE streamFold #-}
182streamFold :: Monad m => (Stream m a -> m b) -> SmallArray a -> m b
183streamFold f arr = f (read arr)
185{-# INLINE_NORMAL reader #-}
186reader :: Monad m => Unfold m (SmallArray a) a
187reader = Unfold step inject
188 where
189 inject arr = return (arr, 0)
190 step (arr, i)
191 | i == length arr = return D.Stop
192 | otherwise =
193 return $
194 case indexSmallArray## arr i of
195 (# x #) -> D.Yield x (arr, i + 1)