fix linear indexing for ReshapedArray

JuliaLang · Jun 17, 2021 · 0e8722c · 0e8722c
1 parent d09a290
commit 0e8722c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/base/reshapedarray.jl b/base/reshapedarray.jl
@@ -172,7 +172,6 @@ function _reshape(v::AbstractVector, dims::Dims{1})
 end
 # General reshape
 function _reshape(parent::AbstractArray, dims::Dims)
-    require_one_based_indexing(parent)
     n = length(parent)
     prod(dims) == n || _throw_dmrs(n, "size", dims)
     __reshape((parent, IndexStyle(parent)), dims)
@@ -227,7 +226,8 @@ offset_if_vec(i::Integer, axs::Tuple) = i
 
 @inline function getindex(A::ReshapedArrayLF, index::Int)
     @boundscheck checkbounds(A, index)
-    @inbounds ret = parent(A)[index]
+    indexparent = index - firstindex(A) + firstindex(parent(A))
+    @inbounds ret = parent(A)[indexparent]
     ret
 end
 @inline function getindex(A::ReshapedArray{T,N}, indices::Vararg{Int,N}) where {T,N}
@@ -251,7 +251,8 @@ end
 
 @inline function setindex!(A::ReshapedArrayLF, val, index::Int)
     @boundscheck checkbounds(A, index)
-    @inbounds parent(A)[index] = val
+    indexparent = index - firstindex(A) + firstindex(parent(A))
+    @inbounds parent(A)[indexparent] = val
     val
 end
 @inline function setindex!(A::ReshapedArray{T,N}, val, indices::Vararg{Int,N}) where {T,N}

diff --git a/test/abstractarray.jl b/test/abstractarray.jl
@@ -1414,5 +1414,29 @@ end
 end
 
 @testset "reshape for offset arrays" begin
-    @test_throws ArgumentError reshape(Base.IdentityUnitRange(0:1), (2,1))
+    r = reshape(Base.IdentityUnitRange(0:1), (2,1))
+    @test r[eachindex(r)] == [0:1;]
+
+    struct ZeroBasedArray{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
+        a :: A
+        function ZeroBasedArray(a::AbstractArray)
+            Base.require_one_based_indexing(a)
+            new{eltype(a), ndims(a), typeof(a)}(a)
+        end
+    end
+    Base.parent(z::ZeroBasedArray) = z.a
+    Base.size(z::ZeroBasedArray) = size(parent(z))
+    Base.axes(z::ZeroBasedArray) = map(x -> Base.IdentityUnitRange(0:x - 1), size(parent(z)))
+    Base.getindex(z::ZeroBasedArray{<:Any, N}, i::Vararg{Int,N}) where {N} = parent(z)[map(x -> x + 1, i)...]
+    Base.setindex!(z::ZeroBasedArray{<:Any, N}, val, i::Vararg{Int,N}) where {N} = parent(z)[map(x -> x + 1, i)...] = val
+
+    z = ZeroBasedArray(collect(1:4))
+    r2 = reshape(z, (4, 1))
+    @test r2[CartesianIndices(r2)] == r2[LinearIndices(r2)]
+    r2[firstindex(r2)] = 34
+    @test z[0] == 34
+    r2[eachindex(r2)] = r2 .* 2
+    for (i, j) in zip(eachindex(r2), eachindex(z))
+        @test r2[i] == z[j]
+    end
 end