Revision 9b8e5755d39111293d8613be5355668efe19dfe7 authored by Simon Byrne on 02 December 2023, 20:33:33 UTC, committed by Simon Byrne on 02 December 2023, 20:33:33 UTC
1 parent f2df9d3
OffsetArrays.jl
# This file is a part of Julia. License is MIT: https://julialang.org/license
# OffsetArrays (arrays with indexing that doesn't start at 1)
# This test file is designed to exercise support for generic indexing,
# even though offset arrays aren't implemented in Base.
# OffsetArrays v1.11.2
# No compat patch and docstrings
module OffsetArrays
using Base: tail, @propagate_inbounds
using Base: IdentityUnitRange
export OffsetArray, OffsetMatrix, OffsetVector
const IIUR = IdentityUnitRange{<:AbstractUnitRange{<:Integer}}
########################################################################################################
# axes.jl
########################################################################################################
struct IdOffsetRange{T<:Integer,I<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
parent::I
offset::T
function IdOffsetRange{T,I}(r::I, offset::T) where {T<:Integer,I<:AbstractUnitRange{T}}
_bool_check(T, r, offset)
new{T,I}(r, offset)
end
#= 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}}
_bool_check(T, r, offset)
new{T,IdOffsetRange{T,I}}(r, offset)
end
end
function _bool_check(::Type{Bool}, r, offset)
# disallow the construction of IdOffsetRange{Bool, UnitRange{Bool}}(true:true, true)
if offset && (first(r) || last(r))
throw(ArgumentError("values = $r and offset = $offset can not produce a boolean range"))
end
return nothing
end
_bool_check(::Type, r, offset) = nothing
# Construction/coercion from arbitrary AbstractUnitRanges
function IdOffsetRange{T,I}(r::AbstractUnitRange, offset::Integer = 0) where {T<:Integer,I<:AbstractUnitRange{T}}
rc, o = offset_coerce(I, r)
return IdOffsetRange{T,I}(rc, convert(T, o+offset)::T)
end
function IdOffsetRange{T}(r::AbstractUnitRange, offset::Integer = 0) where T<:Integer
rc = convert(AbstractUnitRange{T}, r)::AbstractUnitRange{T}
return IdOffsetRange{T,typeof(rc)}(rc, convert(T, offset)::T)
end
IdOffsetRange(r::AbstractUnitRange{T}, offset::Integer = 0) where T<:Integer =
IdOffsetRange{T,typeof(r)}(r, convert(T, offset)::T)
# Coercion from other IdOffsetRanges
IdOffsetRange{T,I}(r::IdOffsetRange{T,I}) where {T<:Integer,I<:AbstractUnitRange{T}} = r
function IdOffsetRange{T,I}(r::IdOffsetRange, offset::Integer = 0) where {T<:Integer,I<:AbstractUnitRange{T}}
rc, offset_rc = offset_coerce(I, r.parent)
return IdOffsetRange{T,I}(rc, convert(T, r.offset + offset + offset_rc)::T)
end
IdOffsetRange{T}(r::IdOffsetRange{T}) where {T<:Integer} = r
function IdOffsetRange{T}(r::IdOffsetRange, offset::Integer = 0) where T<:Integer
return IdOffsetRange{T}(r.parent, r.offset + offset)
end
IdOffsetRange(r::IdOffsetRange) = r
# Constructor to make `show` round-trippable
function IdOffsetRange(; values::AbstractUnitRange{<:Integer}, indices::AbstractUnitRange{<:Integer})
length(values) == length(indices) || throw(ArgumentError("values and indices must have the same length"))
offset = first(indices) - 1
return IdOffsetRange(values .- offset, offset)
end
# Conversions to an AbstractUnitRange{Int} (and to an OrdinalRange{Int,Int} on Julia v"1.6") are necessary
# to evaluate CartesianIndices for BigInt ranges, as their axes are also BigInt ranges
Base.AbstractUnitRange{T}(r::IdOffsetRange) where {T<:Integer} = IdOffsetRange{T}(r)
# TODO: uncomment these when Julia is ready
# # Conversion preserves both the values and the indices, throwing an InexactError if this
# # is not possible.
# Base.convert(::Type{IdOffsetRange{T,I}}, r::IdOffsetRange{T,I}) where {T<:Integer,I<:AbstractUnitRange{T}} = r
# Base.convert(::Type{IdOffsetRange{T,I}}, r::IdOffsetRange) where {T<:Integer,I<:AbstractUnitRange{T}} =
# IdOffsetRange{T,I}(convert(I, r.parent), r.offset)
# Base.convert(::Type{IdOffsetRange{T,I}}, r::AbstractUnitRange) where {T<:Integer,I<:AbstractUnitRange{T}} =
# IdOffsetRange{T,I}(convert(I, r), 0)
offset_coerce(::Type{Base.OneTo{T}}, r::Base.OneTo) where T<:Integer = convert(Base.OneTo{T}, r), 0
function offset_coerce(::Type{Base.OneTo{T}}, r::AbstractUnitRange) where T<:Integer
o = first(r) - 1
return Base.OneTo{T}(last(r) - o), o
end
# function offset_coerce(::Type{Base.OneTo{T}}, r::IdOffsetRange) where T<:Integer
# rc, o = offset_coerce(Base.OneTo{T}, r.parent)
# Fallback, specialize this method if `convert(I, r)` doesn't do what you need
offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange =
convert(I, r)::I, 0
@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)
@inline Base.unsafe_indices(r::IdOffsetRange) = (Base.axes1(r),)
@inline Base.length(r::IdOffsetRange) = length(r.parent)
@inline Base.isempty(r::IdOffsetRange) = isempty(r.parent)
Base.reduced_index(i::IdOffsetRange) = typeof(i)(first(i):first(i))
# Workaround for #92 on Julia < 1.4
Base.reduced_index(i::IdentityUnitRange{<:IdOffsetRange}) = typeof(i)(first(i):first(i))
for f in [:firstindex, :lastindex]
@eval @inline Base.$f(r::IdOffsetRange) = $f(r.parent) + r.offset
end
for f in [:first, :last]
# coerce the type to deal with values that get promoted on addition (eg. Bool)
@eval @inline Base.$f(r::IdOffsetRange) = eltype(r)($f(r.parent) + r.offset)
end
# Iteration for an IdOffsetRange
@inline Base.iterate(r::IdOffsetRange, i...) = _iterate(r, i...)
# In general we iterate over the parent term by term and add the offset.
# This might have some performance degradation when coupled with bounds-checking
# See https://github.com/JuliaArrays/OffsetArrays.jl/issues/214
@inline function _iterate(r::IdOffsetRange, i...)
ret = iterate(r.parent, i...)
ret === nothing && return nothing
return (eltype(r)(ret[1] + r.offset), ret[2])
end
# Base.OneTo(n) is known to be exactly equivalent to the range 1:n,
# and has no specialized iteration defined for it,
# so we may add the offset to the range directly and iterate over the result
# This gets around the performance issue described in issue #214
# We use the helper function _addoffset to evaluate the range instead of broadcasting
# just in case this makes it easy for the compiler.
@inline _iterate(r::IdOffsetRange{<:Integer, <:Base.OneTo}, i...) = iterate(_addoffset(r.parent, r.offset), i...)
@inline function Base.getindex(r::IdOffsetRange, i::Integer)
i isa Bool && throw(ArgumentError("invalid index: $i of type Bool"))
@boundscheck checkbounds(r, i)
@inbounds eltype(r)(r.parent[oftype(r.offset, i) - r.offset] + r.offset)
end
# Logical indexing following https://github.com/JuliaLang/julia/pull/31829
#= Helper function to perform logical indxeing for boolean ranges
The code implemented is a branch-free version of the following:
range(first(s) ? first(r) : last(r), length=Int(last(s)))
See https://github.com/JuliaArrays/OffsetArrays.jl/pull/224#discussion_r595635143
Logical indexing does not preserve indices, unlike other forms of vector indexing
=#
@inline function _getindex(r, s::AbstractUnitRange{Bool})
range(first(r) * first(s) + last(r) * !first(s), length=Int(last(s)))
end
@inline function _getindex(r, s::StepRange{Bool})
range(first(r) * first(s) + last(r) * !first(s), step = oneunit(step(s)), length=Int(last(s)))
end
@inline function _getindex(r, s::AbstractUnitRange)
@inbounds rs = r.parent[_subtractoffset(s, r.offset)] .+ r.offset
_indexedby(rs, axes(s))
end
@inline function _getindex(r, s::StepRange)
rs = @inbounds r.parent[s .- r.offset] .+ r.offset
_indexedby(rs, axes(s))
end
for T in [:AbstractUnitRange, :StepRange]
@eval @inline function Base.getindex(r::IdOffsetRange, s::$T{<:Integer})
@boundscheck checkbounds(r, s)
return _getindex(r, s)
end
end
# These methods are necessary to avoid ambiguity
for R in [:IIUR, :IdOffsetRange]
@eval @inline function Base.getindex(r::IdOffsetRange, s::$R)
@boundscheck checkbounds(r, s)
return _getindex(r, s)
end
end
# offset-preserve broadcasting
Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(-), r::IdOffsetRange{T}, x::Integer) where T =
IdOffsetRange{T}(r.parent .- x, r.offset)
Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(+), r::IdOffsetRange{T}, x::Integer) where T =
IdOffsetRange{T}(r.parent .+ x, r.offset)
Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(+), x::Integer, r::IdOffsetRange{T}) where T =
IdOffsetRange{T}(x .+ r.parent, r.offset)
Base.show(io::IO, r::IdOffsetRange) = print(io, IdOffsetRange, "(values=",first(r), ':', last(r),", indices=",first(eachindex(r)),':',last(eachindex(r)), ")")
# Optimizations
@inline Base.checkindex(::Type{Bool}, inds::IdOffsetRange, i::Real) = Base.checkindex(Bool, inds.parent, i - inds.offset)
Base._firstslice(r::IdOffsetRange) = IdOffsetRange(Base._firstslice(r.parent), r.offset)
########################################################################################################
# origin.jl
########################################################################################################
struct Origin{T<:Union{Tuple{Vararg{Int}}, Int}}
index::T
end
Origin(I::Tuple{Vararg{Int}}) = Origin{typeof(I)}(I)
Origin(I::Tuple{Vararg{Number}}) = Origin(map(Int, I))
Origin(I::CartesianIndex) = Origin(Tuple(I))
Origin(I::Number...) = Origin(I)
# Origin(0) != Origin((0, )) but they work the same with broadcasting
Origin(n::Number) = Origin{Int}(Int(n))
Base.Broadcast.broadcastable(o::Origin) = Ref(o)
_showidx(index::Integer) = "(" * string(index) * ")"
_showidx(index::Tuple) = string(index)
Base.show(io::IO, o::Origin) = print(io, "Origin", _showidx(o.index))
########################################################################################################
# utils.jl
########################################################################################################
### Low-level utilities ###
_indexoffset(r::AbstractRange) = first(r) - 1
_indexoffset(i::Integer) = 0
_indexlength(r::AbstractRange) = length(r)
_indexlength(i::Integer) = Int(i)
_indexlength(i::Colon) = Colon()
# utility methods used in reshape
# we don't use _indexlength in this to avoid converting the arguments to Int
_checksize(ind::Integer, s) = ind == s
_checksize(ind::AbstractUnitRange, s) = length(ind) == s
_toaxis(i::Integer) = Base.OneTo(i)
_toaxis(i) = i
_strip_IdOffsetRange(r::IdOffsetRange) = parent(r)
_strip_IdOffsetRange(r) = r
_offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
_offset(axparent::AbstractUnitRange, ::Union{Integer, Colon}) = 1 - first(axparent)
_offsets(A::AbstractArray) = map(ax -> first(ax) - 1, axes(A))
_offsets(A::AbstractArray, B::AbstractArray) = map(_offset, axes(B), axes(A))
abstract type AxisConversionStyle end
struct SingleRange <: AxisConversionStyle end
struct TupleOfRanges <: AxisConversionStyle end
AxisConversionStyle(::Type) = SingleRange()
AxisConversionStyle(::Type{<:CartesianIndices}) = TupleOfRanges()
_convertTupleAbstractUnitRange(x) = _convertTupleAbstractUnitRange(AxisConversionStyle(typeof(x)), x)
_convertTupleAbstractUnitRange(::SingleRange, x) = (convert(AbstractUnitRange{Int}, x),)
_convertTupleAbstractUnitRange(::TupleOfRanges, x) = convert(Tuple{Vararg{AbstractUnitRange{Int}}}, x)
_toAbstractUnitRanges(t::Tuple) = (_convertTupleAbstractUnitRange(first(t))..., _toAbstractUnitRanges(tail(t))...)
_toAbstractUnitRanges(::Tuple{}) = ()
# ensure that the indices are consistent in the constructor
_checkindices(A::AbstractArray, indices, label) = _checkindices(ndims(A), indices, label)
function _checkindices(N::Integer, indices, label)
throw_argumenterror(N, indices, label) = throw(ArgumentError(label*" $indices are not compatible with a $(N)D array"))
N == length(indices) || throw_argumenterror(N, indices, label)
end
@inline _indexedby(r::AbstractVector, ax::Tuple{Any}) = _indexedby(r, ax[1])
@inline _indexedby(r::AbstractUnitRange{<:Integer}, ::Base.OneTo) = no_offset_view(r)
@inline _indexedby(r::AbstractUnitRange{Bool}, ::Base.OneTo) = no_offset_view(r)
@inline _indexedby(r::AbstractVector, ::Base.OneTo) = no_offset_view(r)
@inline function _indexedby(r::AbstractUnitRange{<:Integer}, ax::AbstractUnitRange)
of = convert(eltype(r), first(ax) - 1)
IdOffsetRange(_subtractoffset(r, of), of)
end
@inline _indexedby(r::AbstractUnitRange{Bool}, ax::AbstractUnitRange) = OffsetArray(r, ax)
@inline _indexedby(r::AbstractVector, ax::AbstractUnitRange) = OffsetArray(r, ax)
# These functions are equivalent to the broadcasted operation r .- of
# However these ensure that the result is an AbstractRange even if a specific
# broadcasting behavior is not defined for a custom type
@inline _subtractoffset(r::AbstractUnitRange, of) = UnitRange(first(r) - of, last(r) - of)
@inline _subtractoffset(r::AbstractRange, of) = range(first(r) - of, stop = last(r) - of, step = step(r))
# similar to _subtractoffset, except these evaluate r .+ of
@inline _addoffset(r::AbstractUnitRange, of) = UnitRange(first(r) + of, last(r) + of)
@inline _addoffset(r::AbstractRange, of) = range(first(r) + of, stop = last(r) + of, step = step(r))
_contiguousindexingtype(r::AbstractUnitRange{<:Integer}) = r
_of_eltype(::Type{T}, M::AbstractArray{T}) where {T} = M
_of_eltype(T, M::AbstractArray) = map(T, M)
# filter the arguments to reshape to check if there are any ranges
# If not, we may pop the parent array
_filterreshapeinds(t::Tuple{AbstractUnitRange, Vararg{Any}}) = t
_filterreshapeinds(t::Tuple) = _filterreshapeinds(tail(t))
_filterreshapeinds(t::Tuple{}) = t
_popreshape(A::AbstractArray, ax::Tuple{Vararg{Base.OneTo}}, inds::Tuple{}) = no_offset_view(A)
_popreshape(A::AbstractArray, ax, inds) = A
########################################################################################################
# OffsetArrays.jl
########################################################################################################
# Technically we know the length of CartesianIndices but we need to convert it first, so here we
# don't put it in OffsetAxisKnownLength.
const OffsetAxisKnownLength = Union{Integer,AbstractUnitRange}
const OffsetAxis = Union{OffsetAxisKnownLength,Colon}
const ArrayInitializer = Union{UndefInitializer,Missing,Nothing}
## OffsetArray
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
parent::AA
offsets::NTuple{N,Int}
@inline function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}; checkoverflow = true) where {T, N, AA<:AbstractArray{T,N}}
# allocation of `map` on tuple is optimized away
checkoverflow && map(overflow_check, axes(parent), offsets)
new{T, N, AA}(parent, offsets)
end
end
const OffsetVector{T,AA<:AbstractVector{T}} = OffsetArray{T,1,AA}
const OffsetMatrix{T,AA<:AbstractMatrix{T}} = OffsetArray{T,2,AA}
# checks if the offset may be added to the range without overflowing
function overflow_check(r::AbstractUnitRange, offset::Integer)
Base.hastypemax(eltype(r)) || return nothing
# This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
throw_upper_overflow_error(val) = throw(OverflowError("offset should be <= $(typemax(Int) - val) corresponding to the axis $r, received an offset $offset"))
throw_lower_overflow_error(val) = throw(OverflowError("offset should be >= $(typemin(Int) - val) corresponding to the axis $r, received an offset $offset"))
# With ranges in the picture, first(r) might not necessarily be < last(r)
# we therefore use the min and max of first(r) and last(r) to check for overflow
firstlast_min, firstlast_max = minmax(first(r), last(r))
if offset > 0 && firstlast_max > typemax(Int) - offset
throw_upper_overflow_error(firstlast_max)
elseif offset < 0 && firstlast_min < typemin(Int) - offset
throw_lower_overflow_error(firstlast_min)
end
return nothing
end
# Tuples of integers are treated as offsets
# Empty Tuples are handled here
@inline function OffsetArray(A::AbstractArray, offsets::Tuple{Vararg{Integer}}; kw...)
_checkindices(A, offsets, "offsets")
OffsetArray{eltype(A), ndims(A), typeof(A)}(A, offsets; kw...)
end
# These methods are necessary to disallow incompatible dimensions for
# the OffsetVector and the OffsetMatrix constructors
for (FT, ND) in ((:OffsetVector, :1), (:OffsetMatrix, :2))
@eval @inline function $FT(A::AbstractArray{<:Any,$ND}, offsets::Tuple{Vararg{Integer}}; kw...)
_checkindices(A, offsets, "offsets")
OffsetArray{eltype(A), $ND, typeof(A)}(A, offsets; kw...)
end
FTstr = string(FT)
@eval @inline function $FT(A::AbstractArray, offsets::Tuple{Vararg{Integer}}; kw...)
throw(ArgumentError($FTstr*" requires a "*string($ND)*"D array"))
end
end
## OffsetArray constructors
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
# Nested OffsetArrays may strip off the wrapper and collate the offsets
# empty tuples are handled here
@eval @inline function $FT(A::OffsetArray, offsets::Tuple{Vararg{Int}}; checkoverflow = true)
_checkindices(A, offsets, "offsets")
# ensure that the offsets may be added together without an overflow
checkoverflow && map(overflow_check, axes(A), offsets)
I = map(+, _offsets(A, parent(A)), offsets)
$FT(parent(A), I, checkoverflow = false)
end
@eval @inline function $FT(A::OffsetArray, offsets::Tuple{Integer,Vararg{Integer}}; kw...)
$FT(A, map(Int, offsets); kw...)
end
# In general, indices get converted to AbstractUnitRanges.
# CartesianIndices{N} get converted to N ranges
@eval @inline function $FT(A::AbstractArray, inds::Tuple{Any,Vararg{Any}}; kw...)
$FT(A, _toAbstractUnitRanges(to_indices(A, axes(A), inds)); kw...)
end
# convert ranges to offsets
@eval @inline function $FT(A::AbstractArray, inds::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}; kw...)
_checkindices(A, inds, "indices")
# Performance gain by wrapping the error in a function: see https://github.com/JuliaLang/julia/issues/37558
throw_dimerr(lA, lI) = throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
lA = size(A)
lI = map(length, inds)
lA == lI || throw_dimerr(lA, lI)
$FT(A, map(_offset, axes(A), inds); kw...)
end
@eval @inline $FT(A::AbstractArray, inds...; kw...) = $FT(A, inds; kw...)
@eval @inline $FT(A::AbstractArray; checkoverflow = false) = $FT(A, ntuple(zero, Val(ndims(A))), checkoverflow = checkoverflow)
@eval @inline $FT(A::AbstractArray, origin::Origin; checkoverflow = true) = $FT(A, origin.index .- first.(axes(A)); checkoverflow = checkoverflow)
end
(o::Origin)(A::AbstractArray) = OffsetArray(no_offset_view(A), o)
Origin(A::AbstractArray) = Origin(first.(axes(A)))
# conversion-related methods
@inline OffsetArray{T}(M::AbstractArray, I...; kw...) where {T} = OffsetArray{T,ndims(M)}(M, I...; kw...)
@inline function OffsetArray{T,N}(M::AbstractArray{<:Any,N}, I...; kw...) where {T,N}
M2 = _of_eltype(T, M)
OffsetArray{T,N}(M2, I...; kw...)
end
@inline OffsetArray{T,N}(M::OffsetArray{T,N}, I...; kw...) where {T,N} = OffsetArray(M, I...; kw...)
@inline OffsetArray{T,N}(M::AbstractArray{T,N}, I...; kw...) where {T,N} = OffsetArray{T,N,typeof(M)}(M, I...; kw...)
@inline OffsetArray{T,N,A}(M::AbstractArray{<:Any,N}, I...; kw...) where {T,N,A<:AbstractArray{T,N}} = OffsetArray{T,N,A}(M, I; kw...)
@inline function OffsetArray{T,N,A}(M::AbstractArray{<:Any,N}, I::NTuple{N,Int}; checkoverflow = true) where {T,N,A<:AbstractArray{T,N}}
checkoverflow && map(overflow_check, axes(M), I)
Mv = no_offset_view(M)
MvA = convert(A, Mv)::A
Iof = map(+, _offsets(M), I)
OffsetArray{T,N,A}(MvA, Iof, checkoverflow = false)
end
@inline function OffsetArray{T, N, AA}(parent::AbstractArray{<:Any,N}, offsets::NTuple{N, Integer}; kw...) where {T, N, AA<:AbstractArray{T,N}}
OffsetArray{T, N, AA}(parent, map(Int, offsets)::NTuple{N,Int}; kw...)
end
@inline function OffsetArray{T,N,A}(M::AbstractArray{<:Any,N}, I::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}; kw...) where {T,N,A<:AbstractArray{T,N}}
_checkindices(M, I, "indices")
# Performance gain by wrapping the error in a function: see https://github.com/JuliaLang/julia/issues/37558
throw_dimerr(lA, lI) = throw(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
lM = size(M)
lI = map(length, I)
lM == lI || throw_dimerr(lM, lI)
OffsetArray{T,N,A}(M, map(_offset, axes(M), I); kw...)
end
@inline function OffsetArray{T,N,A}(M::AbstractArray{<:Any,N}, I::Tuple; kw...) where {T,N,A<:AbstractArray{T,N}}
OffsetArray{T,N,A}(M, _toAbstractUnitRanges(to_indices(M, axes(M), I)); kw...)
end
@inline function OffsetArray{T,N,A}(M::AbstractArray{<:Any,N}; kw...) where {T,N,A<:AbstractArray{T,N}}
Mv = no_offset_view(M)
MvA = convert(A, Mv)::A
OffsetArray{T,N,A}(MvA, _offsets(M); kw...)
end
@inline OffsetArray{T,N,A}(M::A; checkoverflow = false) where {T,N,A<:AbstractArray{T,N}} = OffsetArray{T,N,A}(M, ntuple(zero, Val(N)); checkoverflow = checkoverflow)
Base.convert(::Type{T}, M::AbstractArray) where {T<:OffsetArray} = M isa T ? M : T(M)
@inline AbstractArray{T,N}(M::OffsetArray{S,N}) where {T,S,N} = OffsetArray{T}(M)
# array initialization
@inline function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{OffsetAxisKnownLength}}; kw...) where {T,N}
_checkindices(N, inds, "indices")
AA = Array{T,N}(init, map(_indexlength, inds))
OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds); kw...)
end
@inline function OffsetArray{T, N}(init::ArrayInitializer, inds::Tuple; kw...) where {T, N}
OffsetArray{T, N}(init, _toAbstractUnitRanges(inds); kw...)
end
@inline OffsetArray{T,N}(init::ArrayInitializer, inds...; kw...) where {T,N} = OffsetArray{T,N}(init, inds; kw...)
@inline OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}; kw...) where {T,N} = OffsetArray{T,N}(init, inds; kw...)
@inline function OffsetArray{T}(init::ArrayInitializer, inds::Tuple; kw...) where {T}
OffsetArray{T}(init, _toAbstractUnitRanges(inds); kw...)
end
@inline OffsetArray{T}(init::ArrayInitializer, inds...; kw...) where {T} = OffsetArray{T}(init, inds; kw...)
Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
parenttype(A::OffsetArray) = parenttype(typeof(A))
Base.parent(A::OffsetArray) = A.parent
# TODO: Ideally we would delegate to the parent's broadcasting implementation, but that
# is currently broken in sufficiently many implementation, namely RecursiveArrayTools, DistributedArrays
# and StaticArrays, that it will take concentrated effort to get this working across the ecosystem.
# The goal would be to have `OffsetArray(CuArray) .+ 1 == OffsetArray{CuArray}`.
# Base.Broadcast.BroadcastStyle(::Type{<:OffsetArray{<:Any, <:Any, AA}}) where AA = Base.Broadcast.BroadcastStyle(AA)
@inline Base.size(A::OffsetArray) = size(parent(A))
@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.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1)) # we only need to specialize this one
# Utils to translate a function to the parent while preserving offsets
unwrap(x) = x, identity
unwrap(x::OffsetArray) = parent(x), data -> OffsetArray(data, x.offsets, checkoverflow = false)
function parent_call(f, x)
parent, wrap_offset = unwrap(x)
wrap_offset(f(parent))
end
Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
similar(parent(A), T, dims)
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))
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)
OffsetArray(P, map(_offset, axes(P), shape))
end
# Try to use the axes to generate the parent array type
# This is useful if the axes have special meanings, such as with static arrays
# This method is hit if at least one axis provided to similar(A, T, axes) is an IdOffsetRange
# For example this is hit when similar(A::OffsetArray) is called,
# which expands to similar(A, eltype(A), axes(A))
_similar_axes_or_length(A, T, ax, ::Any) = similar(A, T, ax)
_similar_axes_or_length(AT, ax, ::Any) = similar(AT, ax)
# Handle the general case by resorting to lengths along each axis
# This is hit if none of the axes provided to similar(A, T, axes) are IdOffsetRanges,
# and if similar(A, T, axes::AX) is not defined for the type AX.
# In this case the best that we can do is to create a mutable array of the correct size
_similar_axes_or_length(A, T, ax::I, ::I) where {I} = similar(A, T, map(_indexlength, ax))
_similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength, ax))
# reshape accepts a single colon
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
AR = reshape(no_offset_view(A), map(_indexlength, inds))
O = OffsetArray(AR, map(_offset, axes(AR), inds))
return _popreshape(O, axes(AR), _filterreshapeinds(inds))
end
# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
# Short-circuit for AbstractVectors if the axes are compatible to get around the Base restriction
# to 1-based vectors
function _reshape(A::AbstractVector, inds::Tuple{OffsetAxis})
@noinline throw_dimerr(ind::Integer) = throw(
DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with length $ind"))
@noinline throw_dimerr(ind) = throw(
DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with indices $ind"))
_checksize(first(inds), size(A,1)) || throw_dimerr(first(inds))
A
end
_reshape(A, inds) = _reshape2(A, inds)
_reshape2(A, inds) = reshape(A, inds)
# avoid a stackoverflow by relegating to the parent if no_offset_view returns an offsetarray
_reshape2(A::OffsetArray, inds) = reshape(parent(A), inds)
_reshape_nov(A, inds) = _reshape(no_offset_view(A), inds)
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
OffsetArray(_reshape(parent(A), inds), map(_toaxis, inds))
# And for non-offset axes, we can just return a reshape of the parent directly
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(A, inds)
Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(A, inds)
Base.reshape(A::OffsetVector, ::Colon) = A
Base.reshape(A::OffsetVector, ::Tuple{Colon}) = A
Base.reshape(A::OffsetArray, ::Colon) = reshape(A, (Colon(),))
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(A, inds)
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(A, inds)
# The following two additional methods for Colon are added to resolve method ambiguities to
# Base: https://github.com/JuliaLang/julia/pull/45387#issuecomment-1132859663
Base.reshape(A::OffsetArray, inds::Colon) = _reshape_nov(A, inds)
Base.reshape(A::OffsetArray, inds::Tuple{Colon}) = _reshape_nov(A, inds)
# permutedims in Base does not preserve axes, and can not be fixed in a non-breaking way
# This is a stopgap solution
Base.permutedims(v::OffsetVector) = reshape(v, (1, axes(v, 1)))
Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
fill!(similar(Array{typeof(v)}, inds), v)
Base.zeros(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
fill!(similar(Array{T}, inds), zero(T))
Base.ones(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
fill!(similar(Array{T}, inds), one(T))
Base.trues(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
fill!(similar(BitArray, inds), true)
Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
fill!(similar(BitArray, inds), false)
Base.zero(A::OffsetArray) = parent_call(zero, A)
Base.fill!(A::OffsetArray, x) = parent_call(Ap -> fill!(Ap, x), A)
## Indexing
# Note this gets the index of the parent *array*, not the index of the parent *range*
# Here's how one can think about this:
# Δi = i - first(r)
# i′ = first(r.parent) + Δi
# and one obtains the result below.
parentindex(r::IdOffsetRange, i) = oftype(r.offset, i) - r.offset
@propagate_inbounds Base.getindex(A::OffsetArray{<:Any,0}) = A.parent[]
@inline function Base.getindex(A::OffsetArray{<:Any,N}, I::Vararg{Int,N}) where N
@boundscheck checkbounds(A, I...)
J = map(parentindex, axes(A), I)
@inbounds parent(A)[J...]
end
@propagate_inbounds Base.getindex(A::OffsetArray{<:Any,N}, c::Vararg{Colon,N}) where N =
parent_call(x -> getindex(x, c...), A)
# With one Colon we use linear indexing.
# In this case we may forward the index to the parent, as the information about the axes is lost
# The exception to this is with OffsetVectors where the axis information is preserved,
# but that case is handled by getindex(::OffsetArray{<:Any,N}, ::Vararg{Colon,N})
@propagate_inbounds Base.getindex(A::OffsetArray, c::Colon) = A.parent[:]
@inline function Base.getindex(A::OffsetVector, i::Int)
@boundscheck checkbounds(A, i)
@inbounds parent(A)[parentindex(Base.axes1(A), i)]
end
@propagate_inbounds Base.getindex(A::OffsetArray, i::Int) = parent(A)[i]
@inline function Base.setindex!(A::OffsetArray{T,N}, val, I::Vararg{Int,N}) where {T,N}
@boundscheck checkbounds(A, I...)
J = map(parentindex, axes(A), I)
@inbounds parent(A)[J...] = val
A
end
@inline function Base.setindex!(A::OffsetVector, val, i::Int)
@boundscheck checkbounds(A, i)
@inbounds parent(A)[parentindex(Base.axes1(A), i)] = val
A
end
@propagate_inbounds function Base.setindex!(A::OffsetArray, val, i::Int)
parent(A)[i] = val
A
end
@inline Base.iterate(a::OffsetArray, i...) = iterate(parent(a), i...)
Base.in(x, A::OffsetArray) = in(x, parent(A))
Base.copy(A::OffsetArray) = parent_call(copy, A)
Base.strides(A::OffsetArray) = strides(parent(A))
Base.elsize(::Type{OffsetArray{T,N,A}}) where {T,N,A} = Base.elsize(A)
Base.cconvert(::Type{Ptr{T}}, A::OffsetArray{T}) where {T} = Base.cconvert(Ptr{T}, parent(A))
# For fast broadcasting: ref https://discourse.julialang.org/t/why-is-there-a-performance-hit-on-broadcasting-with-offsetarrays/32194
Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)
### Special handling for AbstractRange
const OffsetRange{T} = OffsetVector{T,<:AbstractRange{T}}
const OffsetUnitRange{T} = OffsetVector{T,<:AbstractUnitRange{T}}
Base.step(a::OffsetRange) = step(parent(a))
Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, or::OffsetRange) = Base.checkindex(Bool, inds, parent(or))
# Certain special methods for linear indexing with integer ranges (or OffsetRanges)
# These may bypass the default getindex(A, I...) pathway if the parent types permit this
# For example AbstractUnitRanges and Arrays have special linear indexing behavior defined
# If both the arguments are offset, we may unwrap the indices to call (::OffsetArray)[::AbstractRange{Int}]
@propagate_inbounds function Base.getindex(A::OffsetArray, r::OffsetRange{Int})
_indexedby(A[parent(r)], axes(r))
end
# If the indices are offset, we may unwrap them and pass the parent to getindex
@propagate_inbounds function Base.getindex(A::AbstractRange, r::OffsetRange{Int})
_indexedby(A[parent(r)], axes(r))
end
# An OffsetUnitRange might use the rapid getindex(::Array, ::AbstractUnitRange{Int}) for contiguous indexing
@propagate_inbounds function Base.getindex(A::Array, r::OffsetUnitRange{Int})
B = A[_contiguousindexingtype(parent(r))]
OffsetArray(B, axes(r), checkoverflow = false)
end
# Linear Indexing of OffsetArrays with AbstractUnitRanges may use the faster contiguous indexing methods
@inline function Base.getindex(A::OffsetArray, r::AbstractUnitRange{Int})
@boundscheck checkbounds(A, r)
# nD OffsetArrays do not have their linear indices shifted, so we may forward the indices provided to the parent
@inbounds B = parent(A)[_contiguousindexingtype(r)]
_indexedby(B, axes(r))
end
@inline function Base.getindex(A::OffsetVector, r::AbstractUnitRange{Int})
@boundscheck checkbounds(A, r)
# OffsetVectors may have their linear indices shifted, so we subtract the offset from the indices provided
@inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
_indexedby(B, axes(r))
end
# This method added mainly to index an OffsetRange with another range
@inline function Base.getindex(A::OffsetVector, r::AbstractRange{Int})
@boundscheck checkbounds(A, r)
@inbounds B = parent(A)[_subtractoffset(r, A.offsets[1])]
_indexedby(B, axes(r))
end
# In general we would pass through getindex(A, I...) which calls to_indices(A, I) and finally to_index(I)
# An OffsetUnitRange{Int} has an equivalent IdOffsetRange with the same values and axes,
# something similar also holds for OffsetUnitRange{BigInt}
# We may replace the former with the latter in an indexing operation to obtain a performance boost
@inline function Base.to_index(r::OffsetUnitRange{<:Union{Int,BigInt}})
of = first(axes(r,1)) - 1
IdOffsetRange(_subtractoffset(parent(r), of), of)
end
@inline function _boundscheck_index_retaining_axes(r, s)
@boundscheck checkbounds(r, s)
@inbounds pr = r[UnitRange(s)]
_indexedby(pr, axes(s))
end
@inline _boundscheck_return(r, s) = (@boundscheck checkbounds(r, s); s)
for OR in [:IIUR, :IdOffsetRange]
for R in [:StepRange, :StepRangeLen, :LinRange, :UnitRange]
@eval @inline Base.getindex(r::$R, s::$OR) = _boundscheck_index_retaining_axes(r, s)
end
# this method is needed for ambiguity resolution
@eval @inline function Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::$OR) where T
_boundscheck_index_retaining_axes(r, s)
end
end
Base.getindex(r::Base.OneTo, s::IdOffsetRange) = _boundscheck_index_retaining_axes(r, s)
# These methods are added to avoid ambiguities with Base.
# The ones involving Base types should be ported to Base and version-limited here
@inline Base.getindex(r::IdentityUnitRange, s::IIUR) = _boundscheck_return(r, s)
@inline Base.getindex(r::IdentityUnitRange, s::IdOffsetRange) = _boundscheck_return(r, s)
if IdentityUnitRange !== Base.Slice
@inline Base.getindex(r::Base.Slice, s::IIUR) = _boundscheck_return(r, s)
@inline Base.getindex(r::Base.Slice, s::IdOffsetRange) = _boundscheck_return(r, s)
end
# eltype conversion
# This may use specialized map methods for the parent
Base.map(::Type{T}, O::OffsetArray) where {T} = parent_call(x -> map(T, x), O)
Base.map(::Type{T}, r::IdOffsetRange) where {T<:Real} = _indexedby(map(T, UnitRange(r)), axes(r))
if eltype(IIUR) === Int
# This is type-piracy, but there is no way to convert an IdentityUnitRange to a non-Int type in Base
Base.map(::Type{T}, r::IdentityUnitRange) where {T<:Real} = _indexedby(map(T, UnitRange(r)), axes(r))
end
# mapreduce is faster with an IdOffsetRange than with an OffsetUnitRange
# We therefore convert OffsetUnitRanges to IdOffsetRanges with the same values and axes
function Base.mapreduce(f, op, A1::OffsetUnitRange{<:Integer}, As::OffsetUnitRange{<:Integer}...; kw...)
As = (A1, As...)
ofs = map(A -> first(axes(A,1)) - 1, As)
AIds = map((A, of) -> IdOffsetRange(_subtractoffset(parent(A), of), of), As, ofs)
mapreduce(f, op, AIds...; kw...)
end
# Optimize certain reductions that treat an OffsetVector as a list
for f in [:minimum, :maximum, :extrema, :sum]
@eval Base.$f(r::OffsetRange) = $f(parent(r))
end
function Base.show(io::IO, r::OffsetRange)
show(io, r.parent)
print(io, " with indices ", UnitRange(axes(r, 1)))
end
Base.show(io::IO, ::MIME"text/plain", r::OffsetRange) = show(io, r)
### Some mutating functions defined only for OffsetVector ###
Base.resize!(A::OffsetVector, nl::Integer) = (resize!(A.parent, nl); A)
Base.push!(A::OffsetVector, x...) = (push!(A.parent, x...); A)
Base.pop!(A::OffsetVector) = pop!(A.parent)
Base.append!(A::OffsetVector, items) = (append!(A.parent, items); A)
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)
# These functions keep the summary compact
function Base.inds2string(inds::Tuple{Vararg{Union{IdOffsetRange, IdentityUnitRange{<:IdOffsetRange}}}})
Base.inds2string(map(UnitRange, inds))
end
Base.showindices(io::IO, ind1::IdOffsetRange, inds::IdOffsetRange...) = Base.showindices(io, map(UnitRange, (ind1, inds...))...)
function Base.showarg(io::IO, @nospecialize(a::OffsetArray), toplevel)
print(io, "OffsetArray(")
Base.showarg(io, parent(a), false)
Base.showindices(io, axes(a)...)
print(io, ')')
if toplevel
print(io, " with eltype ", eltype(a))
end
end
function Base.replace_in_print_matrix(A::OffsetArray{<:Any,2}, i::Integer, j::Integer, s::AbstractString)
J = map(parentindex, axes(A), (i,j))
Base.replace_in_print_matrix(parent(A), J..., s)
end
function Base.replace_in_print_matrix(A::OffsetArray{<:Any,1}, i::Integer, j::Integer, s::AbstractString)
ip = parentindex(axes(A,1), i)
Base.replace_in_print_matrix(parent(A), ip, j, s)
end
no_offset_view(A::OffsetArray) = no_offset_view(parent(A))
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)
#####
# center/centered
# These two helpers are deliberately not exported; their meaning can be very different in
# other scenarios and will be very likely to cause name conflicts if exported.
#####
function center(A::AbstractArray, r::RoundingMode=RoundDown)
map(axes(A)) do inds
round(Int, (length(inds)-1)/2, r) + first(inds)
end
end
centered(A::AbstractArray, cp::Dims=center(A)) = OffsetArray(A, .-cp)
centered(A::AbstractArray, i::CartesianIndex) = centered(A, Tuple(i))
# we may pass the searchsorted* functions to the parent, and wrap the offset
for f in [:searchsortedfirst, :searchsortedlast, :searchsorted]
_safe_f = Symbol("_safe_" * String(f))
@eval function $_safe_f(v::OffsetArray, x, ilo, ihi, o::Base.Ordering)
offset = firstindex(v) - firstindex(parent(v))
$f(parent(v), x, ilo - offset, ihi - offset, o) .+ offset
end
@eval Base.$f(v::OffsetVector, x, ilo::T, ihi::T, o::Base.Ordering) where T<:Integer =
$_safe_f(v, x, ilo, ihi, o)
end
##
# Deprecations
##
# This is a bad API design as it introduces counter intuitive results (#250)
@deprecate centered(A::AbstractArray, r::RoundingMode) OffsetArray(A, .-center(A, r)) false
end # module
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...