Fusion-friendly *fromListN

[meooow25]

Today *fromListN functions look like

https://github.com/haskell/primitive/blob/efde5a7aaef17bc2e3f8f15a522bb64f9146a587/Data/Primitive/Array.hs#L587-L600

This is not fusion-friendly, so simple use cases like arrayFromListN n [1..n] construct the list (though lazily) and are slower than they could be. This is especially bad for PrimArray, for which the list must be created with boxed values which are promptly unboxed when writing to the ByteArray#. The only way for a user to avoid this is to manually write a loop with the mutable version of the array.

The situation can be improved by replacing the recursive function go with foldr. A change looks like

 arrayFromListN :: Int -> [a] -> Array a
 arrayFromListN n l =
   createArray n (die "fromListN" "uninitialized element") $ \sma ->
-    let go !ix [] = if ix == n
+    let z !ix = if ix == n
           then return ()
           else die "fromListN" "list length less than specified size"
-        go !ix (x : xs) = if ix < n
+        f x k !ix = if ix < n
           then do
             writeArray sma ix x
-            go (ix+1) xs
+            k (ix+1)
           else die "fromListN" "list length greater than specified size"
-    in go 0 l
+    in foldr f z l 0
+{-# INLINE arrayFromListN #-}

foldr can then fuse with the build creating the list to get a good loop. The effect is very clear in some simple benchmarks.

  mk1:     OK
    1.74 ms ±  89 μs, 7.6 MB allocated, 1.4 MB copied,  15 MB peak memory
  mk2:     OK
    375  μs ±  13 μs, 2.3 MB allocated, 296 KB copied,  20 MB peak memory, 0.22x
  primMk1: OK
    3.21 ms ± 283 μs,  15 MB allocated,  39 KB copied,  20 MB peak memory
  primMk2: OK
    39.2 μs ± 1.4 μs, 781 KB allocated,  23 B  copied,  20 MB peak memory, 0.01x

Benchmark code

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
import Control.Monad.ST
import Data.Primitive
import Data.Primitive.Array
import Data.Primitive.PrimArray
import Test.Tasty.Bench

main :: IO ()
main = do
  defaultMain
    [ bench "mk1" $ whnf mk1 100000
    , bcompare "mk1" $ bench "mk2" $ whnf mk2 100000
    , bench "primMk1" $ whnf primMk1 100000
    , bcompare "primMk1" $ bench "primMk2" $ whnf primMk2 100000
    ]

mk1 :: Int -> Array Int
mk1 n = arrayFromListN n [1..n]
{-# NOINLINE mk1 #-}

mk2 :: Int -> Array Int
mk2 n = arrayFromListN2 n [1..n]
{-# NOINLINE mk2 #-}

primMk1 :: Int -> PrimArray Int
primMk1 n = primArrayFromListN n [1..n]
{-# NOINLINE primMk1 #-}

primMk2 :: Int -> PrimArray Int
primMk2 n = primArrayFromListN2 n [1..n]
{-# NOINLINE primMk2 #-}

arrayFromListN2 :: Int -> [a] -> Array a
arrayFromListN2 n l =
  createArray n (die "fromListN" "uninitialized element") $ \sma ->
    let z !ix = if ix == n
          then return ()
          else die "fromListN" "list length less than specified size"
        f x k !ix = if ix < n
          then do
            writeArray sma ix x
            k (ix+1)
          else die "fromListN" "list length greater than specified size"
    in foldr f z l 0
{-# INLINE arrayFromListN2 #-}

-- | Create a 'PrimArray' from a list of a known length. If the length
-- of the list does not match the given length, this throws an exception.
primArrayFromListN2 :: Prim a => Int -> [a] -> PrimArray a
primArrayFromListN2 len vs = createPrimArray len $ \arr ->
  let z !ix = if ix == len
        then return ()
        else die "fromListN" "list length less than specified size"
      f a k !ix = if ix < len
        then do
          writePrimArray arr ix a
          k (ix + 1)
        else die "fromListN" "list length greater than specified size"
  in foldr f z vs 0
{-# INLINE primArrayFromListN2 #-}

-- Not yet released
createPrimArray
  :: Prim a => Int -> (forall s. MutablePrimArray s a -> ST s ()) -> PrimArray a
{-# INLINE createPrimArray #-}
createPrimArray 0 _ = emptyPrimArray
createPrimArray n f = runPrimArray $ do
  marr <- newPrimArray n
  f marr
  pure marr

die :: String -> String -> a
die fun problem = error $ "Data.Primitive.Array." ++ fun ++ ": " ++ problem

How do you feel about making this change?