@@ -42,9 +42,15 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
4242
4343 -- * Random access
4444 , getIndexUnsafe
45+ , getIndex
46+ , getIndexRev
47+ , getIndices
4548
4649 -- * Mutation
4750 , putIndexUnsafe
51+ , putIndex
52+ , modifyIndexUnsafe
53+ , modifyIndex
4854 , snocUnsafe
4955 , snocWith
5056 , snoc
@@ -56,6 +62,10 @@ module Streamly.Internal.Data.Array.Foreign.Mut.Type
5662 , asBytes
5763 , asPtrUnsafe
5864
65+ -- * Subarrays
66+ , getSliceUnsafe
67+ , getSlice
68+
5969 -- * Folding
6070 , foldl'
6171 , foldr
@@ -381,6 +391,63 @@ putIndexUnsafe Array{..} i x =
381391 assert (i >= 0 && elemPtr `plusPtr` elemSize <= aEnd) (return ())
382392 liftIO $ poke elemPtr x
383393
394+invalidIndex :: String -> Int -> a
395+invalidIndex label i =
396+ error $ label ++ ": invalid array index " ++ show i
397+
398+{-# INLINE putIndexPtr #-}
399+putIndexPtr :: forall m a. (MonadIO m, Storable a) =>
400+ Ptr a -> Ptr a -> Int -> a -> m ()
401+putIndexPtr ptr end i x = do
402+ let elemSize = sizeOf (undefined :: a)
403+ elemPtr = ptr `plusPtr` (elemSize * i)
404+ if i >= 0 && elemPtr `plusPtr` elemSize <= end
405+ then liftIO $ poke elemPtr x
406+ else invalidIndex "putIndexPtr" i
407+
408+-- | /O(1)/ Write the given element at the given index in the array.
409+-- Performs in-place mutation of the array.
410+--
411+-- >>> putIndex arr ix val = Array.modifyIndex arr ix (const (val, ()))
412+--
413+-- /Pre-release/
414+{-# INLINE putIndex #-}
415+putIndex :: (MonadIO m, Storable a) => Array a -> Int -> a -> m ()
416+putIndex arr i x =
417+ unsafeWithForeignPtrM (aStart arr) $ \p -> putIndexPtr p (aEnd arr) i x
418+
419+-- | Modify a given index of an array using a modifier function.
420+--
421+-- /Pre-release/
422+modifyIndexUnsafe :: forall m a b. (MonadIO m, Storable a) =>
423+ Array a -> Int -> (a -> (a, b)) -> m b
424+modifyIndexUnsafe Array{..} i f = do
425+ liftIO $ unsafeWithForeignPtr aStart $ \ptr -> do
426+ let elemSize = sizeOf (undefined :: a)
427+ elemPtr = ptr `plusPtr` (elemSize * i)
428+ assert (i >= 0 && elemPtr `plusPtr` elemSize <= aEnd) (return ())
429+ r <- peek elemPtr
430+ let (x, res) = f r
431+ poke elemPtr x
432+ return res
433+
434+-- | Modify a given index of an array using a modifier function.
435+--
436+-- /Pre-release/
437+modifyIndex :: forall m a b. (MonadIO m, Storable a) =>
438+ Array a -> Int -> (a -> (a, b)) -> m b
439+modifyIndex Array{..} i f = do
440+ liftIO $ unsafeWithForeignPtr aStart $ \ptr -> do
441+ let elemSize = sizeOf (undefined :: a)
442+ elemPtr = ptr `plusPtr` (elemSize * i)
443+ if i >= 0 && elemPtr `plusPtr` elemSize <= aEnd
444+ then do
445+ r <- peek elemPtr
446+ let (x, res) = f r
447+ poke elemPtr x
448+ return res
449+ else invalidIndex "modifyIndex" i
450+
384451-------------------------------------------------------------------------------
385452-- Rounding
386453-------------------------------------------------------------------------------
@@ -619,6 +686,116 @@ getIndexUnsafe Array {..} i =
619686 assert (i >= 0 && elemPtr `plusPtr` elemSize <= aEnd) (return ())
620687 liftIO $ peek elemPtr
621688
689+{-# INLINE getIndexPtr #-}
690+getIndexPtr :: forall m a. (MonadIO m, Storable a) =>
691+ Ptr a -> Ptr a -> Int -> m a
692+getIndexPtr ptr end i = do
693+ let elemSize = sizeOf (undefined :: a)
694+ elemPtr = ptr `plusPtr` (elemSize * i)
695+ if i >= 0 && elemPtr `plusPtr` elemSize <= end
696+ then liftIO $ peek elemPtr
697+ else invalidIndex "getIndexPtr" i
698+
699+-- | /O(1)/ Lookup the element at the given index. Index starts from 0.
700+--
701+{-# INLINE getIndex #-}
702+getIndex :: (MonadIO m, Storable a) => Array a -> Int -> m a
703+getIndex arr i =
704+ unsafeWithForeignPtrM (aStart arr) $ \p -> getIndexPtr p (aEnd arr) i
705+
706+{-# INLINE getIndexPtrRev #-}
707+getIndexPtrRev :: forall m a. (MonadIO m, Storable a) =>
708+ Ptr a -> Ptr a -> Int -> m a
709+getIndexPtrRev ptr end i = do
710+ let elemSize = sizeOf (undefined :: a)
711+ elemPtr = end `plusPtr` negate (elemSize * (i + 1))
712+ if i >= 0 && elemPtr >= ptr
713+ then liftIO $ peek elemPtr
714+ else invalidIndex "getIndexPtrRev" i
715+
716+-- | /O(1)/ Lookup the element at the given index from the end of the array.
717+-- Index starts from 0.
718+--
719+-- Slightly faster than computing the forward index and using getIndex.
720+--
721+{-# INLINE getIndexRev #-}
722+getIndexRev :: (MonadIO m, Storable a) => Array a -> Int -> m a
723+getIndexRev arr i =
724+ unsafeWithForeignPtrM (aStart arr) $ \p -> getIndexPtrRev p (aEnd arr) i
725+
726+data GetIndicesState fp end st = GetIndicesState fp end st
727+
728+-- | Given an unfold that generates array indices, read the elements on those
729+-- indices from the supplied Array. An error is thrown if an index is out of
730+-- bounds.
731+--
732+-- /Pre-release/
733+{-# INLINE getIndices #-}
734+getIndices :: (MonadIO m, Storable a) =>
735+ Unfold m (Array a) Int -> Unfold m (Array a) a
736+getIndices (Unfold stepi injecti) = Unfold step inject
737+
738+ where
739+
740+ inject arr@(Array fp@(ForeignPtr _ _) (Ptr end) _) = do
741+ st <- injecti arr
742+ return $ GetIndicesState fp (Ptr end) st
743+
744+ {-# INLINE_LATE step #-}
745+ step (GetIndicesState fp@(ForeignPtr start _) end st) = do
746+ r <- stepi st
747+ case r of
748+ D.Yield i s -> do
749+ x <- liftIO $ getIndexPtr (Ptr start) end i
750+ return $ D.Yield x (GetIndicesState fp end s)
751+ D.Skip s -> return $ D.Skip (GetIndicesState fp end s)
752+ D.Stop -> do
753+ liftIO $ touchForeignPtr fp
754+ return D.Stop
755+
756+-------------------------------------------------------------------------------
757+-- Subarrays
758+-------------------------------------------------------------------------------
759+
760+-- | /O(1)/ Slice an array in constant time.
761+--
762+-- Unsafe: The bounds of the slice are not checked.
763+--
764+-- /Unsafe/
765+--
766+-- /Pre-release/
767+{-# INLINE getSliceUnsafe #-}
768+getSliceUnsafe :: forall a. Storable a
769+ => Int -- ^ from index
770+ -> Int -- ^ length of the slice
771+ -> Array a
772+ -> Array a
773+getSliceUnsafe index len (Array fp e _) =
774+ let size = sizeOf (undefined :: a)
775+ fp1 = fp `plusForeignPtr` (index * size)
776+ end = unsafeForeignPtrToPtr fp1 `plusPtr` (len * size)
777+ in assert (index >= 0 && len >= 0 && end <= e) (Array fp1 end end)
778+
779+-- | /O(1)/ Slice an array in constant time. Throws an error if the slice
780+-- extends out of the array bounds.
781+--
782+-- /Pre-release/
783+{-# INLINE getSlice #-}
784+getSlice :: forall a. Storable a =>
785+ Int -- ^ from index
786+ -> Int -- ^ length of the slice
787+ -> Array a
788+ -> Array a
789+getSlice index len (Array fp e _) =
790+ let size = sizeOf (undefined :: a)
791+ fp1 = fp `plusForeignPtr` (index * size)
792+ end = unsafeForeignPtrToPtr fp1 `plusPtr` (len * size)
793+ in if index >= 0 && len >= 0 && end <= e
794+ then Array fp1 end end
795+ else error
796+ $ "getSlice: invalid slice, index "
797+ ++ show index ++ " length " ++ show len
798+
622799-------------------------------------------------------------------------------
623800-- Size
624801-------------------------------------------------------------------------------