561b13b · streamly · git.iko.soy

12(

13 -- * Type

14 -- $arrayNotes

15− Array (..)

15+ MutArray (..)

16

17 -- * Constructing and Writing

18 -- ** Construction

191-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold

192

193-------------------------------------------------------------------------------

194−-- Array Data Type

194+-- MutArray Data Type

195-------------------------------------------------------------------------------

196

197−data Array a =

197+data MutArray a =

198− Array

198+ MutArray

199 { arrContents# :: MutableArray# RealWorld a

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

201

233--

234-- /Pre-release/

235{-# INLINE new #-}

236−new :: forall m a. MonadIO m => Int -> m (Array a)

236+new :: forall m a. MonadIO m => Int -> m (MutArray a)

237new n@(I# n#) =

238 liftIO

239 $ IO

240 $ \s# ->

241 case newArray# n# bottomElement s# of

242 (# s1#, arr# #) ->

243− let ma = Array arr# 0 0 n

243+ let ma = MutArray arr# 0 0 n

244 in (# s1#, ma #)

245

246-------------------------------------------------------------------------------

252--

253-- /Pre-release/

254{-# INLINE putIndexUnsafe #-}

255−putIndexUnsafe :: forall m a. MonadIO m => Int -> Array a -> a -> m ()

255+putIndexUnsafe :: forall m a. MonadIO m => Int -> MutArray a -> a -> m ()

256−putIndexUnsafe i Array {..} x =

256+putIndexUnsafe i MutArray {..} x =

257 liftIO

258 $ IO

259 $ \s# ->

273--

274-- /Pre-release/

275{-# INLINE putIndex #-}

276−putIndex :: MonadIO m => Int -> Array a -> a -> m ()

276+putIndex :: MonadIO m => Int -> MutArray a -> a -> m ()

277−putIndex i arr@Array {..} x =

277+putIndex i arr@MutArray {..} x =

278 if i >= 0 && i < arrLen

279 then putIndexUnsafe i arr x

280 else invalidIndex "putIndex" i

285-- Unsafe because it does not check the bounds of the array.

286--

287-- /Pre-release/

288−modifyIndexUnsafe :: MonadIO m => Int -> Array a -> (a -> (a, b)) -> m b

288+modifyIndexUnsafe :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b

289−modifyIndexUnsafe i Array {..} f = do

289+modifyIndexUnsafe i MutArray {..} f = do

290 liftIO

291 $ IO

292 $ \s# ->

301-- | Modify a given index of an array using a modifier function.

302--

303-- /Pre-release/

304−modifyIndex :: MonadIO m => Int -> Array a -> (a -> (a, b)) -> m b

304+modifyIndex :: MonadIO m => Int -> MutArray a -> (a -> (a, b)) -> m b

305−modifyIndex i arr@Array {..} f = do

305+modifyIndex i arr@MutArray {..} f = do

306 if i >= 0 && i < arrLen

307 then modifyIndexUnsafe i arr f

308 else invalidIndex "modifyIndex" i

316-- the reallocated size is less than the original array it results in a

317-- truncated version of the original array.

318--

319−realloc :: MonadIO m => Int -> Array a -> m (Array a)

319+realloc :: MonadIO m => Int -> MutArray a -> m (MutArray a)

320realloc n arr = do

321 arr1 <- new n

322 let !newLen@(I# newLen#) = min n (arrLen arr)

331 in (# s1#, arr1 {arrLen = newLen, arrTrueLen = n} #)

332

333reallocWith ::

334− MonadIO m => String -> (Int -> Int) -> Int -> Array a -> m (Array a)

334+ MonadIO m => String -> (Int -> Int) -> Int -> MutArray a -> m (MutArray a)

335reallocWith label sizer reqSize arr = do

336 let oldSize = arrLen arr

337 newSize = sizer oldSize

360--

361-- /Internal/

362{-# INLINE snocUnsafe #-}

363−snocUnsafe :: MonadIO m => Array a -> a -> m (Array a)

363+snocUnsafe :: MonadIO m => MutArray a -> a -> m (MutArray a)

364−snocUnsafe arr@Array {..} a = do

364+snocUnsafe arr@MutArray {..} a = do

365 assert (arrStart + arrLen < arrTrueLen) (return ())

366 putIndexUnsafe arrLen arr a

367 return $ arr {arrLen = arrLen + 1}

368

369-- NOINLINE to move it out of the way and not pollute the instruction cache.

370{-# NOINLINE snocWithRealloc #-}

371−snocWithRealloc :: MonadIO m => (Int -> Int) -> Array a -> a -> m (Array a)

371+snocWithRealloc :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)

372snocWithRealloc sizer arr x = do

373 arr1 <- reallocWith "snocWithRealloc" sizer 1 arr

374 snocUnsafe arr1 x

384--

385-- /Pre-release/

386{-# INLINE snocWith #-}

387−snocWith :: MonadIO m => (Int -> Int) -> Array a -> a -> m (Array a)

387+snocWith :: MonadIO m => (Int -> Int) -> MutArray a -> a -> m (MutArray a)

388−snocWith sizer arr@Array {..} x = do

388+snocWith sizer arr@MutArray {..} x = do

389 if arrStart + arrLen < arrTrueLen

390 then snocUnsafe arr x

391 else snocWithRealloc sizer arr x

407--

408-- /Pre-release/

409{-# INLINE snoc #-}

410−snoc :: MonadIO m => Array a -> a -> m (Array a)

410+snoc :: MonadIO m => MutArray a -> a -> m (MutArray a)

411snoc = snocWith (* 2)

412

413-------------------------------------------------------------------------------

418--

419-- Unsafe because it does not check the bounds of the array.

420{-# INLINE_NORMAL getIndexUnsafe #-}

421−getIndexUnsafe :: MonadIO m => Int -> Array a -> m a

421+getIndexUnsafe :: MonadIO m => Int -> MutArray a -> m a

422−getIndexUnsafe n Array {..} =

422+getIndexUnsafe n MutArray {..} =

423 liftIO

424 $ IO

425 $ \s# ->

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

430--

431{-# INLINE getIndex #-}

432−getIndex :: MonadIO m => Int -> Array a -> m a

432+getIndex :: MonadIO m => Int -> MutArray a -> m a

433−getIndex i arr@Array {..} =

433+getIndex i arr@MutArray {..} =

434 if i >= 0 && i < arrLen

435 then getIndexUnsafe i arr

436 else invalidIndex "getIndex" i

452getSliceUnsafe

453 :: Int -- ^ from index

454 -> Int -- ^ length of the slice

455− -> Array a

455+ -> MutArray a

456− -> Array a

456+ -> MutArray a

457−getSliceUnsafe index len arr@Array {..} =

457+getSliceUnsafe index len arr@MutArray {..} =

458 assert (index >= 0 && len >= 0 && index + len <= arrLen)

459 $ arr {arrStart = arrStart + index, arrLen = len}

460

467getSlice

468 :: Int -- ^ from index

469 -> Int -- ^ length of the slice

470− -> Array a

470+ -> MutArray a

471− -> Array a

471+ -> MutArray a

472−getSlice index len arr@Array{..} =

472+getSlice index len arr@MutArray{..} =

473 if index >= 0 && len >= 0 && index + len <= arrLen

474 then arr {arrStart = arrStart + index, arrLen = len}

475 else error

484--

485-- /Pre-release/

486{-# INLINE toList #-}

487−toList :: MonadIO m => Array a -> m [a]

487+toList :: MonadIO m => MutArray a -> m [a]

488−toList arr@Array{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]

488+toList arr@MutArray{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]

489

490-- | Use the 'read' unfold instead.

491--

493--

494-- We can try this if the unfold has any performance issues.

495{-# INLINE_NORMAL toStreamD #-}

496−toStreamD :: MonadIO m => Array a -> D.Stream m a

496+toStreamD :: MonadIO m => MutArray a -> D.Stream m a

497−toStreamD arr@Array{..} =

497+toStreamD arr@MutArray{..} =

498 D.mapM (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)

499

500{-# INLINE toStreamK #-}

501−toStreamK :: MonadIO m => Array a -> K.StreamK m a

501+toStreamK :: MonadIO m => MutArray a -> K.StreamK m a

502−toStreamK arr@Array{..} = K.unfoldrM step 0

502+toStreamK arr@MutArray{..} = K.unfoldrM step 0

503

504 where

505

522--

523-- /Pre-release/

524{-# INLINE_NORMAL writeNUnsafe #-}

525−writeNUnsafe :: MonadIO m => Int -> Fold m a (Array a)

525+writeNUnsafe :: MonadIO m => Int -> Fold m a (MutArray a)

526writeNUnsafe n = Fold step initial return

527

528 where

538--

539-- /Pre-release/

540{-# INLINE_NORMAL writeN #-}

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

541+writeN :: MonadIO m => Int -> Fold m a (MutArray a)

542writeN n = FL.take n $ writeNUnsafe n

543

544-------------------------------------------------------------------------------

548-- | Resumable unfold of an array.

549--

550{-# INLINE_NORMAL producerWith #-}

551−producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (Array a) a

551+producerWith :: Monad m => (forall b. IO b -> m b) -> Producer m (MutArray a) a

552producerWith liftio = Producer step inject extract

553

554 where

570-- | Resumable unfold of an array.

571--

572{-# INLINE_NORMAL producer #-}

573−producer :: MonadIO m => Producer m (Array a) a

573+producer :: MonadIO m => Producer m (MutArray a) a

574producer = producerWith liftIO

575

576-- | Unfold an array into a stream.

577--

578{-# INLINE_NORMAL reader #-}

579−reader :: MonadIO m => Unfold m (Array a) a

579+reader :: MonadIO m => Unfold m (MutArray a) a

580reader = Producer.simplify producer

581

582--------------------------------------------------------------------------------

586-- | Put a sub range of a source array into a subrange of a destination array.

587-- This is not safe as it does not check the bounds.

588{-# INLINE putSliceUnsafe #-}

589−putSliceUnsafe :: MonadIO m => Array a -> Int -> Array a -> Int -> Int -> m ()

589+putSliceUnsafe :: MonadIO m => MutArray a -> Int -> MutArray a -> Int -> Int -> m ()

590putSliceUnsafe src srcStart dst dstStart len = liftIO $ do

591 assertM(len <= arrLen dst)

592 assertM(len <= arrLen src)

600 , () #)

601

602{-# INLINE clone #-}

603−clone :: MonadIO m => Array a -> m (Array a)

603+clone :: MonadIO m => MutArray a -> m (MutArray a)

604clone src = liftIO $ do

605 let len = arrLen src

606 dst <- new len