diff --git a/src/array_partition.jl b/src/array_partition.jl index a1220451..9504b7a9 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -572,7 +572,7 @@ end @inline function unpack_args(i, args::Tuple) return (unpack(args[1], i), unpack_args(i, Base.tail(args))...) end -unpack_args(i, args::Tuple{Any}) = (unpack(args[1], i),) +@inline unpack_args(i, args::Tuple{Any}) = (unpack(args[1], i),) unpack_args(::Any, args::Tuple{}) = () ## utils diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index 93a3e7b8..c5b4debf 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -416,3 +416,14 @@ end struct TestIsnanFunctor end (::TestIsnanFunctor)(x) = isnan(x) @test all(TestIsnanFunctor(), AP[[NaN], [NaN]]) + +@testset "Mixed-eltype in-place broadcast does not allocate" begin + # unpack_args(i, ::Tuple{Any}) was missing @inline, forcing a runtime + # dispatch at the tuple-recursion base case whenever the partitions + # have different element types. + f!(a, b, c, d) = (@. a = b + c * d; nothing) + mk() = ArrayPartition(rand(10), rand(Float32, 20), rand(5)) + a, b, c, d = mk(), mk(), mk(), mk() + f!(a, b, c, d) + @test (@allocated f!(a, b, c, d)) == 0 +end