Optimize newGenArray a little
- 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