Skip to content

Commit c2fd49e

Browse files
johnnychen94johnnychen94
committed
update testhelpers/OffsetArrays
1 parent 4ef8313 commit c2fd49e

File tree

2 files changed

+228
-91
lines changed

2 files changed

+228
-91
lines changed

‎stdlib/Distributed/test/splitrange.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ using Distributed: splitrange
2222
@test splitrange(-1, 1, 4) == Array{UnitRange{Int64},1}([-1:-1,0:0,1:1])
2323

2424
const BASE_TEST_PATH = joinpath(Sys.BINDIR, "..", "share", "julia", "test")
25-
isdefined(Main, :OffsetArrays) || @eval Main include(joinpath($(BASE_TEST_PATH), "testhelpers", "OffsetArrays.jl"))
25+
isdefined(Main, :OffsetArrays) || @eval Main @everywhere include(joinpath($(BASE_TEST_PATH), "testhelpers", "OffsetArrays.jl"))
2626
using .Main.OffsetArrays
2727

2828
oa = OffsetArray([123, -345], (-2,))

‎test/testhelpers/OffsetArrays.jl

Lines changed: 227 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -5,141 +5,278 @@
55
# This test file is designed to exercise support for generic indexing,
66
# even though offset arrays aren't implemented in Base.
77

8-
module OffsetArrays
8+
module OffsetArrays # OffsetArrays@1.0.0 without compat and docstrings
99

10-
using Base: Indices, IndexCartesian, IndexLinear, tail
10+
using Base: Indices, tail, @propagate_inbounds
11+
@static if !isdefined(Base, :IdentityUnitRange)
12+
const IdentityUnitRange = Base.Slice
13+
else
14+
using Base: IdentityUnitRange
15+
end
16+
17+
export OffsetArray, OffsetVector
18+
19+
struct IdOffsetRange{T<:Integer,I<:AbstractUnitRange{T}} <: AbstractUnitRange{T}
20+
parent::I
21+
offset::T
22+
end
23+
IdOffsetRange(r::AbstractUnitRange{T}, offset::Integer) where T =
24+
IdOffsetRange{T,typeof(r)}(r, convert(T, offset))
25+
26+
@inline Base.axes(r::IdOffsetRange) = (Base.axes1(r),)
27+
@inline Base.axes1(r::IdOffsetRange) = IdOffsetRange(Base.axes1(r.parent), r.offset)
28+
@inline Base.unsafe_indices(r::IdOffsetRange) = (r,)
29+
@inline Base.length(r::IdOffsetRange) = length(r.parent)
30+
31+
function Base.iterate(r::IdOffsetRange)
32+
ret = iterate(r.parent)
33+
ret === nothing && return nothing
34+
return (ret[1] + r.offset, ret[2])
35+
end
36+
function Base.iterate(r::IdOffsetRange, i) where T
37+
ret = iterate(r.parent, i)
38+
ret === nothing && return nothing
39+
return (ret[1] + r.offset, ret[2])
40+
end
41+
42+
@inline Base.first(r::IdOffsetRange) = first(r.parent) + r.offset
43+
@inline Base.last(r::IdOffsetRange) = last(r.parent) + r.offset
44+
45+
@propagate_inbounds Base.getindex(r::IdOffsetRange, i::Integer) = r.parent[i - r.offset] + r.offset
46+
@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
47+
return r.parent[s .- r.offset] .+ r.offset
48+
end
49+
50+
Base.show(io::IO, r::IdOffsetRange) = print(io, first(r), ':', last(r))
1151

12-
export OffsetArray
52+
# Optimizations
53+
@inline Base.checkindex(::Type{Bool}, inds::IdOffsetRange, i::Real) = Base.checkindex(Bool, inds.parent, i - inds.offset)
1354

55+
## OffsetArray
1456
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
1557
parent::AA
1658
offsets::NTuple{N,Int}
1759
end
1860
OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
1961

20-
OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N} = OffsetArray{T,N,typeof(A)}(A, offsets)
21-
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} = OffsetArray(A, offsets)
62+
## OffsetArray constructors
2263

