Skip to content

Commit a78ee17

Browse files
johnnychen94johnnychen94
authored
rework OffsetArrays constructor (#148)
* rationalize and rework the OffsetArray constructor for a better and clearer dispatching routes * rename internal helpers with "_" prefix
1 parent f50ebe9 commit a78ee17

File tree

3 files changed

+380
-299
lines changed

3 files changed

+380
-299
lines changed

‎src/OffsetArrays.jl

Lines changed: 70 additions & 139 deletions
Original file line numberDiff line numberDiff line change
@@ -10,62 +10,15 @@ end
1010
export OffsetArray, OffsetMatrix, OffsetVector
1111

1212
include("axes.jl")
13+
include("utils.jl")
1314

14-
## OffsetArray
15-
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
16-
parent::AA
17-
offsets::NTuple{N,Int}
18-
function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}) where {T, N, AA<:AbstractArray}
19-
overflow_check.(axes(parent), offsets)
20-
new{T, N, AA}(parent, offsets)
21-
end
22-
end
23-
OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
24-
OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
25-
26-
function overflow_check(r, offset::T) where T
27-
# This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
28-
throw_upper_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
29-
throw_lower_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
30-
31-
if offset > 0 && last(r) > typemax(T) - offset
32-
throw_upper_overflow_error()
33-
elseif offset < 0 && first(r) < typemin(T) - offset
34-
throw_lower_overflow_error()
35-
end
36-
end
37-
38-
## OffsetArray constructors
39-
40-
offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
41-
offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
42-
43-
function OffsetArray(A::AbstractArray{T,N}, offsets::NTuple{N,Int}) where {T,N}
44-
OffsetArray{T,N,typeof(A)}(A, offsets)
45-
end
46-
OffsetArray(A::AbstractArray{T,0}, offsets::Tuple{}) where T =
47-
OffsetArray{T,0,typeof(A)}(A, ())
48-
49-
OffsetArray(A::AbstractArray{T,N}, offsets::Vararg{Int,N}) where {T,N} =
50-
OffsetArray(A, offsets)
51-
OffsetArray(A::AbstractArray{T,0}) where {T} = OffsetArray(A, ())
52-
15+
# Technically we know the length of CartesianIndices but we need to convert it first, so here we
16+
# don't put it in OffsetAxisKnownLength.
17+
const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange, IdOffsetRange}
18+
const OffsetAxis = Union{OffsetAxisKnownLength, CartesianIndices, Colon}
5319
const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
54-
OffsetArray{T,N}(init::ArrayInitializer, inds::Indices{N}) where {T,N} =
55-
OffsetArray(Array{T,N}(init, map(indexlength, inds)), map(indexoffset, inds))
56-
OffsetArray{T}(init::ArrayInitializer, inds::Indices{N}) where {T,N} = OffsetArray{T,N}(init, inds)
57-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
58-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{AbstractUnitRange,N}) where {T,N} = OffsetArray{T,N}(init, inds)
59-
60-
# OffsetVector constructors
61-
OffsetVector(A::AbstractVector, offset) = OffsetArray(A, offset)
62-
OffsetVector{T}(init::ArrayInitializer, inds::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds)
63-
64-
# OffsetMatrix constructors
65-
OffsetMatrix(A::AbstractMatrix, offset1, offset2) = OffsetArray(A, offset1, offset2)
66-
OffsetMatrix(A::AbstractMatrix, I::CartesianIndices{2}) = OffsetArray(A, I)
67-
OffsetMatrix{T}(init::ArrayInitializer, inds1::AbstractUnitRange, inds2::AbstractUnitRange) where {T} = OffsetArray{T}(init, inds1, inds2)
6820

