on master help sign in
Implement modifyIndices
commitfcb15a693c5c060476e7574b1493b037d00feb66 authorRanjeet Kumar Ranjan <ranjeetr@composewell.com> dateMon 27 Dec 2021, 20:16 (5 years ago) parentb5bf22b
1 file changed+32−15
src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+32−15
src/Streamly/Internal/Data/Array/Foreign/Mut/Type.hs+32−15
@@ -568,9 +568,9 @@ invalidIndex label i =
568568{-# INLINE putIndexPtr #-}
569569putIndexPtr :: forall m a. (MonadIO m, Storable a) =>
570570 Ptr a -> Ptr a -> Int -> a -> m ()
571−putIndexPtr ptr end i x = do
571+putIndexPtr start end i x = do
572572 let elemSize = sizeOf (undefined :: a)
573− elemPtr = ptr `plusPtr` (elemSize * i)
573+ elemPtr = start `plusPtr` (elemSize * i)
574574 if i >= 0 && elemPtr `plusPtr` elemSize <= end
575575 then liftIO $ poke elemPtr x
576576 else invalidIndex "putIndexPtr" i
@@ -621,6 +621,20 @@ modifyIndexUnsafe Array{..} i f = do
621621 poke elemPtr x
622622 return res
623623
624+{-# INLINE modifyIndexPtr #-}
625+modifyIndexPtr :: forall m a b. (MonadIO m, Storable a) =>
626+ Int -> (a -> (a, b)) -> Ptr a -> Ptr a -> m b
627+modifyIndexPtr i f start end = liftIO $ do
628+ let elemSize = sizeOf (undefined :: a)
629+ elemPtr = start `plusPtr` (elemSize * i)
630+ if i >= 0 && elemPtr `plusPtr` elemSize <= end
631+ then do
632+ r <- peek elemPtr
633+ let (x, res) = f r
634+ poke elemPtr x
635+ return res
636+ else invalidIndex "modifyIndex" i
637+
624638-- | Modify a given index of an array using a modifier function.
625639--
626640-- /Pre-release/
@@ -628,22 +642,25 @@ modifyIndex :: forall m a b. (MonadIO m, Storable a) =>
628642 Array a -> Int -> (a -> (a, b)) -> m b
629643modifyIndex Array{..} i f = do
630644 liftIO $ unsafeWithArrayContents arrContents arrStart $ \ptr -> do
631− let elemSize = sizeOf (undefined :: a)
632− elemPtr = ptr `plusPtr` (elemSize * i)
633− if i >= 0 && elemPtr `plusPtr` elemSize <= aEnd
634− then do
635− r <- peek elemPtr
636− let (x, res) = f r
637− poke elemPtr x
638− return res
639− else invalidIndex "modifyIndex" i
645+ modifyIndexPtr i f ptr aEnd
640646
641−-- | Modify the array indices generated by the supplied unfold.
647+-- | Modify the array indices generated by the supplied stream.
642648--
643649-- /Pre-release/
644−modifyIndices :: -- forall m a b. (MonadIO m, Storable a) =>
645− Unfold m (Array a) Int -> Array a -> (a -> a) -> m ()
646−modifyIndices = undefined
650+{-# INLINE modifyIndices #-}
651+modifyIndices :: forall m a. (MonadIO m, Storable a)
652+ => (a -> a) -> Array a -> Fold m Int ()
653+modifyIndices f Array{..} = Fold step initial extract
654+
655+ where
656+
657+ initial = return $ FL.Partial ()
658+
659+ step () i =
660+ let f1 x = (f x, ())
661+ in FL.Partial <$> liftIO (modifyIndexPtr i f1 arrStart aEnd)
662+
663+ extract () = liftIO $ touch arrContents
647664
648665-- | Modify each element of an array using the supplied modifier function.
649666--