core/src/Streamly/Internal/Data/Array/Generic.hs · streamly

git.iko.soy/streamly

on 01fa8dd help sign in

streamly/core/src/Streamly/Internal/Data/Array/Generic.hs

281 lines 8.6 kB history raw

1-- |

2-- Module : Streamly.Internal.Data.Array.Generic

4--

5-- License : BSD-3-Clause

6-- Maintainer : streamly@composewell.com

7-- Stability : pre-release

8-- Portability : GHC

9--

10module Streamly.Internal.Data.Array.Generic

11 ( Array(..)

13 -- * Construction

14 , nil

15 , writeN

16 , write

17 , writeWith

18 , writeLastN

20 , fromStreamN

21 , fromStream

23 , fromListN

24 , fromList

26 -- * Elimination

27 , length

28 , reader

30 , toList

31 , read

32 , readRev

34 , foldl'

35 , foldr

36 , streamFold

37 , fold

39 -- * Random Access

40 , getIndexUnsafe

41 , getSliceUnsafe

42 , strip

43 )

44where

46#include "inline.hs"

48import Control.Monad (replicateM)

49import Control.Monad.IO.Class (MonadIO)

50import GHC.Base (MutableArray#, RealWorld)

51import GHC.IO (unsafePerformIO)

52import Text.Read (readPrec)

54import Streamly.Internal.Data.Fold.Type (Fold(..))

55import Streamly.Internal.Data.Stream.StreamD.Type (Stream)

56import Streamly.Internal.Data.Unfold.Type (Unfold(..))

57import Streamly.Internal.System.IO (unsafeInlineIO)

59import qualified Streamly.Internal.Data.Array.Generic.Mut.Type as MArray

60import qualified Streamly.Internal.Data.Fold.Type as FL

61import qualified Streamly.Internal.Data.Producer.Type as Producer

62import qualified Streamly.Internal.Data.Producer as Producer

63import qualified Streamly.Internal.Data.Ring as RB

64import qualified Streamly.Internal.Data.Stream.StreamD.Type as D

65import qualified Streamly.Internal.Data.Stream.StreamD.Generate as D

66import qualified Text.ParserCombinators.ReadPrec as ReadPrec

68import Prelude hiding (foldr, length, read)

70-------------------------------------------------------------------------------

71-- Array Data Type

72-------------------------------------------------------------------------------

74data Array a =

75 Array

76 { arrContents# :: MutableArray# RealWorld a

77 -- ^ The internal contents of the array representing the entire array.

79 , arrStart :: {-# UNPACK #-}!Int

80 -- ^ The starting index of this slice.

82 , arrLen :: {-# UNPACK #-}!Int

83 -- ^ The length of this slice.

84 }

86unsafeFreeze :: MArray.MutArray a -> Array a

87unsafeFreeze (MArray.MutArray cont# arrS arrL _) = Array cont# arrS arrL

89unsafeThaw :: Array a -> MArray.MutArray a

90unsafeThaw (Array cont# arrS arrL) = MArray.MutArray cont# arrS arrL arrL

92{-# NOINLINE nil #-}

93nil :: Array a

94nil = unsafePerformIO $ unsafeFreeze <$> MArray.nil

96-------------------------------------------------------------------------------

97-- Construction - Folds

98-------------------------------------------------------------------------------

100{-# INLINE_NORMAL writeN #-}

101writeN :: MonadIO m => Int -> Fold m a (Array a)

102writeN = fmap unsafeFreeze <$> MArray.writeN

103

104{-# INLINE_NORMAL writeWith #-}

105writeWith :: MonadIO m => Int -> Fold m a (Array a)

106writeWith elemCount = unsafeFreeze <$> MArray.writeWith elemCount

107

108-- | Fold the whole input to a single array.

109--

110-- /Caution! Do not use this on infinite streams./

111--

112{-# INLINE write #-}

113write :: MonadIO m => Fold m a (Array a)

114write = fmap unsafeFreeze MArray.write

115

116-------------------------------------------------------------------------------

117-- Construction - from streams

118-------------------------------------------------------------------------------

119

120{-# INLINE fromStreamN #-}

121fromStreamN :: MonadIO m => Int -> Stream m a -> m (Array a)

122fromStreamN n = D.fold (writeN n)

123

124{-# INLINE fromStream #-}

125fromStream :: MonadIO m => Stream m a -> m (Array a)

126fromStream = D.fold write

127

128-- XXX Consider foldr/build fusion in toList/fromList

129

130{-# INLINABLE fromListN #-}

131fromListN :: Int -> [a] -> Array a

132fromListN n xs = unsafePerformIO $ fromStreamN n $ D.fromList xs

133

134{-# INLINABLE fromList #-}

135fromList :: [a] -> Array a

136fromList xs = unsafePerformIO $ fromStream $ D.fromList xs

137

138-------------------------------------------------------------------------------

139-- Elimination - Unfolds

140-------------------------------------------------------------------------------

141

142{-# INLINE length #-}

143length :: Array a -> Int

144length = arrLen

145

146{-# INLINE_NORMAL reader #-}

147reader :: Monad m => Unfold m (Array a) a

148reader =

149 Producer.simplify

150 $ Producer.translate unsafeThaw unsafeFreeze

151 $ MArray.producerWith (return . unsafeInlineIO)

152

153-------------------------------------------------------------------------------

154-- Elimination - to streams

155-------------------------------------------------------------------------------

156

157{-# INLINE_NORMAL toList #-}

158toList :: Array a -> [a]

159toList arr = loop 0

160

161 where

162

163 len = length arr

164 loop i | i == len = []

165 loop i = getIndexUnsafe i arr : loop (i + 1)

166

167{-# INLINE_NORMAL read #-}

168read :: Monad m => Array a -> Stream m a

169read arr@Array{..} =

170 D.map (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)

171

172{-# INLINE_NORMAL readRev #-}

173readRev :: Monad m => Array a -> Stream m a

174readRev arr@Array{..} =

175 D.map (`getIndexUnsafe` arr)

176 $ D.enumerateFromThenToIntegral (arrLen - 1) (arrLen - 2) 0

177

178-------------------------------------------------------------------------------

179-- Elimination - using Folds

180-------------------------------------------------------------------------------

181

182{-# INLINE_NORMAL foldl' #-}

183foldl' :: (b -> a -> b) -> b -> Array a -> b

184foldl' f z arr = unsafePerformIO $ D.foldl' f z $ read arr

185

186{-# INLINE_NORMAL foldr #-}

187foldr :: (a -> b -> b) -> b -> Array a -> b

188foldr f z arr = unsafePerformIO $ D.foldr f z $ read arr

189

190{-# INLINE fold #-}

191fold :: Monad m => Fold m a b -> Array a -> m b

192fold f arr = D.fold f (read arr)

193

194{-# INLINE streamFold #-}

195streamFold :: Monad m => (Stream m a -> m b) -> Array a -> m b

196streamFold f arr = f (read arr)

197

198-------------------------------------------------------------------------------

199-- Random reads and writes

200-------------------------------------------------------------------------------

201

202-- | /O(1)/ Lookup the element at the given index. Index starts from 0. Does

203-- not check the bounds.

204--

205-- @since 0.8.0

206{-# INLINE getIndexUnsafe #-}

207getIndexUnsafe :: Int -> Array a -> a

208getIndexUnsafe i arr =

209 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)

210

211{-# INLINE writeLastN #-}

212writeLastN :: MonadIO m => Int -> Fold m a (Array a)

213writeLastN n = FL.rmapM f (RB.writeLastN n)

214

215 where

216

217 f rb = do

218 arr <- RB.toMutArray 0 n rb

219 return $ unsafeFreeze arr

220

221{-# INLINE getSliceUnsafe #-}

222getSliceUnsafe :: Int -> Int -> Array a -> Array a

223getSliceUnsafe offset len (Array cont off1 _) = Array cont (off1 + offset) len

224

225-- XXX This is not efficient as it copies the array. We should support array

226-- slicing so that we can just refer to the underlying array memory instead of

227-- copying.

228

229-- | Truncate the array at the beginning and end as long as the predicate

230-- holds true. Returns a slice of the original array.

231{-# INLINE strip #-}

232strip :: (a -> Bool) -> Array a -> Array a

233strip p arr = unsafeFreeze $ unsafePerformIO $ MArray.strip p (unsafeThaw arr)

234

235-------------------------------------------------------------------------------

236-- Instances

237-------------------------------------------------------------------------------

238

239instance Eq a => Eq (Array a) where

240 {-# INLINE (==) #-}

241 arr1 == arr2 =

242 unsafeInlineIO $! unsafeThaw arr1 `MArray.eq` unsafeThaw arr2

243

244instance Ord a => Ord (Array a) where

245 {-# INLINE compare #-}

246 compare arr1 arr2 =

247 unsafeInlineIO $! unsafeThaw arr1 `MArray.cmp` unsafeThaw arr2

248

249 -- Default definitions defined in base do not have an INLINE on them, so we

250 -- replicate them here with an INLINE.

251 {-# INLINE (<) #-}

252 x < y = case compare x y of { LT -> True; _ -> False }

253

254 {-# INLINE (<=) #-}

255 x <= y = case compare x y of { GT -> False; _ -> True }

256

257 {-# INLINE (>) #-}

258 x > y = case compare x y of { GT -> True; _ -> False }

259

260 {-# INLINE (>=) #-}

261 x >= y = case compare x y of { LT -> False; _ -> True }

262

263 -- These two default methods use '<=' rather than 'compare'

264 -- because the latter is often more expensive

265 {-# INLINE max #-}

266 max x y = if x <= y then y else x

267

268 {-# INLINE min #-}

269 min x y = if x <= y then x else y

270

271instance Show a => Show (Array a) where

272 {-# INLINE show #-}

273 show arr = "fromList " ++ show (toList arr)

274

275instance Read a => Read (Array a) where

276 {-# INLINE readPrec #-}

277 readPrec = do

278 fromListWord <- replicateM 9 ReadPrec.get

279 if fromListWord == "fromList "

280 then fromList <$> readPrec

281 else ReadPrec.pfail