Fix type instability for vec[1, end] indexing

[ChrisRackauckas-Claude]

Summary

• Fixes Type instability for vec[1, end] #525
• Adds a specialized getindex(A, i::Int, re::RaggedEnd) method for type-stable access when using end indexing
• Adds a helper function _ragged_getindex_int_col as a function barrier for the Int column case
• Adds an early check in _ragged_getindex for the RaggedEnd sentinel case

Problem

When accessing a VectorOfArray with [1, end], the result was type-unstable (returning Any):

vecvec = VectorOfArray([fill(2.0, 2) for i in 1:10])
access_end(vec) = vec[1, end]
@code_warntype access_end(vecvec)  # Body::Any

The issue was that lastindex(VA, d) returns RaggedEnd(0, offset) as a sentinel for an already-resolved index, but the getindex dispatch couldn't infer the return type because cols.dim == 0 is a runtime check.

Solution

Add a specialized method for the common vec[i, end] pattern that handles the RaggedEnd type directly:

Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray, i::Int, re::RaggedEnd)
    if re.dim == 0
        return A.u[re.offset][i]
    else
        col = lastindex(A.u)
        resolved_idx = lastindex(A.u[col], re.dim) + re.offset
        return A.u[col][i, resolved_idx]
    end
end

After this fix:

@code_warntype access_end(vecvec)  # Body::Float64

Test plan

• Added test for type stability of vec[1, end] pattern
• All existing tests pass

🤖 Generated with Claude Code

[ChrisRackauckas]

@JoshuaLampert can you take a look?

[JoshuaLampert]

I will try to take a look tomorrow. Today is busy for me.

[Ickaser]

This does fix the case I ran into:

vecvec = VectorOfArray([fill(2.0, 2) for i in 1:10])
access_end(vec) = vec[1, end]
@inferred access_end(vecvec)  # Returns Float64 as expected

There is still another type instability if you switch the indices, though:

vecvec = VectorOfArray([fill(2.0, 2) for i in 1:10])
access_end(vec) = vec[end, 1]
@inferred access_end(vecvec)  # Errors because inference fails
@inferred getindex(vecvec, lastindex(vecvec), 1) # Same error

[JoshuaLampert]

Suggested change

	last_row = lastindex(testva_end)
	@inferred testva_end[last_row, 1]
	@test testva_end[last_row, 1 == 2.0]
	last_row = lastindex(testva_end, 1)
	@inferred testva_end[last_row, 1]
	@test testva_end[last_row, 1] == 2.0

I guess, this is what you wanted to test @Ickaser? lastindex(testva_end) is 10 here, but we want 2, which is what we get by lastindex(testva_end, 1). The last line is a typo I guess.

[Ickaser]

The last line is a typo which should be fixed, but the lastindex(testva_end) was on purpose; I was trying to make an inference test for testva_end[end, 1], which currently isn't inferred. If there is a better way to test that, feel free to replace.

[JoshuaLampert]

Yes, but lastindex(testva_end) returns the lastindex of the underlying vector testva_end.u, which is the dimension you put last when indexing (here 10). When you call testva_end[end, 1] you want the end to resolve to lastindex of the inner dimension, which would be lastindex(testva_end, 1) (here 2). Otherwise there is a BoundsError as you can see here. The inference test for

last_row = lastindex(testva_end, 1)
@inferred testva_end[last_row, 1]

currently fails (but @test testva_end[last_row, 1] == 2.0 passes in contrast to using last_row = lastindex(testva_end)).

[JoshuaLampert]

@ChrisRackauckas, could you let Claude work on fixing inference also for

last_row = lastindex(testva_end, 1)
@inferred testva_end[last_row, 1]

Otherwise, this looks good to me.

[ChrisRackauckas-Claude]

Added type-stable getindex for vec[end, col] pattern as requested.

Changes:

• Added specialized getindex(A::AbstractVectorOfArray, re::RaggedEnd, col::Int) method to handle the vec[end, col] case
• Fixed test to use lastindex(testva_end, 1) for row index (returns RaggedEnd(1, 0)) instead of lastindex(testva_end) (which returns the length of the underlying vector)
• Fixed typo in test assertion (testva_end[last_row, 1 == 2.0] → testva_end[last_row, 1] == 2.0)

Test verification:

testva_end = VectorOfArray([fill(2.0, 2) for i in 1:10])
last_row = lastindex(testva_end, 1)
@inferred testva_end[last_row, 1]  # Now returns Float64, was Any before

All tests pass locally.

[ChrisRackauckas-Claude]

Rebased onto latest master (c98f542) and squashed into a single clean commit.

The PR now includes both type-stable getindex methods:

• getindex(A, i::Int, re::RaggedEnd) for vec[i, end]
• getindex(A, re::RaggedEnd, col::Int) for vec[end, col]

All tests pass.

[JoshuaLampert]

Thanks