Skip to content

Commit

Permalink
to_indices for OffsetUnitRange{<:Integer}
Browse files Browse the repository at this point in the history
  • Loading branch information
@jishnub
jishnub committed Feb 27, 2021
1 parent 57ad360 commit c3ef299
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 30 deletions.
50 changes: 40 additions & 10 deletions src/OffsetArrays.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,28 +333,58 @@ Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest)

### Special handling for AbstractRange

const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}}
const OffsetRange{T} = OffsetVector{T,<:AbstractRange{T}}
const OffsetUnitRange{T} = OffsetVector{T,<:AbstractUnitRange{T}}
const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer

Base.step(a::OffsetRange) = step(parent(a))

@propagate_inbounds function Base.getindex(a::OffsetRange, r::OffsetRange)
_maybewrapoffset(a.parent[r.parent .- a.offsets[1]], axes(r,1))
@inline function Base.getindex(a::OffsetRange, r::OffsetRange)
@boundscheck checkbounds(a, r)
@inbounds pr = a.parent[r.parent .- a.offsets[1]]
_maybewrapoffset(pr, axes(r,1))
end
@propagate_inbounds function Base.getindex(a::OffsetRange, r::IdOffsetRange)
_maybewrapoffset(a.parent[r.parent .+ (r.offset - a.offsets[1])], axes(r,1))
@inline function Base.getindex(a::OffsetRange, r::IdOffsetRange)
@boundscheck checkbounds(a, r)
@inbounds pr = a.parent[r.parent .+ (r.offset - a.offsets[1])]
_maybewrapoffset(pr, axes(r,1))
end
@inline function Base.getindex(a::OffsetRange, r::AbstractRange)
@boundscheck checkbounds(a, r)
@inbounds pr = a.parent[r .- a.offsets[1]]
_maybewrapoffset(pr, axes(r,1))
end
@inline function Base.getindex(a::AbstractRange, r::OffsetRange)
@boundscheck checkbounds(a, r)
@inbounds pr = a[parent(r)]
_maybewrapoffset(pr, axes(r,1))
end

# An OffsetUnitRange{<:Integer} has an equivalent IdOffsetRange with the same values and axes
# We may replace the former with the latter in an indexing operation to obtain a performance boost
function Base.to_indices(A::AbstractArray, ax::Tuple, I::Tuple{OffsetUnitRange{<:Integer}, Vararg{Any}})
or = first(I)
r = parent(or)
of = or.offsets[1]
ior = IdOffsetRange(UnitRange(r) .- of, of)
to_indices(A, ax, (ior, tail(I)...))
end
@propagate_inbounds Base.getindex(a::OffsetRange, r::AbstractRange) = _maybewrapoffset(a.parent[r .- a.offsets[1]], axes(r,1))
@propagate_inbounds Base.getindex(a::AbstractRange, r::OffsetRange) = _maybewrapoffset(a[parent(r)], axes(r,1))

for OR in [:IIUR, :IdOffsetRange]
for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
@eval @propagate_inbounds Base.getindex(r::$R, s::$OR) = _maybewrapoffset(r[UnitRange(s)], axes(s,1))
@eval @inline function Base.getindex(r::$R, s::$OR)
@boundscheck checkbounds(r, s)
@inbounds pr = r[UnitRange(s)]
_maybewrapoffset(pr, axes(s,1))
end
end

# this method is needed for ambiguity resolution
@eval @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T =
_maybewrapoffset(r[UnitRange(s)], axes(s,1))
@eval @inline function Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T
@boundscheck checkbounds(r, s)
@inbounds pr = r[UnitRange(s)]
_maybewrapoffset(pr, axes(s,1))
end
end

