Fix mutability check for array multiplication (#272)

src/interface.jl

@@ -281,6 +281,37 @@ function mutability(x, op, args::Vararg{Any,N}) where {N}
  return mutability(typeof(x), op, typeof.(args)...)
 end
 
+# As a special case, we need to check that the shapes of matrix multiplication
+# match in order for the array to be mutable.
+function mutability(
+ x::AbstractArray,
+ op::typeof(*),
+ args::Vararg{Any,N},
+) where {N}
+ is_mutable = mutability(typeof(x), op, typeof.(args)...) == IsMutable()
+ if is_mutable && size(x) == _size_after_multiply(size.(args)...)
+ return IsMutable()
+ end
+ return IsNotMutable()
+end
+
+function _size_after_multiply(x::NTuple{M,Int}, y::NTuple{N,Int}) where {N,M}
+ if x[end] != y[1]
+ return nothing
+ end
+ return (x[1:end-1]..., y[2:end]...)
+end
+
+_size_after_multiply(::Tuple{}, rhs::NTuple{M,Int}) where {M} = rhs
+_size_after_multiply(lhs::NTuple{M,Int}, ::Tuple{}) where {M} = lhs
+_size_after_multiply(::Tuple{}, ::Tuple{}) = ()
+_size_after_multiply(::Nothing, ::Any) = nothing
+
+function _size_after_multiply(x::NTuple{M,Int}, y::Vararg{Any,N}) where {N,M}
+ head = _size_after_multiply(x, Base.first(y))
+ return _size_after_multiply(head, Base.tail(y)...)
+end
+
 mutability(::Type) = IsNotMutable()
 
 """

test/interface.jl

@@ -187,3 +187,22 @@ end
  end
  end
 end
+
+@testset "issue_271_mutability" begin
+ a = 1
+ x = [1; 2;;]
+ y = [1 2; 3 4]
+ z = [1 2; 3 4; 5 6]
+ @test MA.mutability(x, *, x, x') == MA.IsNotMutable()
+ @test MA.mutability(x, *, x', x) == MA.IsNotMutable()
+ @test MA.mutability(x, *, x, a, x') == MA.IsNotMutable()
+ @test MA.mutability(x, *, x', a, x) == MA.IsNotMutable()
+ @test MA.mutability(y, *, y, y) == MA.IsMutable()
+ @test MA.mutability(y, *, y, y, y) == MA.IsMutable()
+ @test MA.mutability(y, *, y, a, y') == MA.IsMutable()
+ @test MA.mutability(y, *, y', a, y) == MA.IsMutable()
+ @test MA.mutability(y, *, a, a, y) == MA.IsMutable()
+ @test MA.mutability(y, *, y, z', z) == MA.IsMutable()
+ @test MA.mutability(z, *, z, z) == MA.IsNotMutable()
+ @test MA.mutability(z, *, z, z, y) == MA.IsNotMutable()
+end

test/matmul.jl

@@ -142,11 +142,11 @@ end
  @test MA.mul(A, x) == BigInt[3; 3; 3]
  @test MA.mul_to!!(y, A, x) == BigInt[3; 3; 3] && y == BigInt[3; 3; 3]
  @test_throws DimensionMismatch MA.mul(BigInt[1 1; 1 1], BigInt[])
- @test_throws DimensionMismatch MA.mul_to!!(
- BigInt[],
-  BigInt[1 1; 1 1],
-  BigInt[1; 1],
- )
+ @test MA.mul_to!!(BigInt[], BigInt[1 1; 1 1], BigInt[1; 1]) ==
+  BigInt[2, 2]
+ z = BigInt[0, 0]
+ @test MA.mul_to!!(z, BigInt[1 1; 1 1], BigInt[1; 1]) === z
+ @test z == BigInt[2, 2]
 
  @testset "mutability" begin
  alloc_test(() -> MA.promote_operation(*, typeof(A), typeof(x)), 0)
@@ -219,11 +219,11 @@ end
  BigInt[1 1; 1 1],
  zeros(BigInt, 1, 1),
  )
- @test_throws DimensionMismatch MA.mul_to!!(
+ @test MA.mul_to!!(
  zeros(BigInt, 1, 1),
  BigInt[1 1; 1 1],
  zeros(BigInt, 2, 1),
- )
+ ) == zeros(BigInt, 2, 1)
 
  @testset "mutability" begin
  alloc_test(() -> MA.promote_operation(*, typeof(A), typeof(B)), 0)
@@ -422,3 +422,16 @@ Base.:*(m::Monomial, ::Monomial) = m
  @test T == typeof(MA.operate(*, a, b))
  end
 end
+
+@testset "Issue_271" begin
+ A = reshape([1, 2], (2, 1))
+ B = [1 2]
+ C = MA.operate!!(*, A, B)
+ @test A == reshape([1, 2], (2, 1))
+ @test B == [1 2]
+ @test C == A * B
+ D = MA.operate!!(*, B, A)
+ @test A == reshape([1, 2], (2, 1))
+ @test B == [1 2]
+ @test D == B * A
+end