copyto! for broadcasting (#80)

* copyto! for broadcasting

* tidy up tests and add a few

* bump

Project.toml

@@ -1,7 +1,7 @@
 name = "NamedDims"
 uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.2.9"
+version = "0.2.10"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/broadcasting.jl

@@ -55,7 +55,13 @@ function Broadcast.copy(bc::Broadcasted{NamedDimsStyle{S}}) where S
     L = broadcasted_names(bc)
     return NamedDimsArray{L}(data)
 end
-# TODO: copyto! for broadcasting
+
+function Base.copyto!(dest::AbstractArray, bc::Broadcasted{NamedDimsStyle{S}}) where S
+    inner_bc = unwrap_broadcasted(bc)
+    copyto!(dest, inner_bc)
+    L = unify_names(names(dest), broadcasted_names(bc))
+    return NamedDimsArray{L}(dest)
+end
 
 broadcasted_names(bc::Broadcasted) = broadcasted_names(bc.args...)
 function broadcasted_names(a, bs...)

test/broadcasting.jl

@@ -12,6 +12,10 @@ using Tracker
 
         @test nda .+ nda .+ nda == 3ones(3)
         @test names(nda .+ nda .+ nda) == (:a,)
+
+        # in-place
+        @test names(nda .= 0 .* nda .+ 7) == (:a,)
+        @test unname(nda .= 0 .* nda .+ 7) == 7*ones(3)
     end
 
     @testset "partially named dims" begin
@@ -89,6 +93,29 @@ using Tracker
         @test names(nda .+ ones(1,20)) == (:x, :y, :z)
     end
 
+    @testset "in-place assignment .=" begin
+        ab = NamedDimsArray(rand(2,2), (:a, :b))
+        a_ = NamedDimsArray(rand(2,2), (:a, :_))
+        ba = NamedDimsArray(rand(2,2), (:b, :a))
+        ac = NamedDimsArray(rand(2,2), (:a, :c))
+        z = zeros(2,2)
+
+        # https://github.com/invenia/NamedDims.jl/issues/71
+        @test_throws DimensionMismatch z .= ab .+ ba
+        @test_throws DimensionMismatch z .= ab .+ ac
+        @test_throws DimensionMismatch a_ .= ab .+ ac
+        @test_throws DimensionMismatch ab .= a_ .+ ac
+        @test_throws DimensionMismatch ac .= ab .+ ba
+
+        # check that dest is written into:
+        @test names(z .= ab .+ ba') == (:a, :b)
+        @test z == (ab.data .+ ba.data')
+        @test z isa Array  # has not itself magically gained names
+
+        @test names(z .= ab .+ a_) == (:a, :b)
+        @test names(a_ .= ba' .+ ab) == (:a, :b)
+    end
+
 end
 
 @testset "Competing Wrappers" begin