git.iko.soy/streamly
on 36bbaed help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
187 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 , 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 limit = Fold step initial extract
110 where
112 initial = do
113 marr <- liftIO $ newSmallArray limit bottomElement
114 return $ FL.Partial (Tuple' marr 0)
116 step st@(Tuple' marr i) x
117 | i == limit = FL.Done <$> extract st
118 | otherwise = do
119 liftIO $ writeSmallArray marr i x
120 return $ FL.Partial (Tuple' marr (i + 1))
122 extract (Tuple' marr len) = liftIO $ freezeSmallArray marr 0 len
124{-# INLINE_NORMAL fromStreamDN #-}
125fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
126fromStreamDN limit str = do
127 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
128 i <-
129 D.foldlM'
130 (\i x -> i `seq` liftIO (writeSmallArray marr i x) >> return (i + 1))
131 (return 0) $
132 D.take limit str
133 liftIO $ freezeSmallArray marr 0 i
135-- | Create a 'SmallArray' from the first @n@ elements of a list. The
136-- array may hold less than @n@ elements if the length of the list <=
137-- @n@.
138--
139-- It is recommended to use a value of @n@ <= 128. For larger sized
140-- arrays, use an 'Array' from "Streamly.Data.Array" or
141-- "Streamly.Data.Array.Foreign"
142{-# INLINABLE fromListN #-}
143fromListN :: Int -> [a] -> SmallArray a
144fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
146instance NFData a => NFData (SmallArray a) where
147 {-# INLINE rnf #-}
148 rnf = foldl' (\_ x -> rnf x) ()
150-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
151-- array is allocated to size @n@, if the stream terminates before @n@
152-- elements then the array may hold less than @n@ elements.
153--
154-- For optimal performance use this with @n@ <= 128.
155{-# INLINE fromStreamN #-}
156fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
157fromStreamN n (SerialT m) = do
158 when (n < 0) $ error "fromStreamN: negative write count specified"
159 fromStreamDN n $ D.fromStreamK m
161{-# INLINE_EARLY toStream #-}
162toStream :: Monad m => SmallArray a -> SerialT m a
163toStream = SerialT . D.toStreamK . toStreamD
165{-# INLINE_EARLY toStreamRev #-}
166toStreamRev :: Monad m => SmallArray a -> SerialT m a
167toStreamRev = SerialT . D.toStreamK . toStreamDRev
169{-# INLINE fold #-}
170fold :: Monad m => Fold m a b -> SmallArray a -> m b
171fold f arr = D.fold f (toStreamD arr)
173{-# INLINE streamFold #-}
174streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
175streamFold f arr = f (toStream arr)
177{-# INLINE_NORMAL read #-}
178read :: Monad m => Unfold m (SmallArray a) a
179read = Unfold step inject
180 where
181 inject arr = return (arr, 0)
182 step (arr, i)
183 | i == length arr = return D.Stop
184 | otherwise =
185 return $
186 case indexSmallArray## arr i of
187 (# x #) -> D.Yield x (arr, i + 1)