21+
## OffsetArray
6922
"""
7023
OffsetArray(A, indices...)
7124
@@ -116,84 +69,74 @@ ERROR: [...]
11669
```
11770
11871
"""
119-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,AbstractUnitRange}) where {T,N}
120-
axparent = axes(A)
121-
lA = map(length, axparent)
122-
lI = map(length, inds)
123-
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"))
124-
OffsetArray(A, map(offset, axparent, inds))
72+
struct OffsetArray{T,N,AA<:AbstractArray} <: AbstractArray{T,N}
73+
parent::AA
74+
offsets::NTuple{N,Int}
75+
function OffsetArray{T, N, AA}(parent::AA, offsets::NTuple{N, Int}) where {T, N, AA<:AbstractArray}
76+
@boundscheck overflow_check.(axes(parent), offsets)
77+
new{T, N, AA}(parent, offsets)
78+
end
12579
end
126-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{AbstractUnitRange,N}) where {T,N} =
127-
OffsetArray(A, inds)
80+
const OffsetVector{T,AA<:AbstractArray} = OffsetArray{T,1,AA}
81+
const OffsetMatrix{T,AA<:AbstractArray} = OffsetArray{T,2,AA}
12882

129-
uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = uncolonindices(axes(A), inds)
130-
uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), uncolonindices(tail(ax), tail(inds))...)
131-
uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), uncolonindices(tail(ax), tail(inds))...)
132-
uncolonindices(::Tuple{}, ::Tuple{}) = ()
83+
function overflow_check(r, offset::T) where T
84+
# This gives some performance boost https://github.com/JuliaLang/julia/issues/33273
85+
throw_upper_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or less than $(typemax(T) - last(r))"))
86+
throw_lower_overflow_error() = throw(ArgumentError("Boundary overflow detected: offset $offset should be equal or greater than $(typemin(T) - first(r))"))
13387

134-
function OffsetArray(A::AbstractArray{T,N}, inds::NTuple{N,Union{AbstractUnitRange, CartesianIndices{1}, Colon}}) where {T,N}
135-
OffsetArray(A, uncolonindices(A, inds))
136-
end
137-
OffsetArray(A::AbstractArray{T,N}, inds::Vararg{Union{AbstractUnitRange, CartesianIndices{1}, Colon},N}) where {T,N} =
138-
OffsetArray(A, uncolonindices(A, inds))
139-
140-
# Specify offsets using CartesianIndices (issue #71)
141-
# Support a mix of AbstractUnitRanges and CartesianIndices{1}
142-
# Extract the range r from CartesianIndices((r,))
143-
function stripCartesianIndices(inds::Tuple{CartesianIndices{1},Vararg{Any}})
144-
I = first(inds)
145-
Ir = convert(Tuple{AbstractUnitRange{Int}}, I) |> first
146-
(Ir, stripCartesianIndices(tail(inds))...)
147-
end
148-
stripCartesianIndices(inds::Tuple)= (first(inds), stripCartesianIndices(tail(inds))...)
149-
stripCartesianIndices(::Tuple{}) = ()
150-
151-
OffsetArray(A::AbstractArray{<:Any,N}, inds::NTuple{N,Union{CartesianIndices{1}, AbstractUnitRange}}) where {N} =
152-
OffsetArray(A, stripCartesianIndices(inds))
153-
OffsetArray(A::AbstractArray{<:Any,N}, inds::Vararg{Union{CartesianIndices{1}, AbstractUnitRange},N}) where {N} =
154-
OffsetArray(A, inds)
155-
156-
# Support an arbitrary CartesianIndices alongside colons and ranges
157-
# The total number of indices should equal ndims(arr)
158-
# We split the CartesianIndices{N} into N CartesianIndices{1} indices to facilitate dispatch
159-
splitCartesianIndices(c::CartesianIndices{0}) = ()
160-
function splitCartesianIndices(c::CartesianIndices)
161-
c1, ct = Base.IteratorsMD.split(c, Val(1))
162-
(c1, splitCartesianIndices(ct)...)
163-
end
164-
function splitCartesianIndices(t::Tuple{CartesianIndices, Vararg{Any}})
165-
(splitCartesianIndices(first(t))..., splitCartesianIndices(tail(t))...)
166-
end
167-
function splitCartesianIndices(t::Tuple)
168-
(first(t), splitCartesianIndices(tail(t))...)
88+
if offset > 0 && last(r) > typemax(T) - offset
89+
throw_upper_overflow_error()
90+
elseif offset < 0 && first(r) < typemin(T) - offset
91+
throw_lower_overflow_error()
92+
end
16993
end
170-
splitCartesianIndices(::Tuple{}) = ()
94+
## OffsetArray constructors
17195

