Add AbstractOneTo and have OneTo be its subtype (#56902)

Currently, `Base` defines `similar` for `Base.OneTo`, with the
understanding that offset axes will be handled elsewhere. However,
`Base.OneTo` is just one possible one-based range, and there are others
such as `StaticArrays.SOneTo` that exist in the ecosystem. `Base`
doesn't provide a method to handle a combination of different one-based
ranges in `similar`, which leaves the packages in an awkward position:
they need to define methods like
```julia
similar(A::AbstractArray, ::Type{T}, shape::HeterogeneousShapeTuple) where {T} = similar(A, T, homogenize_shape(shape))
```
where `HeterogeneousShapeTuple` is defined as
```julia
Tuple{Vararg{Union{Integer, Base.OneTo, SOneTo}}}
```

https://github.com/JuliaArrays/StaticArrays.jl/blob/07c12450d1b3481dda4b503564ae4a5cb4e27ce4/src/abstractarray.jl#L141-L146
Unfortunately, such methods are borderline type-piracy, as noted in
JuliaArrays/StaticArrays.jl#1248. In
particular, if the narrower `Base` method that handles `Union{Integer,
OneTo}` is removed, then this method explicitly becomes pirating.

A solution to this situation is to have `Base` handle all one-based
ranges, such that arbitrary combinations of one-based ranges hit
fallback methods in `Base`. This PR is a first step in this direction.
We add the abstract type `AbstractOneTo`, and have `OneTo` be its
subtype. We also add methods to `similar` and `reshape` that accept
`AbstractOneTo` arguments. This makes it unnecessary for packages to
dispatch on awkward combinations of `Union{Integer, OneTo}` and custom
one-based axes, as the base implementation would handle such cases
already.

There may be other methods that accept an `AbstractOneTo` instead of a
`OneTo`, but these may be addressed in separate PRs. Also, there may be
one-based ranges that can't subtype `AbstractOneTo`, and a full solution
that accepts such ranges as well needs to be implemented through a
trait. This may also be handled in a separate PR.

---------

Co-authored-by: Tim Holy <[email protected]>

Author: jishnub

NEWS.md

@@ -25,7 +25,8 @@ New library functions
 New library features
 --------------------
 
-`sort(keys(::Dict))` and `sort(values(::Dict))` now automatically collect, they previously threw ([#56978]).
+* `sort(keys(::Dict))` and `sort(values(::Dict))` now automatically collect, they previously threw ([#56978]).
+* `Base.AbstractOneTo` is added as a supertype of one-based axes, with `Base.OneTo` as its subtype ([#56902]).
 
 Standard library changes
 ------------------------

base/abstractarray.jl

@@ -818,15 +818,18 @@ julia> similar(falses(10), Float64, 2, 4)
 See also: [`undef`](@ref), [`isassigned`](@ref).
 """
 similar(a::AbstractArray{T}) where {T}                             = similar(a, T)
-similar(a::AbstractArray, ::Type{T}) where {T}                     = similar(a, T, to_shape(axes(a)))
-similar(a::AbstractArray{T}, dims::Tuple) where {T}                = similar(a, T, to_shape(dims))
-similar(a::AbstractArray{T}, dims::DimOrInd...) where {T}          = similar(a, T, to_shape(dims))
-similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T}  = similar(a, T, to_shape(dims))
+similar(a::AbstractArray, ::Type{T}) where {T}                     = similar(a, T, axes(a))
+similar(a::AbstractArray{T}, dims::Tuple) where {T}                = similar(a, T, dims)
+similar(a::AbstractArray{T}, dims::DimOrInd...) where {T}          = similar(a, T, dims)
+similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T}  = similar(a, T, dims)
 # Similar supports specifying dims as either Integers or AbstractUnitRanges or any mixed combination
 # thereof. Ideally, we'd just convert Integers to OneTos and then call a canonical method with the axes,
 # but we don't want to require all AbstractArray subtypes to dispatch on Base.OneTo. So instead we
 # define this method to convert supported axes to Ints, with the expectation that an offset array
 # package will define a method with dims::Tuple{Union{Integer, UnitRange}, Vararg{Union{Integer, UnitRange}}}
+similar(a::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, AbstractOneTo}, Vararg{Union{Integer, AbstractOneTo}}}) where {T} = similar(a, T, to_shape(dims))
+# legacy method for packages that specialize similar(A::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, OneTo, CustomAxis}, Vararg{Union{Integer, OneTo, CustomAxis}}}
+# leaving this method in ensures that Base owns the more specific method
 similar(a::AbstractArray, ::Type{T}, dims::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T} = similar(a, T, to_shape(dims))
 # similar creates an Array by default
 similar(a::AbstractArray, ::Type{T}, dims::Dims{N}) where {T,N}    = Array{T,N}(undef, dims)
@@ -837,7 +840,7 @@ to_shape(dims::DimsOrInds) = map(to_shape, dims)::DimsOrInds
 # each dimension
 to_shape(i::Int) = i
 to_shape(i::Integer) = Int(i)
-to_shape(r::OneTo) = Int(last(r))
+to_shape(r::AbstractOneTo) = Int(last(r))
 to_shape(r::AbstractUnitRange) = r
 
 """
@@ -863,6 +866,8 @@ would create a 1-dimensional logical array whose indices match those
 of the columns of `A`.
 """
 similar(::Type{T}, dims::DimOrInd...) where {T<:AbstractArray} = similar(T, dims)
+similar(::Type{T}, shape::Tuple{Union{Integer, AbstractOneTo}, Vararg{Union{Integer, AbstractOneTo}}}) where {T<:AbstractArray} = similar(T, to_shape(shape))
+# legacy method for packages that specialize similar(::Type{T}, dims::Tuple{Union{Integer, OneTo, CustomAxis}, Vararg{Union{Integer, OneTo, CustomAxis}})
 similar(::Type{T}, shape::Tuple{Union{Integer, OneTo}, Vararg{Union{Integer, OneTo}}}) where {T<:AbstractArray} = similar(T, to_shape(shape))
 similar(::Type{T}, dims::Dims) where {T<:AbstractArray} = T(undef, dims)
 

base/range.jl

@@ -451,14 +451,21 @@ if isdefined(Main, :Base)
     end
 end
 
+"""
+    Base.AbstractOneTo
+
+Abstract type for ranges that start at 1 and have a step size of 1.
+"""
+abstract type AbstractOneTo{T} <: AbstractUnitRange{T} end
+
 """
     Base.OneTo(n)
 
 Define an `AbstractUnitRange` that behaves like `1:n`, with the added
 distinction that the lower limit is guaranteed (by the type system) to
 be 1.
 """
-struct OneTo{T<:Integer} <: AbstractUnitRange{T}
+struct OneTo{T<:Integer} <: AbstractOneTo{T}
     stop::T # invariant: stop >= zero(stop)
     function OneTo{T}(stop) where {T<:Integer}
         throwbool(r)  = (@noinline; throw(ArgumentError("invalid index: $r of type Bool")))

base/reshapedarray.jl

@@ -120,6 +120,9 @@ julia> reshape(1:6, 2, 3)
 reshape
 
 reshape(parent::AbstractArray, dims::IntOrInd...) = reshape(parent, dims)
+reshape(parent::AbstractArray, shp::Tuple{Union{Integer,AbstractOneTo}, Vararg{Union{Integer,AbstractOneTo}}}) = reshape(parent, to_shape(shp))
+# legacy method for packages that specialize reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo,CustomAxis}, Vararg{Union{Integer,OneTo,CustomAxis}}})
+# leaving this method in ensures that Base owns the more specific method
 reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo}, Vararg{Union{Integer,OneTo}}}) = reshape(parent, to_shape(shp))
 reshape(parent::AbstractArray, dims::Tuple{Integer, Vararg{Integer}}) = reshape(parent, map(Int, dims))
 reshape(parent::AbstractArray, dims::Dims)        = _reshape(parent, dims)

test/abstractarray.jl

@@ -2227,3 +2227,24 @@ end
         @test_throws MethodError isreal(G)
     end
 end
+
+@testset "similar/reshape for AbstractOneTo" begin
+    A = [1,2]
+    @testset "reshape" begin
+        @test reshape(A, 2, SizedArrays.SOneTo(1)) == reshape(A, 2, 1)
+        @test reshape(A, Base.OneTo(2), SizedArrays.SOneTo(1)) == reshape(A, 2, 1)
+        @test reshape(A, SizedArrays.SOneTo(1), 2) == reshape(A, 1, 2)
+        @test reshape(A, SizedArrays.SOneTo(1), Base.OneTo(2)) == reshape(A, 1, 2)
+    end
+    @testset "similar" begin
+        b = similar(A, SizedArrays.SOneTo(1), big(2))
+        @test b isa Array{Int, 2}
+        @test size(b) == (1, 2)
+        b = similar(A, SizedArrays.SOneTo(1), Base.OneTo(2))
+        @test b isa Array{Int, 2}
+        @test size(b) == (1, 2)
+        b = similar(A, SizedArrays.SOneTo(1), 2, Base.OneTo(2))
+        @test b isa Array{Int, 3}
+        @test size(b) == (1, 2, 2)
+    end
+end

test/testhelpers/SizedArrays.jl

@@ -14,7 +14,7 @@ import LinearAlgebra: mul!
 
 export SizedArray
 
-struct SOneTo{N} <: AbstractUnitRange{Int} end
+struct SOneTo{N} <: Base.AbstractOneTo{Int} end
 SOneTo(N) = SOneTo{N}()
 Base.length(::SOneTo{N}) where {N} = N
 Base.size(r::SOneTo) = (length(r),)
@@ -58,14 +58,6 @@ Base.parent(S::SizedArray) = S.data
 +(S1::SizedArray{SZ}, S2::SizedArray{SZ}) where {SZ} = SizedArray{SZ}(S1.data + S2.data)
 ==(S1::SizedArray{SZ}, S2::SizedArray{SZ}) where {SZ} = S1.data == S2.data
 
-homogenize_shape(t::Tuple) = (_homogenize_shape(first(t)), homogenize_shape(Base.tail(t))...)
-homogenize_shape(::Tuple{}) = ()
-_homogenize_shape(x::Integer) = x
-_homogenize_shape(x::AbstractUnitRange) = length(x)
-const Dims = Union{Integer, Base.OneTo, SOneTo}
-function Base.similar(::Type{A}, shape::Tuple{Dims, Vararg{Dims}}) where {A<:AbstractArray}
-    similar(A, homogenize_shape(shape))
-end
 function Base.similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray}
     R = similar(A, length.(shape))
     SizedArray{length.(shape)}(R)
@@ -74,6 +66,10 @@ function Base.similar(x::SizedArray, ::Type{T}, shape::Tuple{SOneTo, Vararg{SOne
     sz = map(length, shape)
     SizedArray{sz}(similar(parent(x), T, sz))
 end
+function Base.reshape(x::AbstractArray, shape::Tuple{SOneTo, Vararg{SOneTo}})
+    sz = map(length, shape)
+    SizedArray{length.(sz)}(reshape(x, length.(sz)))
+end
 
 const SizedMatrixLike = Union{SizedMatrix, Transpose{<:Any, <:SizedMatrix}, Adjoint{<:Any, <:SizedMatrix}}