Add Data.MArray.{modifyArray,modifyArray'}

Something I've defined and used many times that feels like it should be provided by the module.

This is a commonly available function in similar modules, for example:

• Data.IORef.modifyIORef
• Data.STRef.modifySTRef
• Data.Vector.Mutable.modify

I'm not sure if this INLINE is necessary, but let's keep it since most functions in array are inlined.

Inlining small array operations is usually very beneficial but I have seen GHC not be eager enough to do that on its own.

For instance, here is a benchmark with a somewhat silly example:

{-# LANGUAGE BangPatterns #-}
import Criterion.Main
import Data.Array.Base
import Data.Array.ST
import Data.Foldable

main :: IO ()
main = defaultMain
    [ bench "testMissingInline" $ whnf testMissingInline n
    , bench "testInline" $ whnf testInline n
    ]
  where
    n = 100000

testMissingInline :: Int -> UArray Int Int
testMissingInline n = runSTUArray $ do
    a <- newArray (1,n) 0
    for_ [1..n] $ \i ->
        modifyArray'MissingInline a i (+1)
    for_ [n,n-1..1] $ \i ->
        modifyArray'MissingInline a i (+1)
    pure a
{-# NOINLINE testMissingInline #-}

testInline :: Int -> UArray Int Int
testInline n = runSTUArray $ do
    a <- newArray (1,n) 0
    for_ [1..n] $ \i ->
        modifyArray'Inline a i (+1)
    for_ [n,n-1..1] $ \i ->
        modifyArray'Inline a i (+1)
    pure a
{-# NOINLINE testInline #-}

modifyArray'MissingInline :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m ()
modifyArray'MissingInline marr i f = do
    (l,u) <- getBounds marr
    n <- getNumElements marr
    let idx = safeIndex (l,u) n i
    x <- unsafeRead marr idx
    let !x' = f x
    unsafeWrite marr idx x'

modifyArray'Inline :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m ()
modifyArray'Inline marr i f = do
    (l,u) <- getBounds marr
    n <- getNumElements marr
    let idx = safeIndex (l,u) n i
    x <- unsafeRead marr idx
    let !x' = f x
    unsafeWrite marr idx x'
{-# INLINE modifyArray'Inline #-}

On GHC 9.2.5 with -O2 I get

benchmarking testMissingInline
time                 981.3 μs   (977.3 μs .. 984.5 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 978.3 μs   (975.9 μs .. 981.7 μs)
std dev              9.496 μs   (7.526 μs .. 12.20 μs)

benchmarking testInline
time                 374.5 μs   (373.4 μs .. 375.5 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 374.7 μs   (374.0 μs .. 375.3 μs)
std dev              2.143 μs   (1.696 μs .. 2.890 μs)

Awesome! Thanks for this convincing demo!

added 2 commits

• 4053f7cf - 1 commit from branch ghc/packages:master
• dae282ef - Merge branch 'master' into 'modifyArray'

Compare with previous version

approved this merge request

merged

Thanks for the fast merge!