to_dims methods

Author: Tokazama

src/OffsetArrays.jl

@@ -122,17 +122,6 @@ struct OffsetArray{T,N,AA<:AbstractArray{T,N}} <: AbstractArray{T,N}
 end
 
 ArrayInterface.parent_type(::Type{O}) where {T,N,A<:AbstractArray{T,N},O<:OffsetArrays.OffsetArray{T,N,A}} = A
-function _offset_axis_type(::Type{T}, dim::StaticInt{D}) where {T,D}
-    OffsetArrays.IdOffsetRange{Int,ArrayInterface.axes_types(T, dim)}
-end
-function ArrayInterface.axes_types(::Type{T}) where {T<:OffsetArrays.OffsetArray}
-    ArrayInterface.Static.eachop_tuple(_offset_axis_type, Static.nstatic(Val(ndims(T))), ArrayInterface.parent_type(T))
-end
-@inline ArrayInterface.axes(A::OffsetArrays.OffsetArray) = Base.axes(A)
-@inline _axes(A::OffsetArrays.OffsetArray, dim::Integer) = Base.axes(A, dim)
-@inline function ArrayInterface.axes(A::OffsetArrays.OffsetArray{T,N}, ::StaticInt{M}) where {T,M,N}
-    _axes(A, StaticInt{M}(), gt(StaticInt{M}(),StaticInt{N}()))
-end
 
 """
     OffsetVector(v, index)
@@ -299,10 +288,20 @@ Base.parent(A::OffsetArray) = A.parent
 # Base.Broadcast.BroadcastStyle(::Type{<:OffsetArray{<:Any, <:Any, AA}}) where AA = Base.Broadcast.BroadcastStyle(AA)
 
 @inline Base.size(A::OffsetArray) = size(parent(A))
+@inline Base.size(A::OffsetArray, dim) = size(parent(A), ArrayInterface.to_dims(A, dim))
 
-@inline Base.axes(A::OffsetArray) = map(IdOffsetRange, axes(parent(A)), A.offsets)
-@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
+@inline Base.axes(A::OffsetArray) = map(IdOffsetRange, ArrayInterface.axes(parent(A)), A.offsets)
+Base.axes(A::OffsetArray, d) = axes(A, ArrayInterface.to_dims(A, d))
+@inline function Base.axes(A::OffsetArray, d::Union{Int,StaticInt})
+    d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
+end
 @inline Base.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1))  # we only need to specialize this one
+function _offset_axis_type(::Type{T}, dim::StaticInt{D}) where {T,D}
+    OffsetArrays.IdOffsetRange{Int,ArrayInterface.axes_types(T, dim)}
+end
+function ArrayInterface.axes_types(::Type{T}) where {T<:OffsetArrays.OffsetArray}
+    Static.eachop_tuple(_offset_axis_type, Static.nstatic(Val(ndims(T))), ArrayInterface.parent_type(T))
+end
 
 # Issue 128
 # See https://github.com/JuliaLang/julia/issues/37274 for the issue reported in Base

src/axes.jl

@@ -164,6 +164,7 @@ end
 offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange =
     convert(I, r)::I, 0
 
+ArrayInterface.parent_type(::Type{<:IdOffsetRange{<:Any,I}}) where {I} = I
 @inline Base.parent(r::IdOffsetRange) = r.parent
 @inline Base.axes(r::IdOffsetRange) = (Base.axes1(r),)
 @inline Base.axes1(r::IdOffsetRange) = IdOffsetRange(Base.axes1(r.parent), r.offset)