Skip to content

Commit

Permalink
Make fromListN functions good consumers
Browse files Browse the repository at this point in the history
...in terms of list fusion.
  • Loading branch information
meooow25 authored and andrewthad committed Nov 19, 2024
1 parent 8260d8f commit 75c5365
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 20 deletions.
20 changes: 15 additions & 5 deletions Data/Primitive/Array.hs
Original file line number Diff line number Diff line change
Expand Up @@ -586,18 +586,28 @@ mapArray' f a =

-- | Create an array from a list of a known length. If the length
-- of the list does not match the given length, this throws an exception.

-- Note [fromListN]
-- ~~~~~~~~~~~~~~~~
-- We want arrayFromListN to be a "good consumer" in list fusion, so we define
-- the function using foldr and inline it to help fire fusion rules.
-- If fusion occurs with a "good producer", it may reduce to a fold on some
-- structure. In certain cases (such as for Data.Set) GHC is not be able to
-- optimize the index to an unboxed Int# (see GHC #24628), so we explicitly use
-- an Int# here.
arrayFromListN :: Int -> [a] -> Array a
{-# INLINE arrayFromListN #-}
arrayFromListN n l =
createArray n (die "fromListN" "uninitialized element") $ \sma ->
let go !ix [] = if ix == n
let z ix# = if I# ix# == n
then return ()
else die "fromListN" "list length less than specified size"
go !ix (x : xs) = if ix < n
f x k = GHC.Exts.oneShot $ \ix# -> if I# ix# < n
then do
writeArray sma ix x
go (ix+1) xs
writeArray sma (I# ix#) x
k (ix# +# 1#)
else die "fromListN" "list length greater than specified size"
in go 0 l
in foldr f z l 0#

-- | Create an array from a list.
arrayFromList :: [a] -> Array a
Expand Down
13 changes: 8 additions & 5 deletions Data/Primitive/ByteArray.hs
Original file line number Diff line number Diff line change
Expand Up @@ -378,17 +378,20 @@ byteArrayFromList xs = byteArrayFromListN (length xs) xs

-- | Create a 'ByteArray' from a list of a known length. If the length
-- of the list does not match the given length, this throws an exception.

-- See Note [fromListN] in Data.Primitive.Array
byteArrayFromListN :: forall a. Prim a => Int -> [a] -> ByteArray
{-# INLINE byteArrayFromListN #-}
byteArrayFromListN n ys = createByteArray (n * sizeOfType @a) $ \marr ->
let go !ix [] = if ix == n
let z ix# = if I# ix# == n
then return ()
else die "byteArrayFromListN" "list length less than specified size"
go !ix (x : xs) = if ix < n
f x k = GHC.Exts.oneShot $ \ix# -> if I# ix# < n
then do
writeByteArray marr ix x
go (ix + 1) xs
writeByteArray marr (I# ix#) x
k (ix# +# 1#)
else die "byteArrayFromListN" "list length greater than specified size"
in go 0 ys
in foldr f z ys 0#

unI# :: Int -> Int#
unI# (I# n#) = n#
Expand Down
13 changes: 8 additions & 5 deletions Data/Primitive/PrimArray.hs
Original file line number Diff line number Diff line change
Expand Up @@ -234,17 +234,20 @@ primArrayFromList vs = primArrayFromListN (L.length vs) vs

-- | 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.

-- See Note [fromListN] in Data.Primitive.Array
primArrayFromListN :: forall a. Prim a => Int -> [a] -> PrimArray a
{-# INLINE primArrayFromListN #-}
primArrayFromListN len vs = createPrimArray len $ \arr ->
let go [] !ix = if ix == len
let z ix# = if I# ix# == len
then return ()
else die "fromListN" "list length less than specified size"
go (a : as) !ix = if ix < len
f a k = GHC.Exts.oneShot $ \ix# -> if I# ix# < len
then do
writePrimArray arr ix a
go as (ix + 1)
writePrimArray arr (I# ix#) a
k (ix# +# 1#)
else die "fromListN" "list length greater than specified size"
in go vs 0
in foldr f z vs 0#

-- | Convert a 'PrimArray' to a list.
{-# INLINE primArrayToList #-}
Expand Down
11 changes: 6 additions & 5 deletions Data/Primitive/SmallArray.hs
Original file line number Diff line number Diff line change
Expand Up @@ -924,18 +924,19 @@ instance (Typeable s, Typeable a) => Data (SmallMutableArray s a) where
-- | Create a 'SmallArray' from a list of a known length. If the length
-- of the list does not match the given length, this throws an exception.
smallArrayFromListN :: Int -> [a] -> SmallArray a
{-# INLINE smallArrayFromListN #-}
smallArrayFromListN n l =
createSmallArray n
(die "smallArrayFromListN" "uninitialized element") $ \sma ->
let go !ix [] = if ix == n
let z ix# = if I# ix# == n
then return ()
else die "smallArrayFromListN" "list length less than specified size"
go !ix (x : xs) = if ix < n
f x k = GHC.Exts.oneShot $ \ix# -> if I# ix# < n
then do
writeSmallArray sma ix x
go (ix + 1) xs
writeSmallArray sma (I# ix#) x
k (ix# +# 1#)
else die "smallArrayFromListN" "list length greater than specified size"
in go 0 l
in foldr f z l 0#

-- | Create a 'SmallArray' from a list.
smallArrayFromList :: [a] -> SmallArray a
Expand Down

0 comments on commit 75c5365

Please sign in to comment.