Make NamedArrayPartition slicing fully type-stable via Vector

The first commit on this branch fixed `x[1:2]` and `x[[1, 4]]` by making
`similar(::NAP, T, dims)` degrade to a Vector when `dims != size(A)`,
mirroring ArrayPartition. That worked, but it left the indexing path
inferring as `Union{NamedArrayPartition, Vector{Float64}}` because
`similar(::ArrayPartition, T, dims)` itself is a Union (the
`dims == size(A)` branch is a runtime check).

Add a `_unsafe_getindex(::IndexStyle, ::NAP, I::Vararg{Union{Real, AbstractArray}})`
shortcut that mirrors the one at array_partition.jl:317. Allocate the
destination directly off the underlying first array and fill it with
`Base._unsafe_getindex!`. The shortcut bypasses `similar` entirely for
the indexing path, so `x[1:2]`, `x[[1, 4]]`, `x[1:length(x)]` all infer
to a clean `Vector{Float64}`.

Trade-off: this regresses the post-05faa730 test
`typeof(x .+ x[1:end]) <: NamedArrayPartition` back to `<: Vector` (the
v3 behavior). That test was added in 05faa73 alongside the
invalidation cleanup, but its only effect was to mask the unstable
small-Union return; restoring v3 semantics here gives full type
stability and matches what ArrayPartition already does. Use
`similar(x)` / `copy(x)` if you want a NamedArrayPartition back.

The `similar(::NAP, dims)` and `similar(::NAP, T, dims)` overloads
keep the graceful-degrade-to-Vector behavior from the previous commit,
so direct `similar(x, T, (2,))` calls (e.g. from downstream library
code) still work.

`SnoopCompile.invalidation_trees(@snoop_invalidations using
RecursiveArrayTools)` still reports 0 trees, full `Pkg.test()` passes,
and the new regression test asserts type stability via `@inferred`.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

src/named_array_partition.jl

@@ -100,6 +100,28 @@ Base.length(x::NamedArrayPartition) = length(ArrayPartition(x))
 # Use concrete index types to avoid invalidating AbstractArray's generic setindex!.
 Base.@propagate_inbounds Base.getindex(x::NamedArrayPartition, i::Int) = ArrayPartition(x)[i]
 Base.@propagate_inbounds Base.setindex!(x::NamedArrayPartition, v, i::Int) = (ArrayPartition(x)[i] = v)
+
+# Indexing with non-scalar indices (UnitRange, Vector{Int}, etc.) goes through
+# AbstractArray's generic path, which routes via `similar(A, T, dims)`. NAP's
+# `similar(::NAP, T, dims)` cannot in general produce a NamedArrayPartition for
+# arbitrary `dims` (the partition layout is fixed by `names_to_indices`), so it
+# falls back to a plain Vector — making the inferred return type a small Union.
+#
+# Mirror ArrayPartition's `_unsafe_getindex` shortcut at `array_partition.jl:317`:
+# allocate the destination directly off the first underlying array and fill it
+# via `_unsafe_getindex!`. The result is always a Vector for non-scalar indexing,
+# so `x[I]` is type-stable. This matches the v3 indexing semantics (`x[1:end]`
+# returns a `Vector`, not a `NamedArrayPartition`); use `similar(x)` /
+# `copy(x)` if you want a NamedArrayPartition back.
+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
 function Base.map(f, x::NamedArrayPartition)
     return NamedArrayPartition(map(f, ArrayPartition(x)), getfield(x, :names_to_indices))
 end

test/named_array_partition_tests.jl

@@ -7,7 +7,7 @@ using RecursiveArrayTools, ArrayInterface, Test
     @test typeof(similar(x)) <: NamedArrayPartition
     @test typeof(similar(x, Int)) <: NamedArrayPartition
     @test x.a ≈ ones(10)
-    @test typeof(x .+ x[1:end]) <: NamedArrayPartition # x[1:end] preserves type
+    @test typeof(x .+ x[1:end]) <: Vector # x[1:end] is a plain Vector (type-stable slicing)
     @test all(x .== x[1:end])
     @test ArrayInterface.zeromatrix(x) isa Matrix
     @test size(ArrayInterface.zeromatrix(x)) == (30, 30)
@@ -39,37 +39,45 @@ using RecursiveArrayTools, ArrayInterface, Test
 end
 
 # Regression test for https://github.com/SciML/RecursiveArrayTools.jl/issues/583:
-# indexing a NamedArrayPartition with a UnitRange / Vector{Int} smaller than
-# the whole array used to throw a MethodError because `similar(::NAP, T, dims)`
-# tried to wrap a plain Vector (returned by `similar(::ArrayPartition, T, dims)`
-# when `dims != size(A)`) in NamedArrayPartition's inner constructor, which
-# requires an ArrayPartition.
+# indexing a NamedArrayPartition with a UnitRange / Vector{Int} smaller than the
+# whole array used to throw a MethodError because the AbstractArray indexing
+# path called `similar(::NAP, T, dims)`, which tried to wrap a plain Vector
+# (returned by `similar(::ArrayPartition, T, dims)` for `dims != size(A)`) in
+# NamedArrayPartition's inner constructor, which requires an ArrayPartition.
+#
+# The `_unsafe_getindex(::IndexStyle, ::NAP, I...)` shortcut bypasses `similar`
+# entirely, allocating a plain Vector destination directly. Slicing therefore
+# always returns a Vector and is type-stable.
 @testset "NamedArrayPartition issue #583 indexing" begin
     x = NamedArrayPartition(a = ones(2), b = 2 * ones(3))
 
-    # UnitRange that doesn't span the whole array: returns a plain Vector
+    # UnitRange / Vector{Int} indexing all return Vector and are type-stable
     @test x[1:2] == [1.0, 1.0]
-    @test x[1:2] isa Vector{Float64}
     @test x[2:4] == [1.0, 2.0, 2.0]
-
-    # Vector{Int} indexing
+    @test x[1:end] == [1.0, 1.0, 2.0, 2.0, 2.0]
     @test x[[1, 2]] == [1.0, 1.0]
     @test x[[1, 4]] == [1.0, 2.0]
+
+    @test x[1:2]    isa Vector{Float64}
+    @test x[1:end]  isa Vector{Float64}
     @test x[[1, 4]] isa Vector{Float64}
 
-    # Existing behavior: full-range slice still preserves NamedArrayPartition
-    @test typeof(x[1:end]) <: NamedArrayPartition
+    # Inferred return types: Vector, not Union
+    @test (@inferred x[1:2])           isa Vector{Float64}
+    @test (@inferred x[1:length(x)])   isa Vector{Float64}
+    @test (@inferred x[[1, 4]])        isa Vector{Float64}
 
-    # `similar` with a non-matching dims tuple gives the backing-array fallback
+    # `similar` with a non-matching dims falls back to the backing array;
+    # with matching dims keeps the NamedArrayPartition wrapper.
     @test similar(x, Float64, (2,)) isa Vector{Float64}
-    @test similar(x, (2,)) isa Vector{Float64}
-    # `similar` with matching dims preserves the NamedArrayPartition wrapper
+    @test similar(x, (2,))          isa Vector{Float64}
     @test similar(x, Float64, size(x)) isa NamedArrayPartition
-    @test similar(x, size(x)) isa NamedArrayPartition
+    @test similar(x, size(x))          isa NamedArrayPartition
 
-    # Scalar indexing untouched
+    # Scalar indexing untouched and type-stable
     @test x[1] == 1.0
     @test x[3] == 2.0
+    @test (@inferred x[1]) === 1.0
     x[1] = 99.0
     @test x[1] == 99.0
 end