Fix OffsetArray-correct indexing in ArrayPartition

Project.toml

@@ -61,6 +61,7 @@ LinearAlgebra = "1.10"
 Measurements = "2.11"
 MonteCarloMeasurements = "1.2"
 Mooncake = "0.5"
+OffsetArrays = "1.17"
 Pkg = "1"
 Polyester = "0.7.16"
 PrecompileTools = "1.2.1"
@@ -90,6 +91,7 @@ FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -105,4 +107,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["FastBroadcast", "KernelAbstractions", "Measurements", "Pkg", "Polyester", "Random", "RecursiveArrayToolsShorthandConstructors", "SafeTestsets", "SciMLTesting", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful"]
+test = ["FastBroadcast", "KernelAbstractions", "Measurements", "OffsetArrays", "Pkg", "Polyester", "Random", "RecursiveArrayToolsShorthandConstructors", "SafeTestsets", "SciMLTesting", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful"]

src/array_partition.jl

@@ -238,12 +238,15 @@ Base.filter(f, A::ArrayPartition) = ArrayPartition(map(x -> filter(f, x), A.x))
 for type in [AbstractArray, PermutedDimsArray]
     @eval function Base.copyto!(dest::$(type), A::ArrayPartition)
         @assert length(dest) == length(A)
-        cur = 1
+        cur = firstindex(dest)
         @inbounds for i in 1:length(A.x)
             if A.x[i] isa Number
                 dest[cur:(cur + length(A.x[i]) - 1)] .= A.x[i]
             else
-                dest[cur:(cur + length(A.x[i]) - 1)] .= vec(A.x[i])
+                # copyto!(dest, doffs, src) copies src by linear index regardless of
+                # src's own indexing, so a non-1-based partition (e.g. OffsetArray)
+                # still lands at the right destination offset.
+                copyto!(dest, cur, A.x[i])
             end
             cur += length(A.x[i])
         end
@@ -308,9 +311,10 @@ end
 Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Int)
     @boundscheck checkbounds(A, i)
     @inbounds for j in 1:length(A.x)
-        i -= length(A.x[j])
+        x = A.x[j]
+        i -= length(x)
         if i <= 0
-            return A.x[j][length(A.x[j]) + i]
+            return x[firstindex(x) + length(x) + i - 1]
         end
     end
     throw(BoundsError(A, i))
@@ -337,9 +341,10 @@ Base.getindex(A::ArrayPartition{T, S}, ::Colon) where {T, S} = T[a for a in Chai
 Base.@propagate_inbounds function Base.setindex!(A::ArrayPartition, v, i::Int)
     @boundscheck checkbounds(A, i)
     return @inbounds for j in 1:length(A.x)
-        i -= length(A.x[j])
+        x = A.x[j]
+        i -= length(x)
         if i <= 0
-            A.x[j][length(A.x[j]) + i] = v
+            x[firstindex(x) + length(x) + i - 1] = v
             break
         end
     end

test/Core/partitions_test.jl

@@ -1,4 +1,4 @@
-using RecursiveArrayTools, Test, Statistics, ArrayInterface, Adapt
+using RecursiveArrayTools, Test, Statistics, ArrayInterface, Adapt, OffsetArrays
 using RecursiveArrayToolsShorthandConstructors
 
 @test length(ArrayPartition()) == 0
@@ -416,3 +416,26 @@ end
 struct TestIsnanFunctor end
 (::TestIsnanFunctor)(x) = isnan(x)
 @test all(TestIsnanFunctor(), AP[[NaN], [NaN]])
+
+@testset "OffsetArray partitions index correctly" begin
+    # getindex/setindex! and copyto!(::AbstractArray, ::ArrayPartition) assumed every
+    # partition (and the destination) were 1-based, silently reading/writing the wrong
+    # element for a partition with non-1-based indices.
+    o = OffsetArray([10.0, 20.0, 30.0], 0:2)
+    A = ArrayPartition([1.0, 2.0], o)
+    @test A[3] == 10.0
+    @test A[4] == 20.0
+    @test A[5] == 30.0
+    @test collect(A) == [1.0, 2.0, 10.0, 20.0, 30.0]
+
+    o2 = OffsetArray(zeros(3), 0:2)
+    B = ArrayPartition([1.0, 2.0], o2)
+    B[5] = 99.0
+    @test parent(o2) == [0.0, 0.0, 99.0]
+    @test B[5] == 99.0
+
+    C = ArrayPartition([1.0, 2.0], [3.0, 4.0])
+    dest = OffsetArray(zeros(4), 0:3)
+    copyto!(dest, C)
+    @test parent(dest) == [1.0, 2.0, 3.0, 4.0]
+end