172-
function OffsetArray(A::AbstractArray, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}}})
173-
OffsetArray(A, splitCartesianIndices(inds))
174-
end
175-
function OffsetArray(A::AbstractArray, inds::Vararg{Union{AbstractUnitRange, CartesianIndices, Colon}})
176-
OffsetArray(A, inds)
96+
for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
97+
# The only route out to inner constructor
98+
@eval function $FT(A::AbstractArray{T, N}, offsets::NTuple{N, Integer}) where {T, N}
99+
ndims(A) == N || throw(DimensionMismatch("The number of offsets $(N) should equal ndims(A) = $(ndims(A))"))
100+
OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
101+
end
102+
# nested OffsetArrays
103+
@eval $FT(A::OffsetArray{T, N}, offsets::NTuple{N, Integer}) where {T,N} = $FT(parent(A), A.offsets .+ offsets)
104+
# convert ranges to offsets
105+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N,OffsetAxisKnownLength}) where {T,N}
106+
axparent = axes(A)
107+
lA = map(length, axparent)
108+
lI = map(length, inds)
109+
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"))
110+
$FT(A, map(_offset, axparent, inds))
111+
end
112+
# lower CartesianIndices and Colon
113+
@eval function $FT(A::AbstractArray{T}, inds::NTuple{N, OffsetAxis}) where {T, N}
114+
indsN = _uncolonindices(A, _expandCartesianIndices(inds))
115+
$FT(A, indsN)
116+
end
117+
@eval $FT(A::AbstractArray{T}, inds::Vararg{OffsetAxis,N}) where {T, N} = $FT(A, inds)
177118
end
178119

179-
# Add methods to initialize OffsetArrays using CartesianIndices (issue #71)
180-
function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T,N}
181-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
182-
OffsetArray{T,N}(init, indsN)
120+
# array initialization
121+
function OffsetArray{T,N}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N}
122+
AA = Array{T,N}(init, map(_indexlength, inds))
123+
OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds))
183124
end
184-
function OffsetArray{T}(init::ArrayInitializer, inds::Tuple{Vararg{Union{AbstractUnitRange, CartesianIndices}}}) where {T}
185-
indsN = stripCartesianIndices(splitCartesianIndices(inds))
186-
OffsetArray{T}(init, indsN)
125+
function OffsetArray{T, N}(init::ArrayInitializer, inds::NTuple{NT, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N, NT}
126+
# NT is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
127+
indsN = _expandCartesianIndices(inds)
128+
length(indsN) == N || throw(DimensionMismatch("The number of offsets $(length(indsN)) should equal ndims(A) = $N"))
129+
OffsetArray{T, N}(init, indsN)
187130
end
188-
OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T,N} = OffsetArray{T,N}(init, inds)
189-
OffsetArray{T}(init::ArrayInitializer, inds::Vararg{Union{AbstractUnitRange, CartesianIndices}}) where {T} = OffsetArray{T}(init, inds)
131+
OffsetArray{T,N}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T,N} = OffsetArray{T,N}(init, inds)
190132

191-
# avoid a level of indirection when nesting OffsetArrays
192-
function OffsetArray(A::OffsetArray, offsets::NTuple{N,Int}) where {N}
193-
OffsetArray(parent(A), offsets .+ A.offsets)
133+
OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
134+
function OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
135+
# N is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
136+
indsN = _expandCartesianIndices(inds)
137+
OffsetArray{T, length(indsN)}(init, indsN)
194138
end
195-
OffsetArray(A::OffsetArray{T,0}, inds::Tuple{}) where {T} = OffsetArray(parent(A), ())
196-
# OffsetArray(A::OffsetArray{T,N}, inds::Tuple{}) where {T,N} = error("this should never be called")
139+
OffsetArray{T}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T} = OffsetArray{T}(init, inds)
197140

198141
Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
199142
parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA
@@ -211,27 +154,24 @@ Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
211154
@inline Base.axes(A::OffsetArray, d) = d <= ndims(A) ? IdOffsetRange(axes(parent(A), d), A.offsets[d]) : IdOffsetRange(axes(parent(A), d))
212155
@inline Base.axes1(A::OffsetArray{T,0}) where {T} = IdOffsetRange(axes(parent(A), 1)) # we only need to specialize this one
213156