# This is technically breaking, so it might be incorporated in the next major release
Expand Down
27 changes: 16 additions & 11 deletions src/axes.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ julia> ro[-1]
-1
julia> ro[3]
ERROR: BoundsError: attempt to access 3-element UnitRange{$Int} at index [5]
ERROR: BoundsError: attempt to access 3-element OffsetArrays.IdOffsetRange{Int64,UnitRange{Int64}} with indices -1:1 at index [3]
```
If the range doesn't start at 1, the values may be different from the indices:
Expand All @@ -35,7 +35,7 @@ julia> ro[-1]
9
julia> ro[3]
ERROR: BoundsError: attempt to access 3-element UnitRange{$Int} at index [5]
ERROR: BoundsError: attempt to access 3-element OffsetArrays.IdOffsetRange{Int64,UnitRange{Int64}} with indices -1:1 at index [3]
```
# Extended help
Expand Down Expand Up @@ -78,10 +78,10 @@ struct IdOffsetRange{T<:Integer,I<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
offset::T

IdOffsetRange{T,I}(r::I, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}} = new{T,I}(r, offset)
#= This method is necessary to avoid a StackOverflowError in IdOffsetRange{T,I}(r::IdOffsetRange, offset::Integer).
The type signature in that method is more specific than IdOffsetRange{T,I}(r::I, offset::T),
so it ends up calling itself if I <: IdOffsetRange.

#= This method is necessary to avoid a StackOverflowError in IdOffsetRange{T,I}(r::IdOffsetRange, offset::Integer).
The type signature in that method is more specific than IdOffsetRange{T,I}(r::I, offset::T),
so it ends up calling itself if I <: IdOffsetRange.
=#
function IdOffsetRange{T,IdOffsetRange{T,I}}(r::IdOffsetRange{T,I}, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}}
new{T,IdOffsetRange{T,I}}(r, offset)
Expand Down Expand Up @@ -155,14 +155,19 @@ end
return (ret[1] + r.offset, ret[2])
end

@propagate_inbounds Base.getindex(r::IdOffsetRange, i::Integer) = r.parent[i - r.offset] + r.offset
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
pr = r.parent[s .- r.offset] .+ r.offset
@inline function Base.getindex(r::IdOffsetRange, i::Integer)
@boundscheck checkbounds(r, i)
@inbounds r.parent[i - r.offset] + r.offset
end
@inline function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
@boundscheck checkbounds(r, s)
@inbounds pr = r.parent[s .- r.offset] .+ r.offset
_maybewrapoffset(pr, axes(s,1))
end
# The following method is required to avoid falling back to getindex(::AbstractUnitRange, ::StepRange{<:Integer})
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
rs = r.parent[s .- r.offset] .+ r.offset
@inline function Base.getindex(r::IdOffsetRange, s::StepRange{<:Integer})
@boundscheck checkbounds(r, s)
@inbounds rs = r.parent[s .- r.offset] .+ r.offset
return no_offset_view(rs)
end

Expand Down
18 changes: 9 additions & 9 deletions src/utils.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,17 @@ _offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
"""
OffsetArrays.AxisConversionStyle(typeof(indices))
`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}`
or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices.
This method is called after `to_indices(A::Array, axes(A), indices)` to provide
`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}`
or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices.
This method is called after `to_indices(A::Array, axes(A), indices)` to provide
further information in case `to_indices` does not return a `Tuple` of `AbstractUnitRange{Int}`.
Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`,
which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should
define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define
`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`.
Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`,
which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should
define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define
`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`.
An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of
An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of
`AbstractUnitRange{Int}` while flattening the indices.
# Example
Expand Down Expand Up @@ -76,6 +76,6 @@ _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ::Base.OneTo) = no_offset_view
_maybewrapoffset(r::AbstractVector, ::Base.OneTo) = no_offset_view(r)
function _maybewrapoffset(r::AbstractUnitRange{<:Integer}, ax)
of = first(ax) - 1
IdOffsetRange(UnitRange(r .- of), of)
IdOffsetRange(UnitRange(r) .- of, of)
end
_maybewrapoffset(r::AbstractVector, ax) = OffsetArray(r, ax)

0 comments on commit c3ef299

Please sign in to comment.