Fix nested broadcasting of BlockedArray

Author: mtfishman

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "1.7.0"
+version = "1.7.1"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/blockbroadcast.jl

@@ -131,8 +131,8 @@ end
 @inline function Broadcast.materialize!(dest, bc::Broadcasted{BS}) where {NDims, BS<:AbstractBlockStyle{NDims}}
     dest_reshaped = ndims(dest) == NDims ? dest : reshape(dest, size(bc))
     bc2 = Broadcast.instantiate(
-            Broadcast.Broadcasted{BS}(bc.f, bc.args,
-                map(combine_blockaxes, axes(dest_reshaped), axes(bc))))
+            Broadcast.flatten(Broadcast.Broadcasted{BS}(bc.f, bc.args,
+                map(combine_blockaxes, axes(dest_reshaped), axes(bc)))))
     copyto!(dest_reshaped, bc2)
     return dest
 end

test/test_blockbroadcast.jl

@@ -67,10 +67,26 @@ using StaticArrays
         @test A .+ 1 .+ B == Vector(A) .+ 1 .+ B == Vector(A) .+ 1 .+ Matrix(B)
 
         @testset "preserve structure" begin
-             x = BlockedArray(1:6, Fill(3,2))
-             @test x + x isa BlockedVector{Int,<:AbstractRange}
-             @test 2x + x isa BlockedVector{Int,<:AbstractRange}
-             @test 2 .* (x .+ 1) isa BlockedVector{Int,<:AbstractRange}
+            x = BlockedArray(1:6, Fill(3,2))
+            @test x + x isa BlockedVector{Int,<:AbstractRange}
+            @test 2x + x isa BlockedVector{Int,<:AbstractRange}
+            @test 2 .* (x .+ 1) isa BlockedVector{Int,<:AbstractRange}
+        end
+
+        @testset "nested in-place broadcast" begin
+            x = BlockedVector(randn(4), [2, 2])
+            y = BlockedVector(randn(4), [2, 2])
+            dest = copy(x)
+            dest .+= 2 .* y
+            @test dest ≈ x + 2y
+        end
+
+        @testset "0-dim nested in-place broadcast" begin
+            x = BlockedArray(randn(()))
+            y = BlockedArray(randn(()))
+            dest = copy(x)
+            dest .+= 2 .* y
+            @test dest ≈ x + 2y
         end
     end