@@ -126,7 +126,8 @@ Origin(I1::Int, In::Int...) = Origin((I1, In...))
126126# Origin(0) != Origin((0, )) but they work the same with broadcasting
127127Origin (n:: Int ) = Origin {Int} (n)
128128
129- (o:: Origin )(A:: AbstractArray ) = o. index .- first .(axes (A))
129+ (o:: Origin )(A:: AbstractArray ) = OffsetArray (no_offset_view (A), o)
130+ Origin (A:: AbstractArray ) = Origin (first .(axes (A)))
130131
131132# ## Low-level utilities ###
132133
@@ -139,6 +140,26 @@ _indexlength(i::Colon) = Colon()
139140
140141_offset (axparent:: AbstractUnitRange , ax:: AbstractUnitRange ) = first (ax) - first (axparent)
141142_offset (axparent:: AbstractUnitRange , ax:: Integer ) = 1 - first (axparent)
143+ _offsets (A:: AbstractArray ) = map (ax -> first (ax) - 1 , axes (A))
144+ _offsets (A:: AbstractArray , B:: AbstractArray ) = map (_offset, axes (B), axes (A))
145+
146+ # indexing utilities
147+ # These functions are equivalent to the broadcasted operation r .- of
148+ # However these ensure that the result is an AbstractRange even if a specific
149+ # broadcasting behavior is not defined for a custom type
150+ @inline _subtractoffset (r:: AbstractUnitRange , of) = UnitRange (first (r) - of, last (r) - of)
151+ @inline _subtractoffset (r:: AbstractRange , of) = range (first (r) - of, stop = last (r) - of, step = step (r))
152+
153+ @inline _indexedby (r:: AbstractVector , ax:: Tuple{Any} ) = _indexedby (r, ax[1 ])
154+ @inline _indexedby (r:: AbstractUnitRange{<:Integer} , :: Base.OneTo ) = no_offset_view (r)
155+ @inline _indexedby (r:: AbstractUnitRange{Bool} , :: Base.OneTo ) = no_offset_view (r)
156+ @inline _indexedby (r:: AbstractVector , :: Base.OneTo ) = no_offset_view (r)
157+ @inline function _indexedby (r:: AbstractUnitRange{<:Integer} , ax:: AbstractUnitRange )
158+ of = convert (eltype (r), first (ax) - 1 )
159+ IdOffsetRange (_subtractoffset (r, of), of)
160+ end
161+ @inline _indexedby (r:: AbstractUnitRange{Bool} , ax:: AbstractUnitRange ) = OffsetArray (r, ax)
162+ @inline _indexedby (r:: AbstractVector , ax:: AbstractUnitRange ) = OffsetArray (r, ax)
142163
143164abstract type AxisConversionStyle end
144165struct SingleRange <: AxisConversionStyle end
@@ -172,9 +193,10 @@ const ArrayInitializer = Union{UndefInitializer,Missing,Nothing}
172193struct OffsetArray{T,N,AA<: AbstractArray } <: AbstractArray{T,N}
173194 parent:: AA
174195 offsets:: NTuple{N,Int}
175- function OffsetArray {T,N,AA} (parent:: AA , offsets:: NTuple{N,Int} ) where {T,N,AA <: AbstractArray }
176- @boundscheck overflow_check .(axes (parent), offsets)
177- new {T,N,AA} (parent, offsets)
196+ @inline function OffsetArray {T, N, AA} (parent:: AA , offsets:: NTuple{N, Int} ; checkoverflow = true ) where {T, N, AA<: AbstractArray{T,N} }
197+ # allocation of `map` on tuple is optimized away
198+ checkoverflow && map (overflow_check, axes (parent), offsets)
199+ new {T, N, AA} (parent, offsets)
178200 end
179201end
180202
@@ -196,54 +218,106 @@ end
196218
197219# Tuples of integers are treated as offsets
198220# Empty Tuples are handled here
199- function OffsetArray (A:: AbstractArray , offsets:: Tuple{Vararg{Integer}} )
221+ @inline function OffsetArray (A:: AbstractArray , offsets:: Tuple{Vararg{Integer}} ; kw ... )
200222 _checkindices (A, offsets, " offsets" )
201- OffsetArray {eltype(A),ndims(A),typeof(A)} (A, offsets)
223+ OffsetArray {eltype(A), ndims(A), typeof(A)} (A, offsets; kw ... )
202224end
203225
204226# These methods are necessary to disallow incompatible dimensions for
205227# the OffsetVector and the OffsetMatrix constructors
206228for (FT, ND) in ((:OffsetVector , :1 ), (:OffsetMatrix , :2 ))
207- @eval function $FT (A:: AbstractArray{<:Any,$ND} , offsets:: Tuple{Vararg{Integer}} )
229+ @eval @inline function $FT (A:: AbstractArray{<:Any,$ND} , offsets:: Tuple{Vararg{Integer}} ; kw ... )
208230 _checkindices (A, offsets, " offsets" )
209- OffsetArray {eltype(A),$ND,typeof(A)} (A, offsets)
231+ OffsetArray {eltype(A), $ND, typeof(A)} (A, offsets; kw ... )
210232 end
211233 FTstr = string (FT)
212- @eval function $FT (A:: AbstractArray , offsets:: Tuple{Vararg{Integer}} )
213- throw (ArgumentError ($ FTstr * " requires a " * string ($ ND) * " D array" ))
234+ @eval @inline function $FT (A:: AbstractArray , offsets:: Tuple{Vararg{Integer}} ; kw ... )
235+ throw (ArgumentError ($ FTstr* " requires a " * string ($ ND)* " D array" ))
214236 end
215237end
216238
217239# # OffsetArray constructors
218240for FT in (:OffsetArray , :OffsetVector , :OffsetMatrix )
219241 # Nested OffsetArrays may strip off the wrapper and collate the offsets
220- @eval function $FT (A:: OffsetArray , offsets:: Tuple{Vararg{Integer}} )
242+ # empty tuples are handled here
243+ @eval @inline function $FT (A:: OffsetArray , offsets:: Tuple{Vararg{Int}} ; checkoverflow = true )
221244 _checkindices (A, offsets, " offsets" )
222- $ FT (parent (A), map (+ , A. offsets, offsets))
245+ # ensure that the offsets may be added together without an overflow
246+ checkoverflow && map (overflow_check, axes (A), offsets)
247+ I = map (+ , _offsets (A, parent (A)), offsets)
248+ $ FT (parent (A), I, checkoverflow = false )
249+ end
250+ @eval @inline function $FT (A:: OffsetArray , offsets:: Tuple{Integer,Vararg{Integer}} ; kw... )
251+ $ FT (A, map (Int, offsets); kw... )
223252 end
224253
225254 # In general, indices get converted to AbstractUnitRanges.
226255 # CartesianIndices{N} get converted to N ranges
227- @eval function $FT (A:: AbstractArray , inds:: Tuple{Any,Vararg{Any}} )
228- $ FT (A, _toAbstractUnitRanges (to_indices (A, axes (A), inds)))
256+ @eval @inline function $FT (A:: AbstractArray , inds:: Tuple{Any,Vararg{Any}} ; kw ... )
257+ $ FT (A, _toAbstractUnitRanges (to_indices (A, axes (A), inds)); kw ... )
229258 end
230259
231260 # convert ranges to offsets
232- @eval function $FT (A:: AbstractArray , inds:: Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}} )
261+ @eval @inline function $FT (A:: AbstractArray , inds:: Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}} ; kw ... )
233262 _checkindices (A, inds, " indices" )
234263 # Performance gain by wrapping the error in a function: see https://github.com/JuliaLang/julia/issues/37558
235264 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" ))
236265 lA = size (A)
237266 lI = map (length, inds)
238267 lA == lI || throw_dimerr (lA, lI)
239- $ FT (A, map (_offset, axes (A), inds))
268+ $ FT (A, map (_offset, axes (A), inds); kw ... )
240269 end
241270
242- @eval $ FT (A:: AbstractArray , inds:: Vararg ) = $ FT (A, inds)
271+ @eval @inline $ FT (A:: AbstractArray , inds... ; kw... ) = $ FT (A, inds; kw... )
272+ @eval @inline $ FT (A:: AbstractArray ; checkoverflow = false ) = $ FT (A, ntuple (zero, Val (ndims (A))), checkoverflow = checkoverflow)
243273
244- @eval $ FT (A:: AbstractArray , origin:: Origin ) = $ FT (A, origin (A) )
274+ @eval @inline $ FT (A:: AbstractArray , origin:: Origin ; checkoverflow = true ) = $ FT (A, origin. index .- first .( axes (A)); checkoverflow = checkoverflow )
245275end
246276
277+ # conversion-related methods
278+ @inline OffsetArray {T} (M:: AbstractArray , I... ; kw... ) where {T} = OffsetArray {T,ndims(M)} (M, I... ; kw... )
279+
280+ @inline function OffsetArray {T,N} (M:: AbstractArray{<:Any,N} , I... ; kw... ) where {T,N}
281+ M2 = _of_eltype (T, M)
282+ OffsetArray {T,N} (M2, I... ; kw... )
283+ end
284+ @inline OffsetArray {T,N} (M:: OffsetArray{T,N} , I... ; kw... ) where {T,N} = OffsetArray (M, I... ; kw... )
285+ @inline OffsetArray {T,N} (M:: AbstractArray{T,N} , I... ; kw... ) where {T,N} = OffsetArray {T,N,typeof(M)} (M, I... ; kw... )
286+
287+ @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... )
288+ @inline function OffsetArray {T,N,A} (M:: AbstractArray{<:Any,N} , I:: NTuple{N,Int} ; checkoverflow = true ) where {T,N,A<: AbstractArray{T,N} }
289+ checkoverflow && map (overflow_check, axes (M), I)
290+ Mv = no_offset_view (M)
291+ MvA = convert (A, Mv):: A
292+ Iof = map (+ , _offsets (M), I)
293+ OffsetArray {T,N,A} (MvA, Iof, checkoverflow = false )
294+ end
295+ @inline function OffsetArray {T, N, AA} (parent:: AbstractArray{<:Any,N} , offsets:: NTuple{N, Integer} ; kw... ) where {T, N, AA<: AbstractArray{T,N} }
296+ OffsetArray {T, N, AA} (parent, map (Int, offsets):: NTuple{N,Int} ; kw... )
297+ end
298+ @inline function OffsetArray {T,N,A} (M:: AbstractArray{<:Any,N} , I:: Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}} ; kw... ) where {T,N,A<: AbstractArray{T,N} }
299+ _checkindices (M, I, " indices" )
300+ # Performance gain by wrapping the error in a function: see https://github.com/JuliaLang/julia/issues/37558
301+ 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" ))
302+ lM = size (M)
303+ lI = map (length, I)
304+ lM == lI || throw_dimerr (lM, lI)
305+ OffsetArray {T,N,A} (M, map (_offset, axes (M), I); kw... )
306+ end
307+ @inline function OffsetArray {T,N,A} (M:: AbstractArray{<:Any,N} , I:: Tuple ; kw... ) where {T,N,A<: AbstractArray{T,N} }
308+ OffsetArray {T,N,A} (M, _toAbstractUnitRanges (to_indices (M, axes (M), I)); kw... )
309+ end
310+ @inline function OffsetArray {T,N,A} (M:: AbstractArray{<:Any,N} ; kw... ) where {T,N,A<: AbstractArray{T,N} }
311+ Mv = no_offset_view (M)
312+ MvA = convert (A, Mv):: A
313+ OffsetArray {T,N,A} (MvA, _offsets (M); kw... )
314+ end
315+ @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)
316+
317+ Base. convert (:: Type{T} , M:: AbstractArray ) where {T<: OffsetArray } = M isa T ? M : T (M)
318+
319+ @inline AbstractArray {T,N} (M:: OffsetArray{S,N} ) where {T,S,N} = OffsetArray {T} (M)
320+
247321# array initialization
248322function OffsetArray {T,N} (init:: ArrayInitializer , inds:: Tuple{Vararg{OffsetAxisKnownLength}} ) where {T,N}
249323 _checkindices (N, inds, " indices" )
@@ -277,6 +351,14 @@ Base.eachindex(::IndexLinear, A::OffsetVector) = axes(A, 1)
277351@inline Base. axes (A:: OffsetArray , d) = d <= ndims (A) ? IdOffsetRange (axes (parent (A), d), A. offsets[d]) : IdOffsetRange (axes (parent (A), d))
278352@inline Base. axes1 (A:: OffsetArray{T,0} ) where {T} = IdOffsetRange (axes (parent (A), 1 )) # we only need to specialize this one
279353
354+ # Utils to translate a function to the parent while preserving offsets
355+ unwrap (x) = x, identity
356+ unwrap (x:: OffsetArray ) = parent (x), data -> OffsetArray (data, x. offsets, checkoverflow = false )
357+ function parent_call (f, x)
358+ parent, wrap_offset = unwrap (x)
359+ wrap_offset (f (parent))
360+ end
361+
280362Base. similar (A:: OffsetArray , :: Type{T} , dims:: Dims ) where T =
281363 similar (parent (A), T, dims)
282364function Base. similar (A:: AbstractArray , :: Type{T} , inds:: Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}} ) where T
@@ -329,6 +411,8 @@ Base.falses(inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
329411# and one obtains the result below.
330412parentindex (r:: IdOffsetRange , i) = i - r. offset
331413
414+ @propagate_inbounds Base. getindex (A:: OffsetArray{<:Any,0} ) = A. parent[]
415+
332416@inline function Base. getindex (A:: OffsetArray{T,N} , I:: Vararg{Int,N} ) where {T,N}
333417 @boundscheck checkbounds (A, I... )
334418 J = map (parentindex, axes (A), I)
341425end
342426@propagate_inbounds Base. getindex (A:: OffsetArray , i:: Int ) = parent (A)[i]
343427
428+ @propagate_inbounds Base. getindex (A:: OffsetArray{<:Any,N} , c:: Vararg{Colon,N} ) where N =
429+ parent_call (x -> getindex (x, c... ), A)
430+
431+ # With one Colon we use linear indexing.
432+ # In this case we may forward the index to the parent, as the information about the axes is lost
433+ # The exception to this is with OffsetVectors where the axis information is preserved,
434+ # but that case is handled by getindex(::OffsetArray{<:Any,N}, ::Vararg{Colon,N})
435+ @propagate_inbounds Base. getindex (A:: OffsetArray , c:: Colon ) = A. parent[:]
436+
437+ const OffsetRange{T} = OffsetVector{T,<: AbstractRange{T} }
438+ const OffsetUnitRange{T} = OffsetVector{T,<: AbstractUnitRange{T} }
439+
440+ # An OffsetUnitRange might use the rapid getindex(::Array, ::AbstractUnitRange{Int}) for contiguous indexing
441+ @propagate_inbounds function Base. getindex (A:: Array , r:: OffsetUnitRange{Int} )
442+ B = A[parent (r)]
443+ OffsetArray (B, axes (r), checkoverflow = false )
444+ end
445+
446+ # Linear Indexing of OffsetArrays with AbstractUnitRanges may use the faster contiguous indexing methods
447+ @inline function Base. getindex (A:: OffsetArray , r:: AbstractUnitRange{Int} )
448+ @boundscheck checkbounds (A, r)
449+ # nD OffsetArrays do not have their linear indices shifted, so we may forward the indices provided to the parent
450+ @inbounds B = parent (A)[r]
451+ _indexedby (B, axes (r))
452+ end
453+ @inline function Base. getindex (A:: OffsetVector , r:: AbstractUnitRange{Int} )
454+ @boundscheck checkbounds (A, r)
455+ # OffsetVectors may have their linear indices shifted, so we subtract the offset from the indices provided
456+ @inbounds B = parent (A)[_subtractoffset (r, A. offsets[1 ])]
457+ _indexedby (B, axes (r))
458+ end
459+
460+ # This method added mainly to index an OffsetRange with another range
461+ @inline function Base. getindex (A:: OffsetVector , r:: AbstractRange{Int} )
462+ @boundscheck checkbounds (A, r)
463+ @inbounds B = parent (A)[_subtractoffset (r, A. offsets[1 ])]
464+ _indexedby (B, axes (r))
465+ end
466+
467+
344468@inline function Base. setindex! (A:: OffsetArray{T,N} , val, I:: Vararg{Int,N} ) where {T,N}
345469 @boundscheck checkbounds (A, I... )
346470 J = @inbounds map (parentindex, axes (A), I)
@@ -363,36 +487,37 @@ Base.dataids(A::OffsetArray) = Base.dataids(parent(A))
363487Broadcast. broadcast_unalias (dest:: OffsetArray , src:: OffsetArray ) = parent (dest) === parent (src) ? src : Broadcast. unalias (dest, src)
364488
365489# ## Special handling for AbstractRange
366-
367- const OffsetRange{T} = OffsetArray{T,1 ,<: AbstractRange{T} }
368490const IIUR = IdentityUnitRange{S} where S<: AbstractUnitRange{T} where T<: Integer
369491
370492Base. step (a:: OffsetRange ) = step (parent (a))
371493
372494@propagate_inbounds Base. getindex (a:: OffsetRange , r:: OffsetRange ) = OffsetArray (a[parent (r)], r. offsets)
373- @propagate_inbounds function Base. getindex (a:: OffsetRange , r:: IdOffsetRange )
374- OffsetArray (a. parent[r. parent .+ (r. offset - a. offsets[1 ])], r. offset)
375- end
376- @propagate_inbounds Base. getindex (r:: OffsetRange , s:: IIUR ) =
377- OffsetArray (r[s. indices], s)
378- @propagate_inbounds Base. getindex (a:: OffsetRange , r:: AbstractRange ) = a. parent[r .- a. offsets[1 ]]
379- @propagate_inbounds Base. getindex (a:: AbstractRange , r:: OffsetRange ) = OffsetArray (a[parent (r)], r. offsets)
380-
381- @propagate_inbounds Base. getindex (r:: UnitRange , s:: IIUR ) =
382- OffsetArray (r[s. indices], s)
383495
384- @propagate_inbounds Base. getindex (r:: StepRange , s:: IIUR ) =
385- OffsetArray (r[s. indices], s)
496+ @inline function _boundscheck_index_retaining_axes (r, s)
497+ @boundscheck checkbounds (r, s)
498+ @inbounds pr = r[UnitRange (s)]
499+ _indexedby (pr, axes (s))
500+ end
501+ @inline _boundscheck_return (r, s) = (@boundscheck checkbounds (r, s); s)
386502
387- # this method is needed for ambiguity resolution
388- @propagate_inbounds Base. getindex (r:: StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision} , s:: IIUR ) where T =
389- OffsetArray (r[s. indices], s)
503+ for OR in [:IIUR , :IdOffsetRange ]
504+ for R in [:StepRange , :StepRangeLen , :LinRange , :UnitRange ]
505+ @eval @inline Base. getindex (r:: $R , s:: $OR ) = _boundscheck_index_retaining_axes (r, s)
506+ end
390507
391- @propagate_inbounds Base. getindex (r:: StepRangeLen{T} , s:: IIUR ) where {T} =
392- OffsetArray (r[s. indices], s)
508+ # this method is needed for ambiguity resolution
509+ @eval @inline function Base. getindex (r:: StepRangeLen{T,<:Base.TwicePrecision,<:Base.TwicePrecision} , s:: $OR ) where T
510+ _boundscheck_index_retaining_axes (r, s)
511+ end
512+ end
513+ Base. getindex (r:: Base.OneTo , s:: IdOffsetRange ) = _boundscheck_index_retaining_axes (r, s)
393514
394- @propagate_inbounds Base. getindex (r:: LinRange , s:: IIUR ) =
395- OffsetArray (r[s. indices], s)
515+ # These methods are added to avoid ambiguities with Base.
516+ # The ones involving Base types should be ported to Base and version-limited here
517+ @inline Base. getindex (r:: IdentityUnitRange , s:: IIUR ) = _boundscheck_return (r, s)
518+ @inline Base. getindex (r:: IdentityUnitRange , s:: IdOffsetRange ) = _boundscheck_return (r, s)
519+ @inline Base. getindex (r:: Base.Slice , s:: IIUR ) = _boundscheck_return (r, s)
520+ @inline Base. getindex (r:: Base.Slice , s:: IdOffsetRange ) = _boundscheck_return (r, s)
396521
397522function Base. show (io:: IO , r:: OffsetRange )
398523 show (io, r. parent)
0 commit comments