diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index af4da6e9..709f686c 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -955,10 +955,35 @@ end Base.@propagate_inbounds function Base.getindex( A::AbstractVectorOfArray{T, N}, I::Vararg{Int, N} ) where {T, N} - @boundscheck checkbounds(A, I...) col = I[N] inner_I = Base.front(I) - u_col = A.u[col] + u = A.u + # Fast path avoiding the O(n) `size(A)` scan below. `size(u_col, d) <= + # size(A, d)` always holds, so a 1-based bound against `size(u_col, d)` is + # sufficient for an in-bounds index into `A`. Must not use + # `checkbounds(u_col, ...)`: that respects a non-1-based inner axis (e.g. + # OffsetArray) and would wrongly accept an index `A`'s 1-based contract + # must reject. N == 1 stores scalars in `u` directly (no inner array), + # where `size(A)` is already O(1); skip the fast path there. + # `A`'s column dimension is likewise 1-based over `length(u)` regardless + # of `u`'s own axes (e.g. an OffsetVector container): `1 <= col <= + # length(u)` matches that contract, and `checkbounds(Bool, u, col)` is + # still required since `u`'s own axes need not be `1:length(u)`. + if N > 1 && 1 <= col <= length(u) && checkbounds(Bool, u, col) + u_col = @inbounds u[col] + in_inner_bounds = true + for d in 1:(N - 1) + if !(1 <= inner_I[d] <= size(u_col, d)) + in_inner_bounds = false + break + end + end + if in_inner_bounds + return @inbounds u_col[inner_I...] + end + end + @boundscheck checkbounds(A, I...) + u_col = u[col] # Return zero for indices outside ragged storage for d in 1:(N - 1) if inner_I[d] > size(u_col, d) diff --git a/test/Core/basic_indexing.jl b/test/Core/basic_indexing.jl index f3b05cbd..ae5208e9 100644 --- a/test/Core/basic_indexing.jl +++ b/test/Core/basic_indexing.jl @@ -1,6 +1,27 @@ using RecursiveArrayTools, Test using RecursiveArrayToolsShorthandConstructors +# Counts calls to `size` on its wrapped vector, used to check that an inner +# array is not scanned once per column of an `AbstractVectorOfArray`. +const counter = Ref(0) +struct CountingVec{T} <: AbstractVector{T} + data::Vector{T} +end +Base.size(c::CountingVec) = (counter[] += 1; size(c.data)) +Base.getindex(c::CountingVec, i::Int) = c.data[i] +Base.IndexStyle(::Type{<:CountingVec}) = IndexLinear() + +# A vector whose own indices do not start at 1, e.g. `OffsetVec(v, -1)` has +# axes `0:length(v)-1`, like `OffsetArrays.OffsetArray(v, 0:length(v)-1)`. +struct OffsetVec{T} <: AbstractVector{T} + data::Vector{T} + offset::Int +end +Base.axes(o::OffsetVec) = (o.offset .+ (1:length(o.data)),) +Base.size(o::OffsetVec) = size(o.data) +Base.getindex(o::OffsetVec, i::Int) = o.data[i - o.offset] +Base.IndexStyle(::Type{<:OffsetVec}) = IndexLinear() + # Example Problem recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] testa = cat(recs..., dims = 2) @@ -394,3 +415,62 @@ x = VectorOfArray(StructArray{SVector{1, Float64}}(ntuple(_ -> [1.0, 2.0], 1))) y = 2 * x @. x = y @test all(all.(y .== x)) + +@testset "getindex O(1) fast path keeps BoundsError/ragged semantics" begin + # count `size` calls on the inner arrays to check the fast path is O(1) + # without relying on timing + counter[] = 0 + n = 1000 + A = VectorOfArray([CountingVec(rand(2)) for _ in 1:n]) + @test A[2, 7] == A.u[7][2] + @test counter[] == 1 # exactly the selected column's own size, not one per inner array + + # ragged: zero-padding inside the max size is unchanged + R = VA[[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]] + @test R[1, 1] == 1.0 + @test R[2, 1] == 2.0 + @test R[3, 2] == 0.0 # inner array 2 (length 1) padded with zero at row 3 + @test R[1, 3] == 4.0 + + # truly out-of-range indices still throw, on both dimensions + @test_throws BoundsError R[4, 1] # row 4 exceeds the global max size (3) + @test_throws BoundsError R[1, 4] # column 4 exceeds the number of inner arrays + @test_throws BoundsError R[0, 1] + @test_throws BoundsError R[1, 0] +end + +@testset "getindex fast path rejects a non-1-based inner axis (OffsetArray)" begin + # minimal non-1-based AbstractVector, avoiding an OffsetArrays dependency + O0 = VA[OffsetVec([1.0, 2.0, 3.0], -1), OffsetVec([4.0, 5.0, 6.0], -1)] # axes 0:2 + O2 = VA[OffsetVec([1.0, 2.0, 3.0], 1), OffsetVec([4.0, 5.0, 6.0], 1)] # axes 2:4 + @test_throws BoundsError O0[0, 1] + @test_throws BoundsError O2[4, 1] +end + +@testset "getindex fast path rejects a non-1-based VA.u container" begin + # VectorOfArray(::AbstractVector{VT}) stores a non-1-based container (like + # OffsetVec) as-is; this is the same column-index mistake as the inner-axis + # case above, one level up. + cols = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]] + C0 = VectorOfArray(OffsetVec(cols, -1)) # axes 0:3 + C2 = VectorOfArray(OffsetVec(cols, 1)) # axes 2:5 + + @test_throws BoundsError C0[1, 0] + @test_throws BoundsError C0[2, 0] + @test_throws BoundsError C0[CartesianIndex(1, 0)] + @test_throws BoundsError C2[1, 5] + @test_throws BoundsError C2[2, 5] + + # a column inside both `1:length(u)` and `u`'s own axes still returns the + # stored value + @test C0[1, 1] == 3.0 + @test C0[1, 3] == 7.0 + @test C2[1, 2] == 1.0 + @test C2[1, 4] == 5.0 + + # a column inside `1:length(u)` but outside `u`'s own axes still throws, + # matching master (indexing `u` itself fails natively there) + @test_throws BoundsError C0[1, 4] + @test_throws BoundsError C2[1, 1] + @test_throws BoundsError C2[1, 6] +end