Fix for setindex!! with array values (#16)

* fixed type instability in setindex! with array values

* removed TODO comment

* renamed one is_linear_index to test_linear_index_only + return false
for vector

* fix issue with `possible` for `_setindex!` not handling invalid array values

* move things around and add further description to the code

* fixed comment typo

* fix typos in tests

* drop eltype restrictions

* add a non-<:Real test case

---------

Co-authored-by: Mason Protter <mason.protter@icloud.com>

src/base.jl

@@ -481,13 +481,50 @@ Base.@propagate_inbounds _setindex!(xs, v, I...) = (setindex!(xs, v, I...); xs)
 
 pure(::typeof(_setindex!)) = NoBang._setindex
 possible(::typeof(_setindex!), ::Union{Tuple,NamedTuple,Empty}, ::Vararg) = false
-possible(::typeof(_setindex!), ::C, ::T, ::Vararg) where {C <: AbstractArray, T} =
-    implements(setindex!, C) && promote_type(eltype(C), T) <: eltype(C)
 possible(::typeof(_setindex!), ::C, ::V, ::K) where {C <: AbstractDict, V, K} =
     implements(setindex!, C) &&
     promote_type(keytype(C), K) <: keytype(C) &&
     promote_type(valtype(C), V) <: valtype(C)
 
+# Default implementation for `_setindex!` with `AbstractArray`.
+# But this will return `false` even in cases such as
+#
+#     setindex!!([1, 2, 3], [4, 5, 6], :)
+#
+# because `promote_type(eltype(C), T) <: eltype(C)` is `false`.
+possible(::typeof(_setindex!), ::C, ::T, ::Vararg) where {C <: AbstractArray, T} =
+    implements(setindex!, C) && promote_type(eltype(C), T) <: eltype(C)
+
+# To address this, we specialize on the case where `T<:AbstractArray`.
+# In addition, we need to support a wide range of indexing behaviors:
+#
+# We also need to ensure that the dimensionality of the index is
+# valid, i.e. that we're not returning `true` in cases such as
+#
+#     setindex!!([1, 2, 3], [4, 5], 1)
+#
+# which should return `false`.
+_index_dimension(::Any) = 0
+_index_dimension(::Colon) = 1
+_index_dimension(::AbstractVector) = 1
+_index_dimension(indices::Tuple) = sum(map(_index_dimension, indices))
+
+function possible(
+    ::typeof(_setindex!),
+    ::C,
+    ::T,
+    indices::Vararg
+    ) where {M, C <: AbstractArray{<:Any}, T <: AbstractArray{<:Any,M}}
+    return implements(setindex!, C) &&
+        promote_type(eltype(C), eltype(T)) <: eltype(C) &&
+        # This will still return `false` for scenarios such as
+        #
+        #     setindex!!([1, 2, 3], [4, 5, 6], :, 1)
+        #
+        # which are in fact valid. However, this cases are rare.
+        (_index_dimension(indices) == M || _index_dimension(indices) == 1)
+end
+
 """
     resize!!(vector::AbstractVector, n::Integer) -> vector′
 """

test/test_setindex.jl

@@ -2,6 +2,26 @@ module TestSetIndex
 
 include("preamble.jl")
 
+# Some utility methods for testing `setindex!`.
+test_linear_index_only(::Tuple, ::AbstractArray) = false
+test_linear_index_only(inds::NTuple{1}, ::AbstractArray) = true
+test_linear_index_only(inds::NTuple{1}, ::AbstractVector) = false
+
+replace_colon_with_axis(inds::Tuple, x) = ntuple(length(inds)) do i
+    inds[i] isa Colon ? axes(x, i) : inds[i]
+end
+replace_colon_with_vector(inds::Tuple, x) = ntuple(length(inds)) do i
+    inds[i] isa Colon ? collect(axes(x, i)) : inds[i]
+end
+replace_colon_with_range(inds::Tuple, x) = ntuple(length(inds)) do i
+    inds[i] isa Colon ? (1:size(x, i)) : inds[i]
+end
+replace_colon_with_booleans(inds::Tuple, x) = ntuple(length(inds)) do i
+    inds[i] isa Colon ? trues(size(x, i)) : inds[i]
+end
+
+replace_colon_with_range_linear(inds::NTuple{1}, x::AbstractArray) = inds[1] isa Colon ? (1:length(x),) : inds
+
 @testset begin
     @test setindex!!((1, 2, 3), :two, 2) === (1, :two, 3)
     @test setindex!!((a=1, b=2, c=3), :two, :b) === (a=1, b=:two, c=3)
@@ -13,11 +33,55 @@ include("preamble.jl")
 end
 
 @testset "mutation" begin
-    @testset for args in [
-        ([1, 2, 3], 20, 2),
-        (Dict(:a=>1, :b=>2), 10, :a),
+    @testset "without type expansion" begin
+        for args in [
+            ([1, 2, 3], 20, 2),
+            (Dict(:a => 1, :b => 2), 10, :a),
+            ([(1,),(2,),(3,)], [(4,), (5,), (6,)], :, 1)
+        ]
+            @test setindex!!(args...) === args[1]
+        end
+    end
+
+    @testset "with type expansion" begin
+        @test setindex!!([1, 2, 3], [4, 5], 1) == [[4, 5], 2, 3]
+        @test setindex!!([1, 2, 3], [4, 5, 6], :, 1) == [4, 5, 6]
+    end
+end
+
+@testset "slices" begin
+    @testset "$(typeof(x)) with $(src_idx)" for (x, src_idx) in [
+        # Vector.
+        (randn(2), (:,)),
+        (randn(2), (1:2,)),
+        # Matrix.
+        (randn(2, 3), (:,)),
+        (randn(2, 3), (:, 1)),
+        (randn(2, 3), (:, 1:3)),
+        # 3D array.
+        (randn(2, 3, 4), (:, 1, :)),
+        (randn(2, 3, 4), (:, 1:3, :)),
+        (randn(2, 3, 4), (1, 1:3, :)),
     ]
-        @test setindex!!(args...) === args[1]
+        # Base case.
+        @test @inferred(setindex!!(x, x[src_idx...], src_idx...)) === x
+
+        # If we have `Colon` in the index, we replace this with other equivalent indices.
+        if any(Base.Fix2(isa, Colon), src_idx)
+            if test_linear_index_only(src_idx, x)
+                # With range instead of `Colon`.
+                @test @inferred(setindex!!(x, x[src_idx...], replace_colon_with_range_linear(src_idx, x)...)) === x
+            else
+                # With axis instead of `Colon`.
+                @test @inferred(setindex!!(x, x[src_idx...], replace_colon_with_axis(src_idx, x)...)) === x
+                # With range instead of `Colon`.
+                @test @inferred(setindex!!(x, x[src_idx...], replace_colon_with_range(src_idx, x)...)) === x
+                # With vectors instead of `Colon`.
+                @test @inferred(setindex!!(x, x[src_idx...], replace_colon_with_vector(src_idx, x)...)) === x
+                # With boolean index instead of `Colon`.
+                @test @inferred(setindex!!(x, x[src_idx...], replace_colon_with_booleans(src_idx, x)...)) === x
+            end
+        end
     end
 end