on master help sign in
Expose some array APIs
commit3187e3794f5d97d0ce966722e022301db9bf27c8 authorRanjeet Kumar Ranjan <ranjeetr@composewell.com> dateThu 13 Apr 2023, 15:25 (4 years ago) parent1d5e891
4 files changed+34−8
core/src/Streamly/Data/Array/Generic.hs+1 core/src/Streamly/Data/MutArray/Generic.hs+7 core/src/Streamly/Internal/Data/Array/Generic.hs+9 core/src/Streamly/Internal/Data/Array/Generic/Mut/Type.hs+17−8
core/src/Streamly/Data/Array/Generic.hs+1
@@ -31,6 +31,7 @@ module Streamly.Data.Array.Generic
3131
3232 -- * Random Access
3333 , A.length
34+ , A.getIndex
3435
3536 -- -- * Folding Arrays
3637 -- , A.streamFold
core/src/Streamly/Data/MutArray/Generic.hs+7
@@ -35,15 +35,22 @@ module Streamly.Data.MutArray.Generic
3535 -- * Unfolds
3636 , reader
3737
38+ -- * Streams
39+ , read
40+
3841 -- * Random reads
3942 , getIndex
4043
4144 -- * Inplace mutation
4245 , putIndex
4346 , modifyIndex
47+
48+ -- * Size
49+ , length
4450 )
4551where
4652
4753import Streamly.Internal.Data.Array.Generic.Mut.Type
54+import Prelude hiding (read, length)
4855
4956#include "DocTestDataMutArrayGeneric.hs"
core/src/Streamly/Internal/Data/Array/Generic.hs+9
@@ -37,6 +37,7 @@ module Streamly.Internal.Data.Array.Generic
3737
3838 -- * Random Access
3939 , getIndexUnsafe
40+ , getIndex
4041 , getSliceUnsafe
4142 , strip
4243 )
@@ -207,6 +208,14 @@ getIndexUnsafe :: Int -> Array a -> a
207208getIndexUnsafe i arr =
208209 unsafePerformIO $ MArray.getIndexUnsafe i (unsafeThaw arr)
209210
211+-- | Lookup the element at the given index. Index starts from 0.
212+--
213+getIndex :: Int -> Array a -> Maybe a
214+getIndex i arr@Array {..} =
215+ if i >= 0 && i < arrLen
216+ then Just $ getIndexUnsafe i arr
217+ else Nothing
218+
210219{-# INLINE writeLastN #-}
211220writeLastN :: MonadIO m => Int -> Fold m a (Array a)
212221writeLastN n = FL.rmapM f (RB.writeLastN n)
core/src/Streamly/Internal/Data/Array/Generic/Mut/Type.hs+17−8
@@ -27,6 +27,7 @@ module Streamly.Internal.Data.Array.Generic.Mut.Type
2727 , writeN
2828 , writeWith
2929 , write
30+ , fromStreamN
3031
3132 -- , writeRevN
3233 -- , writeRev
@@ -91,7 +92,7 @@ module Streamly.Internal.Data.Array.Generic.Mut.Type
9192 , producer -- experimental
9293
9394 -- ** To containers
94− , toStreamD
95+ , read
9596 , readRev
9697 , toStreamK
9798 -- , toStreamKRev
@@ -179,6 +180,7 @@ import GHC.Int (Int(..))
179180import Streamly.Internal.Data.Fold.Type (Fold(..))
180181import Streamly.Internal.Data.Producer.Type (Producer (..))
181182import Streamly.Internal.Data.Unfold.Type (Unfold(..))
183+import Streamly.Internal.Data.Stream.StreamD.Type (Stream)
182184
183185import qualified Streamly.Internal.Data.Fold.Type as FL
184186import qualified Streamly.Internal.Data.Producer as Producer
@@ -527,14 +529,13 @@ getSlice index len arr@MutArray{..} =
527529toList :: MonadIO m => MutArray a -> m [a]
528530toList arr@MutArray{..} = mapM (`getIndexUnsafe` arr) [0 .. (arrLen - 1)]
529531
530−-- | Use the 'read' unfold instead.
532+-- | Generates a stream from the elements of @MutArray@
531533--
532−-- @toStreamD = D.unfold read@
533−--
534−-- We can try this if the unfold has any performance issues.
535−{-# INLINE_NORMAL toStreamD #-}
536−toStreamD :: MonadIO m => MutArray a -> D.Stream m a
537−toStreamD arr@MutArray{..} =
534+-- @read = D.unfold reader@
535+-- /Pre-release/
536+{-# INLINE_NORMAL read #-}
537+read :: MonadIO m => MutArray a -> D.Stream m a
538+read arr@MutArray{..} =
538539 D.mapM (`getIndexUnsafe` arr) $ D.enumerateFromToIntegral 0 (arrLen - 1)
539540
540541-- Check equivalence with StreamK.fromStream . toStreamD and remove
@@ -794,3 +795,11 @@ strip p arr = liftIO $ do
794795 if p r
795796 then getIndexL (idx + 1)
796797 else return idx
798+
799+-- | Create a 'MutArray' from the first @n@ elements of a stream. The
800+-- array is allocated to size @n@, if the stream terminates before @n@
801+-- elements then the array may hold less than @n@ elements.
802+--
803+{-# INLINE fromStreamN #-}
804+fromStreamN :: MonadIO m => Int -> Stream m a -> m (MutArray a)
805+fromStreamN n = D.fold (writeN n)