stricter shapes for setindex

NEWS.md

@@ -22,6 +22,11 @@ New language features
 Language changes
 ----------------
 
+  - The implementation of non-scalar indexed assignment (`A[I...] = X`) now matches its
+    documentation, throwing `DimensionMismatch` on incompatible shapes that the previous
+    implementation would permit, and allowing non-1-based axes on the right-hand value (`X`)
+    that would previously have been forbidden ([#59025]).
+
 Compiler/Runtime improvements
 -----------------------------
 

base/array.jl

@@ -1073,13 +1073,10 @@ function setindex!(A::Array, X::AbstractArray, I::AbstractVector{Int})
     @_propagate_inbounds_meta
     @boundscheck setindex_shape_check(X, length(I))
     @boundscheck checkbounds(A, I)
-    require_one_based_indexing(X)
     X′ = unalias(A, X)
     I′ = unalias(A, I)
-    count = 1
-    for i in I′
-        @inbounds A[i] = X′[count]
-        count += 1
+    for (j, i) in zip(eachindex(X′), I′)
+        @inbounds A[i] = X′[j]
     end
     return A
 end

base/broadcast.jl

@@ -611,7 +611,7 @@ Base.@propagate_inbounds _newindex(ax::Tuple{}, I::Tuple{}) = ()
 end
 
 Base.@propagate_inbounds function Base.getindex(bc::Broadcasted, Is::Vararg{Union{Integer,CartesianIndex},N}) where {N}
-    I = to_index(Base.IteratorsMD.flatten(Is))
+    I = to_index(Base.flatten(Is))
     _getindex(IndexStyle(bc), bc, I)
 end
 @inline function _getindex(::IndexStyle, bc, I)

base/indices.jl

@@ -202,76 +202,58 @@ function promote_shape(a::Indices, b::Indices)
 end
 
 function throw_setindex_mismatch(X, I)
-    if length(I) == 1
-        throw(DimensionMismatch("tried to assign $(length(X)) elements to $(I[1]) destinations"))
+    @noinline
+    pI = filter(!isnegative, I)
+    if length(pI) == 1
+        throw(DimensionMismatch("tried to assign $(length(X)) elements to $(pI[1]) destinations"))
     else
-        throw(DimensionMismatch("tried to assign $(dims2string(size(X))) array to $(dims2string(I)) destination"))
+        throw(DimensionMismatch("tried to assign $(dims2string(size(X))) array to $(dims2string(pI)) destination"))
     end
 end
 
-# check for valid sizes in A[I...] = X where X <: AbstractArray
-# we want to allow dimensions that are equal up to permutation, but only
-# for permutations that leave array elements in the same linear order.
-# those are the permutations that preserve the order of the non-singleton
-# dimensions.
-function setindex_shape_check(X::AbstractArray, I::Integer...)
-    li = ndims(X)
-    lj = length(I)
-    i = j = 1
-    while true
-        ii = length(axes(X,i))
-        jj = I[j]
-        if i == li || j == lj
-            while i < li
-                i += 1
-                ii *= length(axes(X,i))
-            end
-            while j < lj
-                j += 1
-                jj *= I[j]
-            end
-            if ii != jj
-                throw_setindex_mismatch(X, I)
-            end
-            return
-        end
-        if ii == jj
-            i += 1
-            j += 1
-        elseif ii == 1
-            i += 1
-        elseif jj == 1
-            j += 1
+const IntegerOrTuple = Union{Integer, Tuple{Vararg{Integer}}}
+
+_nnprod() = 1
+_nnprod(i::Integer, I::Integer...) =
+    (i == -1) ? _nnprod(I...) : i * _nnprod(I...)
+
+_trailing_dropped() = true
+_trailing_dropped(i::Integer, I::Integer...)  = (i == -1) && _trailing_dropped(I...)
+
+_shapes_match(::Bool, ::Tuple{}) = true
+_shapes_match(isfirstdim::Bool, sz) = _shapes_match(isfirstdim, (), sz...)
+_shapes_match(isfirstdim::Bool, sz::Tuple{}, i::Integer, I::Integer...) =
+    isone(abs(i)) && _shapes_match(isfirstdim, sz, I...)
+function _shapes_match(isfirstdim, sz, i::Integer, I::Integer...)
+    if i == -1
+        return _shapes_match(isfirstdim, sz, I...)
+    else
+        if isfirstdim && _trailing_dropped(I...)
+            # sz comes from a call to size(X) and never contains negatives
+            return prod(sz) == i
         else
-            throw_setindex_mismatch(X, I)
+            return (first(sz) == i) && _shapes_match(false, tail(sz), I...)
         end
     end
 end
 
 setindex_shape_check(X::AbstractArray) =
-    (length(X)==1 || throw_setindex_mismatch(X,()))
+    (length(X) == 1 || throw_setindex_mismatch(X, ()))
 
 setindex_shape_check(X::AbstractArray, i::Integer) =
-    (length(X)==i || throw_setindex_mismatch(X, (i,)))
+    (length(X) == i || throw_setindex_mismatch(X, (i,)))
 
-setindex_shape_check(X::AbstractArray{<:Any, 0}, i::Integer...) =
-    (length(X) == prod(i) || throw_setindex_mismatch(X, i))
+setindex_shape_check(X::AbstractArray{<:Any,0}, I::Integer...) =
+    (length(X) == _nnprod(I...) || throw_setindex_mismatch(X, I))
 
-setindex_shape_check(X::AbstractArray{<:Any,1}, i::Integer) =
-    (length(X)==i || throw_setindex_mismatch(X, (i,)))
+setindex_shape_check(X::AbstractArray{<:Any,1}, I::Integer...) =
+    (length(X) == _nnprod(I...) || throw_setindex_mismatch(X, I))
 
-setindex_shape_check(X::AbstractArray{<:Any,1}, i::Integer, j::Integer) =
-    (length(X)==i*j || throw_setindex_mismatch(X, (i,j)))
+setindex_shape_check(X::AbstractArray, I::IntegerOrTuple...) =
+    setindex_shape_check(X, flatten(I)...)
 
-function setindex_shape_check(X::AbstractArray{<:Any,2}, i::Integer, j::Integer)
-    if length(X) != i*j
-        throw_setindex_mismatch(X, (i,j))
-    end
-    sx1 = length(axes(X,1))
-    if !(i == 1 || i == sx1 || sx1 == 1)
-        throw_setindex_mismatch(X, (i,j))
-    end
-end
+setindex_shape_check(X::AbstractArray, I::Integer...) =
+    _shapes_match(true, size(X), I...) || throw_setindex_mismatch(X, I)
 
 setindex_shape_check(::Any...) =
     throw(ArgumentError("indexed assignment with a single value to possibly many locations is not supported; perhaps use broadcasting `.=` instead?"))

base/multidimensional.jl

@@ -10,7 +10,7 @@ module IteratorsMD
     import .Base: +, -, *, (:)
     import .Base: simd_outer_range, simd_inner_length, simd_index, setindex
     import Core: Tuple
-    using .Base: to_index, fill_to_length, tail, safe_tail
+    using .Base: to_index, fill_to_length, tail, safe_tail, flatten
     using .Base: IndexLinear, IndexCartesian, AbstractCartesianIndex,
         ReshapedArray, ReshapedArrayLF, OneTo, Fix1
     using .Base.Iterators: Reverse, PartitionIterator
@@ -87,9 +87,6 @@ module IteratorsMD
     # Un-nest passed CartesianIndexes
     CartesianIndex{N}(index::CartesianIndex{N}) where {N} = index
     CartesianIndex(index::Union{Integer, CartesianIndex}...) = CartesianIndex(flatten(index))
-    flatten(::Tuple{}) = ()
-    flatten(I::Tuple{Any}) = Tuple(I[1])
-    @inline flatten(I::Tuple) = (Tuple(I[1])..., flatten(tail(I))...)
     CartesianIndex(index::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = CartesianIndex(index...)
     function show(io::IO, i::CartesianIndex)
         print(io, "CartesianIndex(")
@@ -819,6 +816,12 @@ index_shape() = ()
 @inline index_shape(::Real, rest...) = index_shape(rest...)
 @inline index_shape(A::AbstractArray, rest...) = (axes(A)..., index_shape(rest...)...)
 
+index_shape_nested() = ()
+
+# -1 used as signal for dropped dimension. 0 is actually a valid index length
+@inline index_shape_nested(::Real, rest...) = (-1, index_shape_nested(rest...)...)
+@inline index_shape_nested(A::AbstractArray, rest...) = (size(A), index_shape_nested(rest...)...)
+
 """
     LogicalIndex(mask)
 
@@ -1041,8 +1044,8 @@ function _generate_unsafe_setindex!_body(N::Int)
     quote
         x′ = unalias(A, x)
         @nexprs $N d->(I_d = unalias(A, I[d]))
-        idxlens = @ncall $N index_lengths I
-        @ncall $N setindex_shape_check x′ (d->idxlens[d])
+        idxsigs = @ncall $N index_shape_nested I
+        @ncall $N setindex_shape_check x′ (d->idxsigs[d])
         X = eachindex(x′)
         Xy = _prechecked_iterate(X)
         @inbounds @nloops $N i d->I_d begin
@@ -1523,7 +1526,8 @@ end
     N = length(I)
     quote
         idxlens = @ncall $N index_lengths I0 d->I[d]
-        @ncall $N setindex_shape_check X idxlens[1] d->idxlens[d+1]
+        idxsigs = @ncall $N index_shape_nested I0 d->I[d]
+        @ncall $N setindex_shape_check X idxsigs[1] d->idxsigs[d+1]
         isempty(X) && return B
         f0 = indexoffset(I0)+1
         l0 = idxlens[1]

base/tuple.jl

@@ -272,6 +272,10 @@ end
 first(::Tuple{}) = throw(ArgumentError("tuple must be non-empty"))
 first(t::Tuple) = t[1]
 
+flatten(::Tuple{}) = ()
+flatten(I::Tuple{Any}) = Tuple(I[1])
+flatten(I::Tuple) = (@inline; (Tuple(I[1])..., flatten(tail(I))...))
+
 # eltype
 
 # the <: here makes the runtime a bit more complicated (needing to check isdefined), but really helps inference

doc/src/manual/arrays.md

@@ -598,22 +598,30 @@ where each `I_k` may be a scalar integer, an array of integers, or any other
 ranges of the form `a:c` or `a:b:c` to select contiguous or strided
 subsections, and arrays of booleans to select elements at their `true` indices.
 
-If all indices `I_k` are integers, then the value in location `I_1, I_2, ..., I_n` of `A` is
-overwritten with the value of `X`, [`convert`](@ref)ing to the
-[`eltype`](@ref) of `A` if necessary.
-
-
-If any index `I_k` is itself an array, then the right hand side `X` must also be an
-array with the same shape as the result of indexing `A[I_1, I_2, ..., I_n]` or a vector with
-the same number of elements. The value in location `I_1[i_1], I_2[i_2], ..., I_n[i_n]` of
-`A` is overwritten with the value `X[i_1, i_2, ..., i_n]`, converting if necessary. The
-element-wise assignment operator `.=` may be used to [broadcast](@ref Broadcasting) `X`
-across the selected locations:
-
-
-```
-A[I_1, I_2, ..., I_n] .= X
-```
+Just as [indexing](@ref man-array-indexing) `A[I_1, I_2, ..., I_n]` returns either a single
+element (if all indices are scalar) or an array (if any index is nonscalar), the indexed
+assignment of `A[I_1, I_2, ..., I_n] = X` assigns either a single value to a single location
+or an array of value(s) to an array of location(s) within the array `A`.
+
+If all indices `I_k` are scalar, then the value in location `I_1, I_2, ..., I_n` of `A` is assigned
+to the value of `X`, [`convert`](@ref)ing to the [`eltype`](@ref) of `A` if necessary. Otherwise
+if any index is nonscalar, the right hand side `X` must be an array with the same number of
+elements as the number of locations selected by the indices and it must have a compatible shape.
+Each value in `X` is assigned into the corresponding location in `A[I_1, I_2, ..., I_n]`, following
+the rules of [linear indexing and omitted/extra indices](@ref man-number-of-indices).
+
+That is, `X` must be the same shape as the result of indexing `A[I_1, I_2, ..., I_n]` or
+either may be a vector. In the simple case where all the supplied indices are vectors and
+the shapes match, each value in location `I_1[i_1], I_2[i_2], ..., I_n[i_n]` of
+`A` is overwritten with the value `X[i_1, i_2, ..., i_n]`, converting if necessary. If either
+side is a vector, then the elements are assigned with column-major [linear indexing](@ref man-linear-indexing).
+
+The element-wise assignment operator `.=` may be used to [broadcast](@ref Broadcasting) `X`
+across multiple selected locations with `A[I_1, I_2, ..., I_n] .= X`. In the common case where the
+concatenated axes of the indices matches the axes of `X`, this behaves like the nonscalar
+indexed assignment described above (that is, `=` without the broadcasting `.`). Unlike indexed
+assignment, however, any mismatched axes will be broadcast (in accordance with broadcast's rules)
+and linear indexing is not supported.
 
 Just as in [Indexing](@ref man-array-indexing), the `end` keyword may be used
 to represent the last index of each dimension within the indexing brackets, as
@@ -838,7 +846,7 @@ julia> x[vec(mask)] == x[mask] # we can also index with a single Boolean vector
 true
 ```
 
-### Number of indices
+### [Number of indices](@id man-number-of-indices)
 
 #### Cartesian indexing
 
@@ -847,7 +855,7 @@ index selects the position(s) in its particular dimension. For example, in the t
 array `A = rand(4, 3, 2)`, `A[2, 3, 1]` will select the number in the second row of the third
 column in the first "page" of the array. This is often referred to as _cartesian indexing_.
 
-#### Linear indexing
+#### [Linear indexing](@id man-linear-indexing)
 
 When exactly one index `i` is provided, that index no longer represents a location in a
 particular dimension of the array. Instead, it selects the `i`th element using the

test/arrayops.jl

@@ -3256,6 +3256,53 @@ end
     @test setindex!(zeros(2,2), fill(1.0), CI0, 1, 1) == [1.0 0.0; 0.0 0.0]
 end
 
+@testset "setindex! at mismatched shapes" begin
+    idx_options = (
+        1,
+        1:1,
+        2:2,
+        1:0,
+        1:3,
+        [1,2,3],
+        [1 2 3],
+        [1;2;;;],
+        OffsetArray(1:3, -1),
+    )
+    A = zeros(3, 3, 3)
+    seen = Set{UInt}()
+    for (i, j, k) in Iterators.product(Iterators.repeated(idx_options, 3)...)
+        id = hash(sort([i, j, k]; by=hash))
+        in!(id, seen) && continue
+        for I in ((i, j, k), (k, j, i))
+            # don't test scalar setindex here
+            all(x -> x isa Int, I) && continue
+
+            X = ones((length(i) for i in I if i != 1)...)
+            X_trail = reshape(X, size(X)..., 1)
+            X_prepend = reshape(X, 1, size(X)...)
+            vX = vec(X)
+            cX = reshape(vX, length(vX), 1)
+            oX = OffsetArray(X, (size(X) .- 2*(ndims(X) - 1))...)
+
+            for _X in (X, X_trail, X_prepend, vX, cX, oX)
+                AI = view(A, I...)
+
+                if ((isone(ndims(_X)) || isone(ndims(AI))) && (length(_X) == length(AI))) ||
+                    all(d -> size(_X, d) == size(AI, d), 1:max(ndims(_X), ndims(AI)))
+                    @test any(isone, setindex!(A, _X, I...)) || isempty(_X)
+                    fill!(A, 0)
+                    @test any(isone, setindex!(A, _X, I..., 1)) || isempty(_X)
+                    fill!(A, 0)
+                else
+                    @test_throws DimensionMismatch setindex!(A, _X, I...)
+                    @test_throws DimensionMismatch setindex!(A, _X, I..., 1)
+                    @test_throws DimensionMismatch setindex!(A, _X, I..., :)
+                end
+            end
+        end
+    end
+end
+
 @testset "conditionally-throwing version of `checkbounds` should return `nothing` if it returns" begin
     # can not just set `typ = Any` because that would include the predicate (non-throwing) methods of `checkbounds`, because `Type{Bool} <: Any`
     for typ in (AbstractString, AbstractArray, Base.AbstractBroadcasted)