diff --git a/Project.toml b/Project.toml index c7313063..cd2556dd 100644 --- a/Project.toml +++ b/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"] diff --git a/src/array_partition.jl b/src/array_partition.jl index a1220451..b717aa33 100644 --- a/src/array_partition.jl +++ b/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 diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index 93a3e7b8..b92c0087 100644 --- a/test/Core/partitions_test.jl +++ b/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