From 49c463b7c28e1a2494f6d136a32be22a0e6facd9 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 3 Jul 2023 23:39:50 +0530 Subject: [PATCH 1/3] Add isonebased trait to opt-out of OffsetArray wrappers --- src/OffsetArrays.jl | 46 ++++++++++++++++++++++++++++++++------------- 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl index 2a03bf5..53d851d 100644 --- a/src/OffsetArrays.jl +++ b/src/OffsetArrays.jl @@ -317,15 +317,35 @@ end Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T = similar(parent(A), T, dims) + +""" + isonebased(ax::AbstractUnitRange{<:Integer}) + +Return whether `firstindex(ax)` is statically known to be `1`. Custom axis types may extend this method +to ensure that the axis offset is ignored, for example in `similar`. +""" +isonebased(_) = false +isonebased(::Integer) = true +isonebased(::Base.OneTo) = true +isonebased(::IIUR{<:Base.OneTo}) = true + +to_length(i::Integer) = i +to_length(i::AbstractUnitRange) = length(i) + +# Since the following is committing type-piracy, we provide an opt-out mechanism to the users function Base.similar(A::AbstractArray, ::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T - # strip IdOffsetRanges to extract the parent range and use it to generate the array - new_shape = map(_strip_IdOffsetRange, shape) - # route through _similar_axes_or_length to avoid a stack overflow if map(_strip_IdOffsetRange, shape) === shape - # This tries to use new_shape directly in similar if similar(A, T, ::typeof(new_shape)) is defined - # If this fails, it calls similar(A, T, map(_indexlength, new_shape)) to use the size along each axis - # to generate the new array - P = _similar_axes_or_length(A, T, new_shape, shape) - return OffsetArray(P, map(_offset, axes(P), shape)) + if all(isonebased, shape) + return similar(A, T, map(to_length, shape)) + else + # strip IdOffsetRanges to extract the parent range and use it to generate the array + new_shape = map(_strip_IdOffsetRange, shape) + # route through _similar_axes_or_length to avoid a stack overflow if map(_strip_IdOffsetRange, shape) === shape + # This tries to use new_shape directly in similar if similar(A, T, ::typeof(new_shape)) is defined + # If this fails, it calls similar(A, T, map(_indexlength, new_shape)) to use the size along each axis + # to generate the new array + P = _similar_axes_or_length(A, T, new_shape, shape) + return OffsetArray(P, map(_offset, axes(P), shape)) + end end Base.similar(::Type{A}, sz::Tuple{Vararg{Int}}) where {A<:OffsetArray} = similar(Array{eltype(A)}, sz) function Base.similar(::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where {T<:AbstractArray} @@ -699,11 +719,11 @@ no_offset_view(a::Array) = a no_offset_view(i::Number) = i no_offset_view(A::AbstractArray) = _no_offset_view(axes(A), A) _no_offset_view(::Tuple{}, A::AbstractArray{T,0}) where T = A -_no_offset_view(::Tuple{Base.OneTo, Vararg{Base.OneTo}}, A::AbstractArray) = A -# the following method is needed for ambiguity resolution -_no_offset_view(::Tuple{Base.OneTo, Vararg{Base.OneTo}}, A::AbstractUnitRange) = A -_no_offset_view(::Any, A::AbstractArray) = OffsetArray(A, Origin(1)) -_no_offset_view(::Any, A::AbstractUnitRange) = UnitRange(A) +function _no_offset_view(ax::Tuple, A::AbstractArray) + all(isonebased, ax) ? A : __no_offset_view(A) +end +__no_offset_view(A::AbstractArray) = OffsetArray(A, Origin(1)) +__no_offset_view(A::AbstractUnitRange) = UnitRange(A) ##### # center/centered From c38d76e909f6e84c5ff3f2b41bc919d54aa551ca Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Jul 2023 00:08:24 +0530 Subject: [PATCH 2/3] use isonebased in similar(::Type{<:AbstractArray}, ax...) --- src/OffsetArrays.jl | 27 +++++++++++---------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl index 53d851d..b4dfde2 100644 --- a/src/OffsetArrays.jl +++ b/src/OffsetArrays.jl @@ -327,30 +327,25 @@ to ensure that the axis offset is ignored, for example in `similar`. isonebased(_) = false isonebased(::Integer) = true isonebased(::Base.OneTo) = true -isonebased(::IIUR{<:Base.OneTo}) = true - -to_length(i::Integer) = i -to_length(i::AbstractUnitRange) = length(i) +isonebased(r::IIUR) = isonebased(r.indices) # Since the following is committing type-piracy, we provide an opt-out mechanism to the users function Base.similar(A::AbstractArray, ::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T - if all(isonebased, shape) - return similar(A, T, map(to_length, shape)) - else - # strip IdOffsetRanges to extract the parent range and use it to generate the array - new_shape = map(_strip_IdOffsetRange, shape) - # route through _similar_axes_or_length to avoid a stack overflow if map(_strip_IdOffsetRange, shape) === shape - # This tries to use new_shape directly in similar if similar(A, T, ::typeof(new_shape)) is defined - # If this fails, it calls similar(A, T, map(_indexlength, new_shape)) to use the size along each axis - # to generate the new array - P = _similar_axes_or_length(A, T, new_shape, shape) - return OffsetArray(P, map(_offset, axes(P), shape)) - end + # strip IdOffsetRanges to extract the parent range and use it to generate the array + new_shape = map(_strip_IdOffsetRange, shape) + # route through _similar_axes_or_length to avoid a stack overflow if map(_strip_IdOffsetRange, shape) === shape + # This tries to use new_shape directly in similar if similar(A, T, ::typeof(new_shape)) is defined + # If this fails, it calls similar(A, T, map(_indexlength, new_shape)) to use the size along each axis + # to generate the new array + P = _similar_axes_or_length(A, T, new_shape, shape) + all(isonebased, shape) && return P + return OffsetArray(P, map(_offset, axes(P), shape)) end Base.similar(::Type{A}, sz::Tuple{Vararg{Int}}) where {A<:OffsetArray} = similar(Array{eltype(A)}, sz) function Base.similar(::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where {T<:AbstractArray} new_shape = map(_strip_IdOffsetRange, shape) P = _similar_axes_or_length(T, new_shape, shape) + all(isonebased, shape) && return P OffsetArray(P, map(_offset, axes(P), shape)) end # Try to use the axes to generate the parent array type From 71565b6bbb2375ad75f7f5a95967071af1322fe0 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Jul 2023 11:07:01 +0530 Subject: [PATCH 3/3] dispatch on types, and add tests --- src/OffsetArrays.jl | 22 +++++++++++++++------- test/runtests.jl | 10 +++++++++- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl index b4dfde2..0851588 100644 --- a/src/OffsetArrays.jl +++ b/src/OffsetArrays.jl @@ -319,15 +319,23 @@ Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T = similar(parent(A), T, dims) """ - isonebased(ax::AbstractUnitRange{<:Integer}) + isonebased(::Type{T}) where {T<:AbstractUnitRange{<:Integer}} -Return whether `firstindex(ax)` is statically known to be `1`. Custom axis types may extend this method -to ensure that the axis offset is ignored, for example in `similar`. +Return whether `first(ax::T)` is statically known to be `1` for any `ax` of type `T` +(`false` by default). +Custom axis types may extend this method to ensure that the axis offset is ignored, +for example in `similar`. This may strip `OffsetArray` wrappers on occasion, or dispatch +to `Base` methods for 1-based axes. + +For axes that are essentially wrappers around another `AbstractUnitRange`, +and share their indexing with their parents, one may forward the type of the +parent range to `isonebased`. """ -isonebased(_) = false -isonebased(::Integer) = true -isonebased(::Base.OneTo) = true -isonebased(r::IIUR) = isonebased(r.indices) +isonebased(x) = isonebased(typeof(x)) +isonebased(::Type) = false +isonebased(::Type{<:Integer}) = true +isonebased(::Type{<:Base.OneTo}) = true +isonebased(::Type{IdentityUnitRange{T}}) where {T} = isonebased(T) # Since the following is committing type-piracy, we provide an opt-out mechanism to the users function Base.similar(A::AbstractArray, ::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T diff --git a/test/runtests.jl b/test/runtests.jl index 300a18d..4b54cd3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,7 +9,7 @@ using EllipsisNotation using FillArrays using LinearAlgebra using OffsetArrays -using OffsetArrays: IdentityUnitRange, no_offset_view, IIUR, Origin, IdOffsetRange +using OffsetArrays: IdentityUnitRange, no_offset_view, IIUR, Origin, IdOffsetRange, isonebased using StaticArrays using Test @@ -2262,6 +2262,14 @@ Base.size(x::PointlessWrapper) = size(parent(x)) Base.axes(x::PointlessWrapper) = axes(parent(x)) Base.getindex(x::PointlessWrapper, i...) = x.parent[i...] +@testset "isonebased" begin + @test isonebased(Base.OneTo(2)) + @test !isonebased(1:2) + @test isonebased(IdentityUnitRange(Base.OneTo(2))) + @test !isonebased(IdentityUnitRange(1:2)) + @test similar(zeros(1), Int, IdentityUnitRange(Base.OneTo(2))) isa Vector{Int} +end + @testset "no offset view" begin # OffsetArray fallback A = randn(3, 3)