Skip to content

Optimize newGenArray a little

meooow requested to merge meooow/array:newGenArray into master
  • Use unsafeNewArray_ instead of newAray_. We know we will fill the array, and newArray_ wastefully initializes it beforehand.
  • Avoid safeIndexing the range when writing the elements. range generates the elements in the required order, so we can simply use an Int counter.

Example benchmark (GHC 9.2.5, -O2):

benchmarking newGenArray(current)
time                 77.49 μs   (77.16 μs .. 77.75 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 77.24 μs   (77.06 μs .. 77.40 μs)
std dev              569.8 ns   (481.4 ns .. 681.8 ns)

benchmarking newGenArray(new)
time                 26.37 μs   (26.32 μs .. 26.43 μs)
                     1.000 R²   (1.000 R² .. 1.000 R²)
mean                 26.34 μs   (26.30 μs .. 26.39 μs)
std dev              140.6 ns   (110.4 ns .. 198.4 ns)
Source
{-# LANGUAGE BangPatterns #-}
import Criterion.Main
import Data.Array.Base
import Data.Array.IO

main :: IO ()
main = defaultMain
  [ bench "newGenArray(current)" $ whnfIO (newGenArray (1,n) pure :: IO (IOUArray Int Int))
  , bench "newGenArray(new)" $ whnfIO (newGenArray2 (1,n) pure :: IO (IOUArray Int Int))
  ]
  where
    !n = 100000

newGenArray :: (MArray a e m, Ix i) => (i,i) -> (i -> m e) -> m (a i e)
newGenArray (l,u) f = do
    marr <- newArray_ (l,u)
    let n = safeRangeSize (l,u)
    sequence_ [ f i >>= unsafeWrite marr (safeIndex (l,u) n i) | i <- range (l,u)]
    return marr

newGenArray2 :: (MArray a e m, Ix i) => (i,i) -> (i -> m e) -> m (a i e)
newGenArray2 bnds f = do
    let !_ = safeRangeSize bnds  -- To error on a negative size
    marr <- unsafeNewArray_ bnds
    let g ix k !i = do
            x <- f ix
            unsafeWrite marr i x
            k (i+1)
    foldr g (const (return ())) (range bnds) 0
    -- Note: We don't worry about the list being too short or too long
    -- because of the Ix law: rangeSize bnds == length (range bnds)
    return marr
Edited by meooow

Merge request reports