214-
const OffsetAxisKnownLength = Union{Integer, UnitRange, Base.OneTo, IdentityUnitRange, IdOffsetRange}
215-
216157
Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
217158
similar(parent(A), T, dims)
218159
function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
219-
B = similar(A, T, map(indexlength, inds))
220-
return OffsetArray(B, map(offset, axes(B), inds))
160+
B = similar(A, T, map(_indexlength, inds))
161+
return OffsetArray(B, map(_offset, axes(B), inds))
221162
end
222163

223164
# reshape accepts a single colon
224-
const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
225165
Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
226166
function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
227-
AR = reshape(A, map(indexlength, inds))
228-
return OffsetArray(AR, map(offset, axes(AR), inds))
167+
AR = reshape(A, map(_indexlength, inds))
168+
return OffsetArray(AR, map(_offset, axes(AR), inds))
229169
end
230170

231171
# Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
232172
# an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
233173
Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
234-
OffsetArray(reshape(parent(A), map(indexlength, inds)), map(indexoffset, inds))
174+
OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
235175
# And for non-offset axes, we can just return a reshape of the parent directly
236176
Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
237177
Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
@@ -241,8 +181,8 @@ Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), ind
241181
Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
242182

243183
function Base.similar(::Type{T}, shape::Tuple{OffsetAxis,Vararg{OffsetAxis}}) where {T<:AbstractArray}
244-
P = T(undef, map(indexlength, shape))
245-
OffsetArray(P, map(offset, axes(P), shape))
184+
P = T(undef, map(_indexlength, shape))
185+
OffsetArray(P, map(_offset, axes(P), shape))
246186
end
247187

248188
Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
@@ -346,15 +286,6 @@ Base.pop!(A::OffsetVector) = pop!(A.parent)
346286
Base.append!(A::OffsetVector, items) = (append!(A.parent, items); A)
347287
Base.empty!(A::OffsetVector) = (empty!(A.parent); A)
348288

349-
### Low-level utilities ###
350-
351-
indexoffset(r::AbstractRange) = first(r) - 1
352-
indexoffset(i::Integer) = 0
353-
indexoffset(i::Colon) = 0
354-
indexlength(r::AbstractRange) = length(r)
355-
indexlength(i::Integer) = i
356-
indexlength(i::Colon) = Colon()
357-
358289
# These functions keep the summary compact
359290
function Base.inds2string(inds::Tuple{Vararg{Union{IdOffsetRange, IdentityUnitRange{<:IdOffsetRange}}}})
360291
Base.inds2string(map(UnitRange, inds))

‎src/utils.jl

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
### Low-level utilities ###
2+
3+
_indexoffset(r::AbstractRange) = first(r) - 1
4+
_indexoffset(i::Integer) = 0
5+
_indexoffset(i::Colon) = 0
6+
_indexlength(r::AbstractRange) = length(r)
7+
_indexlength(i::Integer) = i
8+
_indexlength(i::Colon) = Colon()
9+
10+
_offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
11+
_offset(axparent::AbstractUnitRange, ax::CartesianIndices) = _offset(axparent, first(ax.indices))
12+
_offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
13+
14+
_uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = _uncolonindices(axes(A), inds)
15+
_uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), _uncolonindices(tail(ax), tail(inds))...)
16+
_uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), _uncolonindices(tail(ax), tail(inds))...)
17+
_uncolonindices(::Tuple{}, ::Tuple{}) = ()
18+
19+
_expandCartesianIndices(inds::Tuple{<:CartesianIndices, Vararg{Any}}) = (convert(Tuple{Vararg{AbstractUnitRange{Int}}}, inds[1])..., _expandCartesianIndices(Base.tail(inds))...)
20+
_expandCartesianIndices(inds::Tuple{Any,Vararg{Any}}) = (inds[1], _expandCartesianIndices(Base.tail(inds))...)
21+
_expandCartesianIndices(::Tuple{}) = ()

0 commit comments

Comments
 (0)