From 9f917603036455e3f133c3d94e42960c24b768e1 Mon Sep 17 00:00:00 2001
From: jishnub <jishnub@users.noreply.github.com>
Date: Wed, 14 Jul 2021 23:16:00 +0400
Subject: [PATCH] check compatible sizes

---
 src/OffsetArrays.jl | 27 ++++++++++++++++++---------
 src/utils.jl        |  6 ++++++
 test/runtests.jl    |  8 ++++++++
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl
index 5324da03..c2ca7545 100644
--- a/src/OffsetArrays.jl
+++ b/src/OffsetArrays.jl
@@ -340,30 +340,39 @@ _similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength
 # reshape accepts a single colon
 Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
 function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
-    AR = reshape(A, map(_indexlength, inds))
+    AR = reshape(no_offset_view(A), map(_indexlength, inds))
     O = OffsetArray(AR, map(_offset, axes(AR), inds))
     return _popreshape(O, axes(AR), _filterreshapeinds(inds))
 end
 
 # Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
 # an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
-# try to pass on the indices as received to the parent.
-# If the parent doesn't define an appropriate method, this will fall back to using the lengths
-# Short-circuit the case with matching indices to circumvent the Base restriction to 1-based indices
-_reshape(A::AbstractVector, ::Tuple{OffsetAxis}) = A
-_reshape(A, inds) = reshape(A, inds)
+# Short-circuit for AbstractVectors if the axes are compatible to get around the Base restriction
+# to 1-based vectors
+function _reshape(A::AbstractVector, inds::Tuple{OffsetAxis})
+    @noinline throw_dimerr(ind::Integer) = throw(
+        DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with length $ind"))
+    @noinline throw_dimerr(ind) = throw(
+        DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with indices $ind"))
+    _checksize(first(inds), size(A,1)) || throw_dimerr(first(inds))
+    A
+end
+_reshape(A, inds) = _reshape2(A, inds)
+_reshape2(A, inds) = reshape(A, inds)
+# avoid a stackoverflow by relegating to the parent if no_offset_view returns an offsetarray
+_reshape2(A::OffsetArray, inds) = reshape(parent(A), inds)
 _reshape_nov(A, inds) = _reshape(no_offset_view(A), inds)
 
 Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
     OffsetArray(_reshape(parent(A), inds), map(_toaxis, inds))
 # And for non-offset axes, we can just return a reshape of the parent directly
-Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(parent(A), inds)
-Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(parent(A), inds)
+Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(A, inds)
+Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(A, inds)
 Base.reshape(A::OffsetVector, ::Colon) = A
 Base.reshape(A::OffsetVector, ::Tuple{Colon}) = A
 Base.reshape(A::OffsetArray, ::Colon) = reshape(A, (Colon(),))
 Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(A, inds)
-Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(parent(A), inds)
+Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(A, inds)
 
 # permutedims in Base does not preserve axes, and can not be fixed in a non-breaking way
 # This is a stopgap solution
diff --git a/src/utils.jl b/src/utils.jl
index f92657cc..a803b73a 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -6,6 +6,12 @@ _indexlength(r::AbstractRange) = length(r)
 _indexlength(i::Integer) = Int(i)
 _indexlength(i::Colon) = Colon()
 
+# utility methods used in reshape
+# we don't use _indexlength in this to avoid converting the arguments to Int
+_checksize(ind::Integer, s) = ind == s
+_checksize(ind::AbstractUnitRange, s) = length(ind) == s
+_checksize(::Colon, s) = true
+
 _toaxis(i::Integer) = Base.OneTo(i)
 _toaxis(i) = i
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 6bde9c27..91d7228e 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1794,6 +1794,14 @@ end
     r2 = reshape(r, 3)
     @test axes(r2, 1) == 1:3
     @test r2 == no_offset_view(r)
+    @test_throws Exception reshape(r, length(r) + 1)
+    @test_throws Exception reshape(r, 1:length(r) + 1)
+    rp = parent(r)
+    @test axes(reshape(rp, 4:6), 1) == 4:6
+    @test axes(reshape(r, (3,1))) == (1:3, 1:1)
+    # the following is broken if rp doesn't define its own reshape
+    # we may fix it here but that's perhaps too much type-piracy
+    @test_broken reshape(rp, length(rp))
 end
 
 @testset "permutedims" begin