git.iko.soy/streamly
on 33b00a6 help sign in
streamly/src/Streamly/Internal/Data/SmallArray.hs
181 lines 5.3 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 #-}
12{-# LANGUAGE CPP #-}
13{-# LANGUAGE MagicHash #-}
14{-# LANGUAGE UnboxedTuples #-}
16#include "inline.hs"
18module Streamly.Internal.Data.SmallArray
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)
50import Data.Primitive.SmallArray
52import Streamly.Internal.Data.Unfold.Types (Unfold(..))
53import Streamly.Internal.Data.Fold.Types (Fold(..))
54import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
55import Streamly.Internal.Data.Stream.Serial (SerialT)
57import qualified Streamly.Internal.Data.Stream.StreamD as D
59{-# NOINLINE bottomElement #-}
60bottomElement :: a
61bottomElement = undefined
63{-# INLINE length #-}
64length :: SmallArray a -> Int
65length arr = sizeofSmallArray arr
67{-# INLINE_NORMAL toStreamD #-}
68toStreamD :: Monad m => SmallArray a -> D.Stream m a
69toStreamD arr = D.Stream step 0
70 where
71 {-# INLINE_LATE step #-}
72 step _ i
73 | i == length arr = return D.Stop
74 | otherwise =
75 return $
76 case indexSmallArray## arr i of
77 (# x #) -> D.Yield x (i + 1)
79{-# INLINE_NORMAL toStreamDRev #-}
80toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
81toStreamDRev arr = D.Stream step (length arr - 1)
82 where
83 {-# INLINE_LATE step #-}
84 step _ i
85 | i < 0 = return D.Stop
86 | otherwise =
87 return $
88 case indexSmallArray## arr i of
89 (# x #) -> D.Yield x (i - 1)
91{-# INLINE_NORMAL foldl' #-}
92foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
93foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
95{-# INLINE_NORMAL foldr #-}
96foldr :: (a -> b -> b) -> b -> SmallArray a -> b
97foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
99-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
100-- 'SmallArray'.
101--
102-- Since we are folding to a 'SmallArray' @n@ should be <= 128, for larger number
103-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Memory.Array".
104{-# INLINE_NORMAL writeN #-}
105writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
106writeN limit = Fold step initial extract
107 where
108 initial = do
109 marr <- liftIO $ newSmallArray limit bottomElement
110 return (marr, 0)
111 step (marr, i) x
112 | i == limit = return (marr, i)
113 | otherwise = do
114 liftIO $ writeSmallArray marr i x
115 return (marr, i + 1)
116 extract (marr, len) = liftIO $ freezeSmallArray marr 0 len
118{-# INLINE_NORMAL fromStreamDN #-}
119fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
120fromStreamDN limit str = do
121 marr <- liftIO $ newSmallArray (max limit 0) bottomElement
122 i <-
123 D.foldlM'
124 (\i x -> i `seq` (liftIO $ writeSmallArray marr i x) >> return (i + 1))
125 0 $
126 D.take limit str
127 liftIO $ freezeSmallArray marr 0 i
129-- | Create a 'SmallArray' from the first @n@ elements of a list. The
130-- array may hold less than @n@ elements if the length of the list <=
131-- @n@.
132--
133-- It is recommended to use a value of @n@ <= 128. For larger sized
134-- arrays, use an 'Array' from "Streamly.Data.Array" or
135-- "Streamly.Memory.Array"
136{-# INLINABLE fromListN #-}
137fromListN :: Int -> [a] -> SmallArray a
138fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
140instance NFData a => NFData (SmallArray a) where
141 {-# INLINE rnf #-}
142 rnf = foldl' (\_ x -> rnf x) ()
144-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
145-- array is allocated to size @n@, if the stream terminates before @n@
146-- elements then the array may hold less than @n@ elements.
147--
148-- For optimal performance use this with @n@ <= 128.
149{-# INLINE fromStreamN #-}
150fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
151fromStreamN n m = do
152 when (n < 0) $ error "fromStreamN: negative write count specified"
153 fromStreamDN n $ D.toStreamD m
155{-# INLINE_EARLY toStream #-}
156toStream :: (Monad m, IsStream t) => SmallArray a -> t m a
157toStream = D.fromStreamD . toStreamD
159{-# INLINE_EARLY toStreamRev #-}
160toStreamRev :: (Monad m, IsStream t) => SmallArray a -> t m a
161toStreamRev = D.fromStreamD . toStreamDRev
163{-# INLINE fold #-}
164fold :: Monad m => Fold m a b -> SmallArray a -> m b
165fold f arr = D.runFold f (toStreamD arr)
167{-# INLINE streamFold #-}
168streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
169streamFold f arr = f (toStream arr)
171{-# INLINE_NORMAL read #-}
172read :: Monad m => Unfold m (SmallArray a) a
173read = Unfold step inject
174 where
175 inject arr = return (arr, 0)
176 step (arr, i)
177 | i == length arr = return D.Stop
178 | otherwise =
179 return $
180 case indexSmallArray## arr i of
181 (# x #) -> D.Yield x (arr, i + 1)