Fix NamedArrayPartition Vector{Int} / UnitRange indexing (#583)

[ChrisRackauckas-Claude]

Fixes #583.

Root cause

05faa730b3 (Eliminate all invalidation trees) correctly narrowed getindex(::NamedArrayPartition, ...) from a catch-all args... form to just Int. Anything else now falls through AbstractArray's generic indexing, which routes through similar(x, T, dims).

That exposed a latent bug in NamedArrayPartition's similar overloads. When dims != size(A), the underlying similar(::ArrayPartition, T, dims) already degrades to a plain Vector — a partition layout no longer fits the requested shape (see src/array_partition.jl:72-78). The NamedArrayPartition overloads then tried to wrap that Vector via the inner constructor

NamedArrayPartition{T, A<:ArrayPartition{T}, NT}(::A, ::NT)

which raised the MethodError reported in #583. The pre-05faa730 catch-all args... getindex bypassed similar entirely, so the bug stayed hidden — getindex(::NamedArrayPartition, 1:2) short-circuited to getindex(ArrayPartition(x), 1:2) and never went through the broken similar path.

Fix

Make similar(::NamedArrayPartition, dims) and similar(::NamedArrayPartition, T, dims) mirror ArrayPartition's own degradation: wrap when the inner similar returns an ArrayPartition, otherwise propagate the fallback array as-is. No new method shadows a wider signature than before — using RecursiveArrayTools still produces 0 invalidation trees (verified locally with SnoopCompile.invalidation_trees).

function Base.similar(A::NamedArrayPartition, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
    inner = similar(getfield(A, :array_partition), T, dims)
    inner isa ArrayPartition || return inner
    return NamedArrayPartition(inner, getfield(A, :names_to_indices))
end

x[1:end] keeps returning NamedArrayPartition (full-range slice → dims == size(A) → similar(::ArrayPartition, T, dims) returns an ArrayPartition). x[1:2] and x[[1, 4]] now return a Vector{Float64} (the same fallback ArrayPartition exposes for non-matching dims).

Test plan

Added a regression @testset in test/named_array_partition_tests.jl:

• Issue MWE: NamedArrayPartition(a=ones(2), b=2*ones(3))[1:2] succeeds and returns the right values.
• Vector{Int} indexing: x[[1, 2]], x[[1, 4]].
• Partial UnitRange: x[2:4].
• Full-range slice still preserves type: x[1:end] isa NamedArrayPartition.
• similar(x, Float64, (2,)) isa Vector{Float64}; similar(x, Float64, size(x)) isa NamedArrayPartition.
• Scalar getindex / setindex! untouched.

Local results:

• Pkg.test() passes end-to-end (NamedArrayPartition Tests | 35 35, all other suites green).
• SnoopCompile.invalidation_trees(@snoop_invalidations using RecursiveArrayTools) reports 0 trees — same as the post-05faa730 baseline.

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

Pushed 9cbf6d0 switching to the always-Vector / fully type-stable approach (option 2 from the analysis in #583).

What changed. Added an _unsafe_getindex(::IndexStyle, ::NamedArrayPartition, I::Vararg{Union{Real, AbstractArray}, N}) shortcut mirroring array_partition.jl:317. The indexing path now bypasses similar entirely:

Base.@propagate_inbounds function Base._unsafe_getindex(
        ::IndexStyle, A::NamedArrayPartition,
        I::Vararg{Union{Real, AbstractArray}, N}
    ) where {N}
    shape = Base.index_shape(I...)
    dest = similar(getfield(A, :array_partition).x[1], shape)
    Base._unsafe_getindex!(dest, A, I...)
    return dest
end

The similar(::NAP, dims) / similar(::NAP, T, dims) graceful-degrade fix from the previous commit stays in place for direct callers.

Inferred return types now:

x[1]              :: Float64        (stable)
x[1:2]            :: Vector{Float64} (stable)
x[1:length(x)]    :: Vector{Float64} (stable)
x[[1, 4]]         :: Vector{Float64} (stable)
ap[1:2]           :: Vector{Float64} (matches AP exactly)

Trade-off, called out in the commit message and the test diff. The post-05faa730 assertion typeof(x .+ x[1:end]) <: NamedArrayPartition regresses to <: Vector, which is the pre-05faa730 (v3) behavior. Users who want a NamedArrayPartition back should use similar(x) / copy(x).

Verified:

• Pkg.test() — all suites green; NamedArrayPartition Tests | 40 40.
• @inferred x[1:2] / @inferred x[1:length(x)] / @inferred x[[1, 4]] all return Vector{Float64} (asserted in the regression test).
• SnoopCompile.invalidation_trees(@snoop_invalidations using RecursiveArrayTools) reports 0 trees, same as the post-05faa730 baseline.