From 4b988447283dfe99d2a111c25c90603575894fde Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 20 Jul 2021 02:42:18 +0400 Subject: [PATCH] Add special methods for Cartesian indexing of CartesianIndices for certain permissible indices (#40344) --- base/multidimensional.jl | 24 +++++++++ base/subarray.jl | 6 +++ test/cartesian.jl | 105 +++++++++++++++++++++++++++++++++++++++ test/ranges.jl | 2 + test/subarray.jl | 19 +++++++ 5 files changed, 156 insertions(+) diff --git a/base/multidimensional.jl b/base/multidimensional.jl index 7c2aa90d6389c..61e6be8602512 100644 --- a/base/multidimensional.jl +++ b/base/multidimensional.jl @@ -353,10 +353,34 @@ module IteratorsMD # AbstractArray implementation Base.axes(iter::CartesianIndices{N,R}) where {N,R} = map(Base.axes1, iter.indices) Base.IndexStyle(::Type{CartesianIndices{N,R}}) where {N,R} = IndexCartesian() + # getindex for a 0D CartesianIndices is necessary for disambiguation + @propagate_inbounds function Base.getindex(iter::CartesianIndices{0,R}) where {R} + CartesianIndex() + end @propagate_inbounds function Base.getindex(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R} CartesianIndex(getindex.(iter.indices, I)) end + # CartesianIndices act as a multidimensional range, so cartesian indexing of CartesianIndices + # with compatible dimensions may be seen as indexing into the component ranges. + # This may use the special indexing behavior implemented for ranges to return another CartesianIndices + @propagate_inbounds function Base.getindex(iter::CartesianIndices{N,R}, + I::Vararg{Union{OrdinalRange{<:Integer, <:Integer}, Colon}, N}) where {N,R} + CartesianIndices(getindex.(iter.indices, I)) + end + @propagate_inbounds function Base.getindex(iter::CartesianIndices{N}, + C::CartesianIndices{N}) where {N} + CartesianIndices(getindex.(iter.indices, C.indices)) + end + + # If dimensions permit, we may index into a CartesianIndices directly instead of constructing a SubArray wrapper + @propagate_inbounds function Base.view(c::CartesianIndices{N}, r::Vararg{Union{OrdinalRange{<:Integer, <:Integer}, Colon},N}) where {N} + getindex(c, r...) + end + @propagate_inbounds function Base.view(c::CartesianIndices{N}, C::CartesianIndices{N}) where {N} + getindex(c, C) + end + ndims(R::CartesianIndices) = ndims(typeof(R)) ndims(::Type{CartesianIndices{N}}) where {N} = N ndims(::Type{CartesianIndices{N,TT}}) where {N,TT} = N diff --git a/base/subarray.jl b/base/subarray.jl index 89a4db4d65790..1524f0bb0950f 100644 --- a/base/subarray.jl +++ b/base/subarray.jl @@ -204,6 +204,12 @@ function view(r1::LinRange, r2::OrdinalRange{<:Integer}) getindex(r1, r2) end +# getindex(r::AbstractRange, ::Colon) returns a copy of the range, and we may do the same for a view +function view(r1::AbstractRange, c::Colon) + @_propagate_inbounds_meta + getindex(r1, c) +end + function unsafe_view(A::AbstractArray, I::Vararg{ViewIndex,N}) where {N} @_inline_meta SubArray(A, I) diff --git a/test/cartesian.jl b/test/cartesian.jl index 8d2651b6f425f..b3cb8315decad 100644 --- a/test/cartesian.jl +++ b/test/cartesian.jl @@ -147,6 +147,14 @@ module TestOffsetArray end @testset "CartesianIndices getindex" begin + @testset "0D array" begin + a = zeros() + c = CartesianIndices(a) + @test a[c] == a + @test c[c] === c + @test c[] == CartesianIndex() + end + @testset "AbstractUnitRange" begin for oinds in [(2, ), (2, 3), (2, 3, 4)] A = rand(1:10, oinds) @@ -159,6 +167,34 @@ end @test all(i->A[i]==A[R[i]], R) @test all(i->A[i]==A[R[i]], collect(R)) @test all(i->i in R, collect(R)) + + # Indexing a CartesianIndices with another CartesianIndices having the same ndims + # forwards the indexing to the component ranges and retains the wrapper + @test R[R] === R + + R_array = collect(R) + + all_onetoone = ntuple(x -> 1:1, Val(ndims(R))) + R2 = R[all_onetoone...] + @test R2 isa CartesianIndices{ndims(R)} + + all_one = ntuple(x -> 1, Val(ndims(R))) + @test R2[all_one...] == R_array[all_one...] + + @test R2 == R_array[all_onetoone...] + + R3 = R[ntuple(x -> Colon(), Val(ndims(R)))...] + @test R3 === R + + # test a mix of Colons and ranges + # up to two leading axes are colons, while the rest are UnitRanges + indstrailing = (1:1 for _ in min(ndims(R), 2)+1:ndims(R)) + R4 = R[(Colon() for _ in 1:min(ndims(R), 2))..., indstrailing...] + @test R4 isa CartesianIndices{ndims(R)} + indsleading = CartesianIndices(axes(A)[1:min(ndims(A), 2)]) + for I in indsleading + @test R4[I, indstrailing...] == R_array[I, indstrailing...] + end end end @@ -173,6 +209,75 @@ end # TODO: A[SR] == A[Linearindices(SR)] should hold for StepRange CartesianIndices @test_broken A[SR] == A[LinearIndices(SR)] + + # Create a CartesianIndices with StepRange indices to test indexing into it + R = CartesianIndices(oinds) + R_array = collect(R) + + all_onetoone = ntuple(x -> 1:1, Val(ndims(R))) + R2 = R[all_onetoone...] + @test R2 isa CartesianIndices{ndims(R)} + + all_one = ntuple(x -> 1, Val(ndims(R))) + @test R2[all_one...] == R_array[all_one...] + @test R2 == R_array[all_onetoone...] + + R3 = R[ntuple(x -> Colon(), Val(ndims(R)))...] + @test R3 === R + + # test a mix of Colons and ranges + # up to two leading axes are colons, while the rest are UnitRanges + indstrailing = (1:1 for _ in min(ndims(R), 2)+1:ndims(R)) + R4 = R[(Colon() for _ in 1:min(ndims(R), 2))..., indstrailing...] + @test R4 isa CartesianIndices{ndims(R)} + indsleading = CartesianIndices(axes(R)[1:min(ndims(R), 2)]) + for I in indsleading + @test R4[I, indstrailing...] == R_array[I, indstrailing...] + end + end + + # CartesianIndices whole indices have a unit step may be their own axes + for oinds in [(1:1:4, ), (1:1:4, 1:1:5), (1:1:4, 1:1:5, 1:1:3)] + R = CartesianIndices(oinds) + @test R[R] === R + # test a mix of UnitRanges and StepRanges + R = CartesianIndices((oinds..., 1:3)) + @test R[R] === R + R = CartesianIndices((1:3, oinds...)) + @test R[R] === R + end + end + + @testset "logical indexing of CartesianIndices with ranges" begin + c = CartesianIndices((1:0, 1:2)) + c2 = c[true:false, 1:2] + @test c2 == c + + for (inds, r) in Any[(1:2, false:true), (1:2, false:true:true), + (1:2:3, false:true), (1:2:3, false:true:true)] + + c = CartesianIndices((inds, 1:2)) + c2 = c[r, 1:2] + @test c2 isa CartesianIndices{ndims(c)} + @test c2[1, :] == c[2, :] + end + + for (inds, r) in Any[(1:1, true:true), (1:1, true:true:true), + (1:1:1, true:true), (1:1:1, true:true:true)] + + c = CartesianIndices((inds, 1:2)) + c2 = c[r, 1:2] + @test c2 isa CartesianIndices{ndims(c)} + @test c2[1, :] == c[1, :] + end + + for (inds, r) in Any[(1:1, false:false), (1:1, false:true:false), + (1:1:1, false:false), (1:1:1, false:true:false)] + + c = CartesianIndices((inds, 1:2)) + c2 = c[r, 1:2] + @test c2 isa CartesianIndices{ndims(c)} + @test size(c2, 1) == 0 end end end diff --git a/test/ranges.jl b/test/ranges.jl index 5f248d618dfef..ad8554d14ea4a 100644 --- a/test/ranges.jl +++ b/test/ranges.jl @@ -1517,6 +1517,8 @@ end @test view(1:10, 1:5) === 1:5 @test view(1:10, 1:2:5) === 1:2:5 @test view(1:2:9, 1:5) === 1:2:9 + @test view(1:10, :) === 1:10 + @test view(1:2:9, :) === 1:2:9 # Ensure we don't hit a fallback `view` if there's a better `getindex` implementation vmt = collect(methods(view, Tuple{AbstractRange, AbstractRange})) diff --git a/test/subarray.jl b/test/subarray.jl index 334211d3e3975..cc8aab94e4c42 100644 --- a/test/subarray.jl +++ b/test/subarray.jl @@ -718,3 +718,22 @@ end s = @view v[1] @test copy(s) == fill([1]) end + +@testset "issue 40314: views of CartesianIndices" begin + c = CartesianIndices((1:2, 1:4)) + @test (@view c[c]) === c + for inds in Any[(1:1, 1:2), (1:1:1, 1:2)] + c2 = @view c[inds...] + @test c2 isa CartesianIndices{2} + for i2 in inds[2], i1 in inds[1] + @test c2[i1, i2] == c[i1, i2] + end + end + for inds in Any[(Colon(), 1:2), (Colon(), 1:1:2)] + c2 = @view c[inds...] + @test c2 isa CartesianIndices{2} + for i2 in inds[2], i1 in axes(c, 1) + @test c2[i1, i2] == c[i1, i2] + end + end +end