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

The invalidation-eliminating change in 05faa73 narrowed
`getindex(::NamedArrayPartition, ...)` from a catch-all `args...` form
to just `Int`. Anything else now falls through `AbstractArray`'s
generic indexing, which routes via `similar(x, T, dims)`.

That exposed a latent bug in the `similar` overloads: when
`dims != size(A)`, `similar(::ArrayPartition, T, dims)` intentionally
degrades to a plain `Vector` (a partition layout no longer fits the
requested shape). The `NamedArrayPartition` overloads then tried to wrap
that `Vector` via the inner constructor
`NamedArrayPartition{T, A<:ArrayPartition{T}, NT}(::A, ::NT)`, which
raised a MethodError. The catch-all `args...` getindex used to bypass
`similar` entirely, so the bug stayed hidden until 05faa73.

Make `similar(::NAP, dims)` and `similar(::NAP, T, dims)` mirror
`ArrayPartition`'s own degradation: wrap when the inner result is an
`ArrayPartition`, otherwise return the fallback array as-is. This
restores `x[1:2]` and `x[[1, 4]]` for partial slices, keeps `x[1:end]`
returning a `NamedArrayPartition` (the full-range slice still hits the
ArrayPartition path), and reintroduces no method that catches a wider
signature than necessary — `using RecursiveArrayTools` still produces
0 invalidation trees, verified with `SnoopCompile.invalidation_trees`.

Closes #583.

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

@@ -34,11 +34,16 @@ function Base.similar(A::NamedArrayPartition)
     )
 end
 
-# return ArrayPartition when possible, otherwise next best thing of the correct size
+# return NamedArrayPartition when the requested dims still match the partition layout;
+# otherwise fall back to the plain backing array of the correct size. ArrayPartition's
+# own `similar(A, dims)` already does this degradation (it returns a Vector when
+# `dims != size(A)`), and we simply propagate that result instead of trying to
+# wrap a non-ArrayPartition in a NamedArrayPartition (which would hit the inner
+# constructor signature `NamedArrayPartition(::A<:ArrayPartition, ::NamedTuple)`).
 function Base.similar(A::NamedArrayPartition, dims::NTuple{N, Int}) where {N}
-    return NamedArrayPartition(
-        similar(getfield(A, :array_partition), dims), getfield(A, :names_to_indices)
-    )
+    inner = similar(getfield(A, :array_partition), dims)
+    inner isa ArrayPartition || return inner
+    return NamedArrayPartition(inner, getfield(A, :names_to_indices))
 end
 
 # similar array partition of common type
@@ -48,11 +53,10 @@ end
     )
 end
 
-# return ArrayPartition when possible, otherwise next best thing of the correct size
 function Base.similar(A::NamedArrayPartition, ::Type{T}, dims::NTuple{N, Int}) where {T, N}
-    return NamedArrayPartition(
-        similar(getfield(A, :array_partition), T, dims), getfield(A, :names_to_indices)
-    )
+    inner = similar(getfield(A, :array_partition), T, dims)
+    inner isa ArrayPartition || return inner
+    return NamedArrayPartition(inner, getfield(A, :names_to_indices))
 end
 
 # similar array partition with different types

test/named_array_partition_tests.jl

@@ -37,3 +37,39 @@ using RecursiveArrayTools, ArrayInterface, Test
     @test typeof((x -> x[1]).(x)) <: NamedArrayPartition
     @test typeof(map(x -> x[1], x)) <: NamedArrayPartition
 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.
+@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
+    @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, 2]] == [1.0, 1.0]
+    @test x[[1, 4]] == [1.0, 2.0]
+    @test x[[1, 4]] isa Vector{Float64}
+
+    # Existing behavior: full-range slice still preserves NamedArrayPartition
+    @test typeof(x[1:end]) <: NamedArrayPartition
+
+    # `similar` with a non-matching dims tuple gives the backing-array fallback
+    @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, Float64, size(x)) isa NamedArrayPartition
+    @test similar(x, size(x)) isa NamedArrayPartition
+
+    # Scalar indexing untouched
+    @test x[1] == 1.0
+    @test x[3] == 2.0
+    x[1] = 99.0
+    @test x[1] == 99.0
+end