23-
OffsetArray{T,N}(::UndefInitializer, inds::Indices{N}) where {T,N} =
24-
OffsetArray{T,N,Array{T,N}}(Array{T,N}(undef, map(length, inds)), map(indsoffset, inds))
25-
OffsetArray{T}(::UndefInitializer, inds::Indices{N}) where {T,N} =
26-
OffsetArray{T,N}(undef, inds)
64+
offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
65+
offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
2766

28-
Base.IndexStyle(::Type{T}) where {T<:OffsetArray} = Base.IndexStyle(parenttype(T))
67+
function OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N}
68+
OffsetArray{T,N,typeof(A)}(A, offsets)
69+
end
70+
OffsetArray(A::AbstractArray{T,0}, offsets::Tuple{}) where T =
71+
OffsetArray{T,0,typeof(A)}(A, ())
72+
73+
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
74+
OffsetArray(A, offsets)
75+
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray(A, ())
76+
77+
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
78+
OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
79+
OffsetArray(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
80+
OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
81+
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
82+
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
83+
84+
# OffsetVector constructors
85+
OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
86+
OffsetVector{T}(init::ArrayInitializer, inds::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds)
87+
88+
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
89+
axparent = axes(A)
90+
lA = map(length, axparent)
91+
lI = map(length, inds)
92+
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"))
93+
OffsetArray(A, map(offset, axparent, inds))
94+
end
95+
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N} =
96+
OffsetArray(A, inds)
97+
98+
# avoid a level of indirection when nesting OffsetArrays
99+
function OffsetArray(A::OffsetArray, inds::NTuple{N,AbstractUnitRange}) where {N}
100+
OffsetArray(parent(A), inds)
101+
end
102+
OffsetArray(A::OffsetArray{T,0}, inds::Tuple{}) where {T} = OffsetArray(parent(A), ())
103+
# OffsetArray(A::OffsetArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
104+
105+
Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
29106
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
30107
parenttype(A::OffsetArray) = parenttype(typeof(A))
31108

32109
Base.parent(A::OffsetArray) = A.parent
33110

34-
Base.size(A::OffsetArray) = size(A.parent)
35-
Base.size(A::OffsetArray, d) = size(A.parent, d)
36111
Base.eachindex(::IndexCartesian, A::OffsetArray) = CartesianIndices(axes(A))
37-
Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
112+
Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
38113

39-
# Implementations of indices and axes1. Since bounds-checking is
40-
# performance-critical and relies on indices, these are usually worth
41-
# optimizing thoroughly.
42-
@inline Base.axes(A::OffsetArray, d) = 1 <= d <= length(A.offsets) ? Base.IdentityUnitRange(axes(parent(A))[d] .+ A.offsets[d]) : Base.IdentityUnitRange(1:1)
43-
@inline Base.axes(A::OffsetArray) = _indices(axes(parent(A)), A.offsets) # would rather use ntuple, but see #15276
44-
@inline _indices(inds, offsets) = (Base.IdentityUnitRange(inds[1] .+ offsets[1]), _indices(tail(inds), tail(offsets))...)
45-
_indices(::Tuple{}, ::Tuple{}) = ()
46-
Base.axes1(A::OffsetArray{T,0}) where {T} = Base.IdentityUnitRange(1:1) # we only need to specialize this one
114+
@inline Base.size(A::OffsetArray) = size(parent(A))
115+
@inline Base.size(A::OffsetArray, d) = size(parent(A), d)
47116

