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

git.iko.soy/streamly

on 5d096df help sign in

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

342 lines 10 kB history raw

1-- |

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

4--

5-- License : BSD-3-Clause

6-- Maintainer : streamly@composewell.com

7-- Portability : GHC

8--

9module Streamly.Internal.Data.Array.Generic

10 ( Array(..)

12 -- * Construction

13 , nil

14 , createOf

15 , create

16 , writeWith

17 , writeLastN

19 , fromStreamN

20 , fromStream

21 , fromPureStream

22 , fromByteStr#

24 , fromListN

25 , fromList

27 , chunksOf

29 -- * Elimination

30 , length

31 , reader

33 , toList

34 , read

35 , readRev

37 , foldl'

38 , foldr

39 , streamFold

40 , fold

42 -- * Random Access

43 , getIndexUnsafe

44 , getIndex

45 , getSliceUnsafe

46 , strip

48 -- * Deprecated

49 , writeN

50 , write

51 )

52where

54#include "inline.hs"

56import Control.Monad (replicateM)

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

58import Data.Functor.Identity (Identity(..))

59import Data.Word (Word8)

60import GHC.Base (MutableArray#, RealWorld)

61import GHC.Exts (Addr#)

62import GHC.IO (unsafePerformIO)

63import Text.Read (readPrec)

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

66import Streamly.Internal.Data.Stream.Type (Stream)

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

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

70import qualified Streamly.Internal.Data.MutArray.Generic as MArray

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

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

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

74import qualified Streamly.Internal.Data.Ring.Generic as RB

75import qualified Streamly.Internal.Data.Stream.Type as D

76import qualified Streamly.Internal.Data.Stream.Generate as D

77import qualified Text.ParserCombinators.ReadPrec as ReadPrec

79import Prelude hiding (Foldable(..), read)

81-------------------------------------------------------------------------------

82-- Array Data Type

83-------------------------------------------------------------------------------

85data Array a =

86 Array

87 { arrContents# :: MutableArray# RealWorld a

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

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

91 -- ^ The starting index of this slice.

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

94 -- ^ The length of this slice.

95 }

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

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

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

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

102

103{-# NOINLINE nil #-}

104nil :: Array a

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

106

107-------------------------------------------------------------------------------

108-- Construction - Folds

109-------------------------------------------------------------------------------

110

111{-# INLINE_NORMAL createOf #-}

112createOf :: MonadIO m => Int -> Fold m a (Array a)

113createOf = fmap unsafeFreeze <$> MArray.createOf

114

115-- XXX Deprecate in major

116-- {-# DEPRECATED writeN "Please use createOf instead." #-}

117{-# INLINE writeN #-}

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

119writeN = createOf

120

121{-# INLINE_NORMAL writeWith #-}

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

123writeWith elemCount = unsafeFreeze <$> MArray.createWith elemCount

124

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

126--

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

128--

129{-# INLINE create #-}

130create :: MonadIO m => Fold m a (Array a)

131create = fmap unsafeFreeze MArray.create

132

133-- XXX Deprecate in major

134-- {-# DEPRECATED write "Please use create instead." #-}

135{-# INLINE write #-}

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

137write = create

138

139fromPureStream :: Stream Identity a -> Array a

140fromPureStream x =

141 unsafePerformIO $ fmap (unsafeFreeze) (MArray.fromPureStream x)

142-- fromPureStream = runIdentity . D.fold (unsafeMakePure write)

143-- fromPureStream = fromList . runIdentity . D.toList

144

145fromByteStr# :: Addr# -> Array Word8

146fromByteStr# addr = fromPureStream (D.fromByteStr# addr)

147

148-------------------------------------------------------------------------------

149-- Stream Ops

150-------------------------------------------------------------------------------

151

152{-# INLINE_NORMAL chunksOf #-}

153chunksOf :: forall m a. MonadIO m

154 => Int -> Stream m a -> Stream m (Array a)

155chunksOf n strm = fmap unsafeFreeze $ MArray.chunksOf n strm

156

157-------------------------------------------------------------------------------

158-- Construction - from streams

159-------------------------------------------------------------------------------

160

161{-# INLINE fromStreamN #-}

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

163fromStreamN n = D.fold (writeN n)

164

165{-# INLINE fromStream #-}

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

167fromStream = D.fold write

168

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

170

171{-# INLINABLE fromListN #-}

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

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

174

175{-# INLINABLE fromList #-}

176fromList :: [a] -> Array a

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

178

179-------------------------------------------------------------------------------

180-- Elimination - Unfolds

181-------------------------------------------------------------------------------

182

183{-# INLINE length #-}

184length :: Array a -> Int

185length = arrLen

186

187{-# INLINE_NORMAL reader #-}

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

189reader =

190 Producer.simplify

191 $ Producer.translate unsafeThaw unsafeFreeze

192 $ MArray.producerWith (return . unsafeInlineIO)

193

194-------------------------------------------------------------------------------

195-- Elimination - to streams

196-------------------------------------------------------------------------------

197

198{-# INLINE_NORMAL toList #-}

199toList :: Array a -> [a]

200toList arr = loop 0

201

202 where

203

204 len = length arr

205 loop i | i == len = []

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

207

208{-# INLINE_NORMAL read #-}

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

210read arr@Array{..} =

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

212

213{-# INLINE_NORMAL readRev #-}

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

215readRev arr@Array{..} =

216 D.map (`getIndexUnsafe` arr)

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

218

219-------------------------------------------------------------------------------

220-- Elimination - using Folds

221-------------------------------------------------------------------------------

222

223{-# INLINE_NORMAL foldl' #-}

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

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

226

227{-# INLINE_NORMAL foldr #-}

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

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

230

231{-# INLINE fold #-}

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

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

234

235{-# INLINE streamFold #-}

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

237streamFold f arr = f (read arr)

238

239-------------------------------------------------------------------------------

240-- Random reads and writes

241-------------------------------------------------------------------------------

242

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

244-- not check the bounds.

245--

246-- @since 0.8.0

247{-# INLINE getIndexUnsafe #-}

248getIndexUnsafe :: Int -> Array a -> a

249getIndexUnsafe i arr =

250 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)

251

252-- | Lookup the element at the given index. Index starts from 0.

253--

254{-# INLINE getIndex #-}

255getIndex :: Int -> Array a -> Maybe a

256getIndex i arr@Array {..} =

257 if i >= 0 && i < arrLen

258 then Just $ getIndexUnsafe i arr

259 else Nothing

260

261-- >>> import qualified Streamly.Data.Stream as Stream

262-- >>> import qualified Streamly.Data.Fold as Fold

263-- >>> import qualified Streamly.Internal.Data.Array.Generic as Array

264-- >>> import Data.Function ((&))

265-- >>> :{

266-- Stream.fromList [1,2,3,4,5::Int]

267-- & Stream.scan (Array.writeLastN 2)

268-- & Stream.fold Fold.toList

269-- :}

270-- [fromList [],fromList [1],fromList [1,2],fromList [2,3],fromList [3,4],fromList [4,5]]

271--

272{-# INLINE writeLastN #-}

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

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

275

276 where

277

278 f rb = do

279 arr <- RB.copyToMutArray 0 n rb

280 return $ unsafeFreeze arr

281

282{-# INLINE getSliceUnsafe #-}

283getSliceUnsafe :: Int -> Int -> Array a -> Array a

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

285

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

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

288-- copying.

289

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

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

292{-# INLINE strip #-}

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

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

295

296-------------------------------------------------------------------------------

297-- Instances

298-------------------------------------------------------------------------------

299

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

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

302 arr1 == arr2 =

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

304

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

306 {-# INLINE compare #-}

307 compare arr1 arr2 =

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

309

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

311 -- replicate them here with an INLINE.

312 {-# INLINE (<) #-}

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

314

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

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

317

318 {-# INLINE (>) #-}

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

320

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

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

323

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

325 -- because the latter is often more expensive

326 {-# INLINE max #-}

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

328

329 {-# INLINE min #-}

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

331

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

333 {-# INLINE show #-}

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

335

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

337 {-# INLINE readPrec #-}

338 readPrec = do

339 fromListWord <- replicateM 9 ReadPrec.get

340 if fromListWord == "fromList "

341 then fromList <$> readPrec

342 else ReadPrec.pfail