git.iko.soy/streamly
on 4e04eff help sign in
streamly/src/Streamly/Internal/Data/Array.hs
151 lines 4.2 kB history raw
1{-# OPTIONS_GHC -fno-warn-orphans #-}
3{-# LANGUAGE CPP #-}
4{-# LANGUAGE MagicHash #-}
5{-# LANGUAGE UnboxedTuples #-}
7#include "inline.hs"
9-- |
10-- Module : Streamly.Internal.Data.Array
11-- Copyright : (c) 2019 Composewell Technologies
12--
13-- License : BSD-3-Clause
14-- Maintainer : streamly@composewell.com
15-- Stability : experimental
16-- Portability : GHC
17--
18module Streamly.Internal.Data.Array
19 ( Array(..)
21 , foldl'
22 , foldr
24 , length
26 , writeN
27 , write
29 , toStreamD
30 , toStreamDRev
32 , fromListN
33 , fromList
34 , fromStreamDN
35 , fromStreamD
36 )
37where
39import Prelude hiding (foldr, length)
40import Control.DeepSeq (NFData(..))
41import Control.Monad.IO.Class (liftIO, MonadIO)
42import GHC.IO (unsafePerformIO)
43import GHC.Base (Int(..))
44import Data.Functor.Identity (runIdentity)
45import Data.Primitive.Array hiding (fromList, fromListN)
46import qualified GHC.Exts as Exts
48import Streamly.Internal.Data.Fold.Types (Fold(..))
50import qualified Streamly.Streams.StreamD as D
52{-# NOINLINE bottomElement #-}
53bottomElement :: a
54bottomElement = undefined
56{-# INLINE_NORMAL toStreamD #-}
57toStreamD :: Monad m => Array a -> D.Stream m a
58toStreamD arr = D.Stream step 0
59 where
60 {-# INLINE_LATE step #-}
61 step _ i
62 | i == length arr = return D.Stop
63 step _ (I# i) =
64 return $
65 case Exts.indexArray# (array# arr) i of
66 (# x #) -> D.Yield x ((I# i) + 1)
68{-# INLINE length #-}
69length :: Array a -> Int
70length arr = sizeofArray arr
72{-# INLINE_NORMAL toStreamDRev #-}
73toStreamDRev :: Monad m => Array a -> D.Stream m a
74toStreamDRev arr = D.Stream step (length arr - 1)
75 where
76 {-# INLINE_LATE step #-}
77 step _ i
78 | i < 0 = return D.Stop
79 step _ (I# i) =
80 return $
81 case Exts.indexArray# (array# arr) i of
82 (# x #) -> D.Yield x ((I# i) - 1)
84{-# INLINE_NORMAL foldl' #-}
85foldl' :: (b -> a -> b) -> b -> Array a -> b
86foldl' f z arr = runIdentity $ D.foldl' f z $ toStreamD arr
88{-# INLINE_NORMAL foldr #-}
89foldr :: (a -> b -> b) -> b -> Array a -> b
90foldr f z arr = runIdentity $ D.foldr f z $ toStreamD arr
92-- writeN n = S.evertM (fromStreamDN n)
93{-# INLINE_NORMAL writeN #-}
94writeN :: MonadIO m => Int -> Fold m a (Array a)
95writeN limit = Fold step initial extract
96 where
97 initial = do
98 marr <- liftIO $ newArray limit bottomElement
99 return (marr, 0)
100 step (marr, i) x
101 | i == limit = return (marr, i)
102 | otherwise = do
103 liftIO $ writeArray marr i x
104 return (marr, i + 1)
105 extract (marr, len) = liftIO $ freezeArray marr 0 len
107{-# INLINE_NORMAL write #-}
108write :: MonadIO m => Fold m a (Array a)
109write = Fold step initial extract
110 where
111 initial = do
112 marr <- liftIO $ newArray 0 bottomElement
113 return (marr, 0, 0)
114 step (marr, i, capacity) x
115 | i == capacity =
116 let newCapacity = max (capacity * 2) 1
117 in do newMarr <- liftIO $ newArray newCapacity bottomElement
118 liftIO $ copyMutableArray newMarr 0 marr 0 i
119 liftIO $ writeArray newMarr i x
120 return (newMarr, i + 1, newCapacity)
121 | otherwise = do
122 liftIO $ writeArray marr i x
123 return (marr, i + 1, capacity)
124 extract (marr, len, _) = liftIO $ freezeArray marr 0 len
126{-# INLINE_NORMAL fromStreamDN #-}
127fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (Array a)
128fromStreamDN limit str = do
129 marr <- liftIO $ newArray (max limit 0) bottomElement
130 i <-
131 D.foldlM'
132 (\i x -> i `seq` (liftIO $ writeArray marr i x) >> return (i + 1))
133 0 $
134 D.take limit str
135 liftIO $ freezeArray marr 0 i
137{-# INLINE fromStreamD #-}
138fromStreamD :: MonadIO m => D.Stream m a -> m (Array a)
139fromStreamD str = D.runFold write str
141{-# INLINABLE fromListN #-}
142fromListN :: Int -> [a] -> Array a
143fromListN n xs = unsafePerformIO $ fromStreamDN n $ D.fromList xs
145{-# INLINABLE fromList #-}
146fromList :: [a] -> Array a
147fromList xs = unsafePerformIO $ fromStreamD $ D.fromList xs
149instance NFData a => NFData (Array a) where
150 {-# INLINE rnf #-}
151 rnf = foldl' (\_ x -> rnf x) ()