48-
const OffsetAxis = Union{Integer, UnitRange, Base.IdentityUnitRange{<:UnitRange}, Base.OneTo}
49-
function Base.similar(A::OffsetArray, T::Type, dims::Dims)
50-
B = similar(parent(A), T, dims)
117+
@inline Base.axes(A::OffsetArray) = map(IdOffsetRange, axes(parent(A)), A.offsets)
118+
@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : Base.OneTo(1)
119+
@inline Base.axes1(A::OffsetArray{T,0}) where {T} = Base.OneTo(1) # we only need to specialize this one
120+
121+
const OffsetAxis = Union{Integer, UnitRange, Base.OneTo, IdentityUnitRange, IdOffsetRange, Colon}
122+
Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
123+
similar(parent(A), T, dims)
124+
function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where T
125+
B = similar(A, T, map(indexlength, inds))
126+
return OffsetArray(B, map(offset, axes(B), inds))
51127
end
52-
function Base.similar(A::AbstractArray, T::Type, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
53-
B = similar(A, T, map(indslength, inds))
54-
OffsetArray(B, map(indsoffset, inds))
128+
129+
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
130+
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
131+
AR = reshape(A, map(indexlength, inds))
132+
return OffsetArray(AR, map(offset, axes(AR), inds))
55133
end
56134

57-
Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray} =
58-
OffsetArray(T(undef, map(indslength, shape)), map(indsoffset, shape))
135+
# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
136+
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
137+
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
138+
OffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
139+
# And for non-offset axes, we can just return a reshape of the parent directly
140+
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
141+
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
142+
Base.reshape(A::OffsetArray, ::Colon) = A
143+
Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), inds)
144+
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
59145

60-
Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) = OffsetArray(reshape(A, map(indslength, inds)), map(indsoffset, inds))
146+
function Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray}
147+
P = T(undef, map(indexlength, shape))
148+
OffsetArray(P, map(offset, axes(P), shape))
149+
end
61150

