stricter shapes for setindex

base/broadcast.jl

@@ -614,7 +614,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,57 @@ 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"))
+    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(I::Integer...) = 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(")
@@ -816,6 +813,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)
 
@@ -1039,8 +1042,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
@@ -1561,7 +1564,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

@@ -267,6 +267,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
 
 eltype(::Type{Tuple{}}) = Bottom

doc/src/manual/arrays.md

@@ -603,17 +603,23 @@ 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
-```
+If any index `I_k` is itself an array, then they select multiple locations to be assigned.
+The right hand side `X` must also be an
+array with the same number of elements as the indices select and 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 +844,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 +853,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

@@ -3219,6 +3219,46 @@ 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;;;]
+    )
+    A = zeros(3, 3, 3)
+    for (i, j, k) in Iterators.product(Iterators.repeated(idx_options, 3)...)
+        I = (i, j, k)
+
+        # 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)
+
+        for _X in (X, X_trail, X_prepend, vX, cX)
+            AI = 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...))
+                @test any(isone, setindex!(A, _X, I..., 1))
+            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
+
 # Throws ArgumentError for negative dimensions in Array
 @test_throws ArgumentError fill('a', -10)