From 34a8f7714ac432b3d99534b59896e4dd4d12f03c Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Sat, 21 Dec 2024 12:08:51 +0530 Subject: [PATCH] Add AbstractOneTo ang have OneTo be its subtype --- base/abstractarray.jl | 7 ++++++- base/range.jl | 9 ++++++++- base/reshapedarray.jl | 3 +++ test/abstractarray.jl | 21 +++++++++++++++++++++ test/testhelpers/SizedArrays.jl | 14 +++++--------- 5 files changed, 43 insertions(+), 11 deletions(-) diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 6d222e7e26281..35b83b59191b9 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -827,6 +827,9 @@ similar(a::AbstractArray, ::Type{T}, dims::DimOrInd...) where {T} = similar(a, # 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) diff --git a/base/range.jl b/base/range.jl index 39428ab741955..8d4c4144c1d05 100644 --- a/base/range.jl +++ b/base/range.jl @@ -451,6 +451,13 @@ if isdefined(Main, :Base) end end +""" + Base.AbstractOneTo + +Abstract type for ranges that start at 1. +""" +abstract type AbstractOneTo{T} <: AbstractUnitRange{T} end + """ Base.OneTo(n) @@ -458,7 +465,7 @@ 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"))) diff --git a/base/reshapedarray.jl b/base/reshapedarray.jl index f65a7d8c9561a..78ed61d66d017 100644 --- a/base/reshapedarray.jl +++ b/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) diff --git a/test/abstractarray.jl b/test/abstractarray.jl index c7ec61704c1bc..f4acb0437ae32 100644 --- a/test/abstractarray.jl +++ b/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 diff --git a/test/testhelpers/SizedArrays.jl b/test/testhelpers/SizedArrays.jl index e52e965a64859..961784b89ab68 100644 --- a/test/testhelpers/SizedArrays.jl +++ b/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}}