diff --git a/base/array.jl b/base/array.jl index dd7e4a9d40b37..59888424fe0da 100644 --- a/base/array.jl +++ b/base/array.jl @@ -1026,13 +1026,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 diff --git a/base/broadcast.jl b/base/broadcast.jl index b86baf08ddfe0..8b04927a85cb1 100644 --- a/base/broadcast.jl +++ b/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) diff --git a/base/indices.jl b/base/indices.jl index 88e48a2b331ee..51843d56de034 100644 --- a/base/indices.jl +++ b/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(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?")) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 9512d5d6dbe8b..8fb646fcb3e95 100644 --- a/base/multidimensional.jl +++ b/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] diff --git a/base/tuple.jl b/base/tuple.jl index 62a07f7ecf6a2..58278ab61b26e 100644 --- a/base/tuple.jl +++ b/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 # the <: here makes the runtime a bit more complicated (needing to check isdefined), but really helps inference diff --git a/doc/src/manual/arrays.md b/doc/src/manual/arrays.md index ba2d261301b40..857235fc5c799 100644 --- a/doc/src/manual/arrays.md +++ b/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 diff --git a/test/arrayops.jl b/test/arrayops.jl index 7d018e991c3a8..50d8702a22fdd 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -3220,6 +3220,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 + # Throws ArgumentError for negative dimensions in Array @test_throws ArgumentError fill('a', -10)