fix broadcast allocations

Author: ChrisRackauckas

src/array_partition.jl

@@ -230,15 +230,15 @@ ArrayPartitionStyle(::S, ::Val{N}) where {S,N} = ArrayPartitionStyle(S(Val(N)))
 ArrayPartitionStyle(::Val{N}) where N = ArrayPartitionStyle{Broadcast.DefaultArrayStyle{N}}()
 
 # promotion rules
-function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle}
+@inline function Broadcast.BroadcastStyle(::ArrayPartitionStyle{AStyle}, ::ArrayPartitionStyle{BStyle}) where {AStyle, BStyle}
     ArrayPartitionStyle(Broadcast.BroadcastStyle(AStyle(), BStyle()))
 end
-Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{0}) = Broadcast.DefaultArrayStyle{1}()
+Broadcast.BroadcastStyle(::ArrayPartitionStyle{Style}, ::Broadcast.DefaultArrayStyle{0}) where Style = ArrayPartitionStyle{Style}()
 Broadcast.BroadcastStyle(::ArrayPartitionStyle, ::Broadcast.DefaultArrayStyle{N}) where N = Broadcast.DefaultArrayStyle{N}()
 
 combine_styles(args::Tuple{})         = Broadcast.DefaultArrayStyle{0}()
-combine_styles(args::Tuple{Any})      = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]))
-combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2]))
+@inline combine_styles(args::Tuple{Any})      = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]))
+@inline combine_styles(args::Tuple{Any, Any}) = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), Broadcast.BroadcastStyle(args[2]))
 @inline combine_styles(args::Tuple)   = Broadcast.result_style(Broadcast.BroadcastStyle(args[1]), combine_styles(Base.tail(args)))
 
 function Broadcast.BroadcastStyle(::Type{ArrayPartition{T,S}}) where {T, S}

test/partitions_test.jl

@@ -119,3 +119,17 @@ for sizes in S
   @test all([x[i] == y[i] for i in eachindex(x)])
   @test all([x[i] == y_array[i] for i in eachindex(x)])
 end
+
+# Non-allocating broadcast
+xce0 = ArrayPartition(zeros(2),[0.])
+xcde0 = copy(xce0)
+function foo(y, x)
+	y .= y .+ x
+end
+foo(xcde0, xce0)
+@test 0 == @allocated foo(xcde0, xce0)
+function foo(y, x)
+	y .= y .+ 2 .* x
+end
+foo(xcde0, xce0)
+@test 0 == @allocated foo(xcde0, xce0)