62151
Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
63-
fill!(OffsetArray(Array{typeof(v), N}(undef, map(indslength, inds)), map(indsoffset, inds)), v)
152+
fill!(OffsetArray(Array{typeof(v), N}(undef, map(indexlength, inds)), map(indexoffset, inds)), v)
64153
Base.zeros(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
65-
fill!(OffsetArray(Array{T, N}(undef, map(indslength, inds)), map(indsoffset, inds)), zero(T))
154+
fill!(OffsetArray(Array{T, N}(undef, map(indexlength, inds)), map(indexoffset, inds)), zero(T))
66155
Base.ones(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =
67-
fill!(OffsetArray(Array{T, N}(undef, map(indslength, inds)), map(indsoffset, inds)), one(T))
156+
fill!(OffsetArray(Array{T, N}(undef, map(indexlength, inds)), map(indexoffset, inds)), one(T))
68157
Base.trues(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
69-
fill!(OffsetArray(BitArray{N}(undef, map(indslength, inds)), map(indsoffset, inds)), true)
158+
fill!(OffsetArray(BitArray{N}(undef, map(indexlength, inds)), map(indexoffset, inds)), true)
70159
Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
71-
fill!(OffsetArray(BitArray{N}(undef, map(indslength, inds)), map(indsoffset, inds)), false)
72-
73-
@inline function Base.getindex(A::OffsetArray{T,N}, I::Vararg{Int,N}) where {T,N}
74-
checkbounds(A, I...)
75-
@inbounds ret = parent(A)[offset(A.offsets, I)...]
76-
ret
77-
end
78-
# Vectors don't support one-based linear indexing; they always use the offsets
79-
@inline function Base.getindex(A::OffsetVector, i::Int)
80-
checkbounds(A, i)
81-
@inbounds ret = parent(A)[offset(A.offsets, (i,))[1]]
82-
ret
83-
end
84-
# But multidimensional arrays allow one-based linear indexing
85-
@inline function Base.getindex(A::OffsetArray, i::Int)
86-
checkbounds(A, i)
87-
@inbounds ret = parent(A)[i]
88-
ret
89-
end
90-
@inline function Base.setindex!(A::OffsetArray{T,N}, val, I::Vararg{Int,N}) where {T,N}
91-
checkbounds(A, I...)
92-
@inbounds parent(A)[offset(A.offsets, I)...] = val
160+
fill!(OffsetArray(BitArray{N}(undef, map(indexlength, inds)), map(indexoffset, inds)), false)
161+
162+
## Indexing
163+
164+
# Note this gets the index of the parent *array*, not the index of the parent *range*
165+
# Here's how one can think about this:
166+
# Δi = i - first(r)
167+
# i′ = first(r.parent) + Δi
168+
# and one obtains the result below.
169+
parentindex(r::IdOffsetRange, i) = i - r.offset
170+
171+
@propagate_inbounds function Base.getindex(A::OffsetArray{T,N}, I::Vararg{Int,N}) where {T,N}
172+
J = map(parentindex, axes(A), I)
173+
return parent(A)[J...]
174+
end
175+
176+
@propagate_inbounds Base.getindex(A::OffsetVector, i::Int) = parent(A)[parentindex(Base.axes1(A), i)]
177+
@propagate_inbounds Base.getindex(A::OffsetArray, i::Int) = parent(A)[i]
178+
179+
@propagate_inbounds function Base.setindex!(A::OffsetArray{T,N}, val, I::Vararg{Int,N}) where {T,N}
180+
@boundscheck checkbounds(A, I...)
181+
J = @inbounds map(parentindex, axes(A), I)
182+
@inbounds parent(A)[J...] = val
93183
val
94184
end
95-
@inline function Base.setindex!(A::OffsetVector, val, i::Int)
96-
checkbounds(A, i)
97-
@inbounds parent(A)[offset(A.offsets, (i,))[1]] = val
185+
186+
@propagate_inbounds function Base.setindex!(A::OffsetVector, val, i::Int)
187+
@boundscheck checkbounds(A, i)
188+
@inbounds parent(A)[parentindex(Base.axes1(A), i)] = val
98189
val
99190
end
100-
@inline function Base.setindex!(A::OffsetArray, val, i::Int)
101-
checkbounds(A, i)
191+
@propagate_inbounds function Base.setindex!(A::OffsetArray, val, i::Int)
192+
@boundscheck checkbounds(A, i)
102193
@inbounds parent(A)[i] = val
103194
val
104195
end
105196

106-
@inline function Base.deleteat!(A::OffsetArray, i::Int)
107-
checkbounds(A, i)
108-
@inbounds deleteat!(parent(A), offset(A.offsets, (i,))[1])
109-
end
197+
# For fast broadcasting: ref https://discourse.julialang.org/t/why-is-there-a-performance-hit-on-broadcasting-with-offsetarrays/32194
198+
Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
199+
Broadcast.broadcast_unalias(dest::OffsetArray, src::OffsetArray) = parent(dest) === parent(src) ? src : Broadcast.unalias(dest, src)
110200

111-
@inline function Base.deleteat!(A::OffsetArray{T,N}, I::Vararg{Int, N}) where {T,N}
112-
checkbounds(A, I...)
113-
@inbounds deleteat!(parent(A), offset(A.offsets, I)...)
114-
end
201+
### Special handling for AbstractRange
115202

116-
@inline function Base.deleteat!(A::OffsetArray, i::UnitRange{Int})
117-
checkbounds(A, first(i))
118-
checkbounds(A, last(i))
119-
first_idx = offset(A.offsets, (first(i),))[1]
120-
last_idx = offset(A.offsets, (last(i),))[1]
121-
@inbounds deleteat!(parent(A), first_idx:last_idx)
122-
end
203+
const OffsetRange{T} = OffsetArray{T,1,<:AbstractRange{T}}
204+
const IIUR = IdentityUnitRange{S} where S<:AbstractUnitRange{T} where T<:Integer
205+
206+
Base.step(a::OffsetRange) = step(parent(a))
207+
208+
Base.getindex(a::OffsetRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
209+
Base.getindex(a::OffsetRange, r::AbstractRange) = a.parent[r .- a.offsets[1]]
210+
Base.getindex(a::AbstractRange, r::OffsetRange) = OffsetArray(a[parent(r)], r.offsets)
123211

124-
function Base.push!(a::OffsetArray{T,1}, item) where T
125-
# convert first so we don't grow the array if the assignment won't work
126-
itemT = convert(T, item)
127-
resize!(a, length(a)+1)
128-
a[end] = itemT
129-
return a
212+
@propagate_inbounds Base.getindex(r::UnitRange, s::IIUR) =
213+
OffsetArray(r[s.indices], s)
214+
215+
@propagate_inbounds Base.getindex(r::StepRange, s::IIUR) =
216+
OffsetArray(r[s.indices], s)
217+
218+
@inline @propagate_inbounds Base.getindex(r::StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision}, s::IIUR) where T =
219+
OffsetArray(r[s.indices], s)
220+
@inline @propagate_inbounds Base.getindex(r::StepRangeLen{T}, s::IIUR) where {T} =
221+
OffsetArray(r[s.indices], s)
222+
223+
@inline @propagate_inbounds Base.getindex(r::LinRange, s::IIUR) =
224+
OffsetArray(r[s.indices], s)
225+
226+
function Base.show(io::IO, r::OffsetRange)
227+
show(io, r.parent)
228+
o = r.offsets[1]
229+
print(io, " with indices ", o+1:o+length(r))
130230
end
231+
Base.show(io::IO, ::MIME"text/plain", r::OffsetRange) = show(io, r)
232+
233+
### Convenience functions ###
131234

132-
# Computing a shifted index (subtracting the offset)
133-
offset(offsets::NTuple{N,Int}, inds::NTuple{N,Int}) where {N} = _offset((), offsets, inds)
134-
_offset(out, ::Tuple{}, ::Tuple{}) = out
135-
@inline _offset(out, offsets, inds) = _offset((out..., inds[1]-offsets[1]), Base.tail(offsets), Base.tail(inds))
235+
Base.fill(x, inds::Tuple{UnitRange,Vararg{UnitRange}}) =
236+
fill!(OffsetArray{typeof(x)}(undef, inds), x)
237+
@inline Base.fill(x, ind1::UnitRange, inds::UnitRange...) = fill(x, (ind1, inds...))
136238

137-
indsoffset(r::AbstractRange) = first(r) - 1
138-
indsoffset(i::Integer) = 0
139-
indslength(r::AbstractRange) = length(r)
140-
indslength(i::Integer) = i
141239

240+
### Some mutating functions defined only for OffsetVector ###
142241

143242
Base.resize!(A::OffsetVector, nl::Integer) = (resize!(A.parent, nl); A)
243+
Base.push!(A::OffsetVector, x...) = (push!(A.parent, x...); A)
244+
Base.pop!(A::OffsetVector) = pop!(A.parent)
245+
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)
246+
247+
### Low-level utilities ###
144248

249+
indexoffset(r::AbstractRange) = first(r) - 1
250+
indexoffset(i::Integer) = 0
251+
indexoffset(i::Colon) = 0
252+
indexlength(r::AbstractRange) = length(r)
253+
indexlength(i::Integer) = i
254+
indexlength(i::Colon) = Colon()
255+
256+
function Base.showarg(io::IO, a::OffsetArray, toplevel)
257+
print(io, "OffsetArray(")
258+
Base.showarg(io, parent(a), false)
259+
if ndims(a) > 0
260+
print(io, ", ")
261+
printindices(io, axes(a)...)
262+
end
263+
print(io, ')')
264+
toplevel && print(io, " with eltype ", eltype(a))
265+
end
266+
printindices(io::IO, ind1, inds...) =
267+
(print(io, _unslice(ind1), ", "); printindices(io, inds...))
268+
printindices(io::IO, ind1) = print(io, _unslice(ind1))
269+
_unslice(x) = x
270+
_unslice(x::IdentityUnitRange) = x.indices
271+
272+
function no_offset_view(A::AbstractArray)
273+
if Base.has_offset_axes(A)
274+
OffsetArray(A, map(r->1-first(r), axes(A)))
275+
else
276+
A
277+
end
145278
end
279+
280+
no_offset_view(A::OffsetArray) = no_offset_view(parent(A))
281+
282+
end # module

0 commit comments

Comments
 (0)