Use to_indices in the OffsetArray constructor (#157)

* use to_indices to simplify constructor dispatch
to_indices converts colons to ranges. This is followed by converting CartesianIndices to ranges.

* Fix indexing IdOffsetRange using offset ranges, ensuring that indexing is idempotent

* Define `firstindex` and `lastindex` for `IdOffsetRange`

* Check for consistency of dimensions in the `OffsetVector` and `OffsetMatrix` constructors.

Author: jishnub

Project.toml

@@ -10,8 +10,9 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CatIndices = "aafaddc9-749c-510e-ac4f-586e18779b91"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra"]
+test = ["Aqua", "CatIndices", "DelimitedFiles", "Documenter", "Test", "LinearAlgebra", "EllipsisNotation"]

docs/src/index.md

@@ -56,7 +56,7 @@ Base.require_one_based_indexing(OA)
 ```
 
 [`OffsetArrays.Origin`](@ref) can be convenient if you want to directly specify the origin of the output
-OffsetArray, it will automatically compute the needed offsets. For example:
+OffsetArray, it will automatically compute the corresponding offsets. For example:
 
 ```@repl index
 OffsetArray(A, OffsetArrays.Origin(-1, -1))

docs/src/internals.md

@@ -143,3 +143,46 @@ OffsetArrays.IdOffsetRange(-3:3)
 julia> Ao[ax, 0][1] == Ao[ax[1], 0]
 true
 ```
+
+## Using custom axis types
+
+While a wide variety of `AbstractUnitRange`s provided by `Base` may be used as indices to construct an `OffsetArray`, at times it might be convenient to define custom types. The `OffsetArray` constructor accepts any type that may be converted to an `AbstractUnitRange`. This proceeds through a two-step process. Let's assume that the constructor called is `OffsetArray(A, indstup)`, where `indstup` is a `Tuple` of indices.
+
+1. At the first step, the constructor calls `to_indices(A, axes(A), indstup)` to lower `indstup` to a `Tuple` of `AbstractUnitRange`s. This step converts --- among other things --- `Colon`s to axis ranges. Custom types may extend `Base.to_indices(A, axes(A), indstup)` with the desired conversion of `indstup` to `Tuple{Vararg{AbstractUnitRange{Int}}}` if this is feasible.
+
+2. At the second step, the result obtained from the previous step treated again to convert it to a `Tuple` of `AbstractUnitRange`s to handle cases where the first step doesn't achieve this. An additional customization option may be specified at this stage: a type may be converted either to a single `AbstractUnitRange{Int}`, or to a `Tuple` of them. A type might specify which of these two behaviours is desired by extending [`OffsetArrays.AxisConversionStyle`](@ref). An example of a type that is acted upon at this stage is `CartesianIndices`, which is converted to a `Tuple` of `AbstractUnitRange`s.
+
+For example, here are a couple of custom type that facilitate zero-based indexing:
+
+```jldoctest; setup = :(using OffsetArrays)
+julia> struct ZeroBasedIndexing end
+
+julia> Base.to_indices(A, inds, ::Tuple{ZeroBasedIndexing}) = map(x -> 0:length(x)-1, inds)
+
+julia> a = zeros(3, 3);
+
+julia> oa = OffsetArray(a, ZeroBasedIndexing());
+
+julia> axes(oa)
+(OffsetArrays.IdOffsetRange(0:2), OffsetArrays.IdOffsetRange(0:2))
+```
+
+In this example we had to define the action of `to_indices` as the type `ZeroBasedIndexing` did not have a familiar hierarchy. Things are even simpler if we subtype `AbstractUnitRange`, in which case we need to define `first` and `length` for the custom range to be able to use it as an axis:
+
+```jldoctest; setup = :(using OffsetArrays)
+julia> struct ZeroTo <: AbstractUnitRange{Int}
+       n :: Int
+       ZeroTo(n) = new(n < 0 ? -1 : n)
+       end
+
+julia> Base.first(::ZeroTo) = 0
+
+julia> Base.length(r::ZeroTo) = r.n + 1
+
+julia> oa = OffsetArray(zeros(2,2), ZeroTo(1), ZeroTo(1));
+
+julia> axes(oa)
+(OffsetArrays.IdOffsetRange(0:1), OffsetArrays.IdOffsetRange(0:1))
+```
+
+Note that zero-based indexing may also be achieved using the pre-defined type [`OffsetArrays.Origin`](@ref).

docs/src/reference.md

@@ -7,4 +7,5 @@ OffsetMatrix
 OffsetArrays.Origin
 OffsetArrays.IdOffsetRange
 OffsetArrays.no_offset_view
+OffsetArrays.AxisConversionStyle
 ```

src/OffsetArrays.jl

@@ -1,6 +1,6 @@
 module OffsetArrays
 
-using Base: Indices, tail, @propagate_inbounds
+using Base: tail, @propagate_inbounds
 @static if !isdefined(Base, :IdentityUnitRange)
     const IdentityUnitRange = Base.Slice
 else
@@ -15,8 +15,8 @@ include("origin.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, IdOffsetRange}
-const OffsetAxis = Union{OffsetAxisKnownLength, CartesianIndices, Colon}
+const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange}
+const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
 const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
 
 ## OffsetArray
@@ -51,7 +51,7 @@ julia> OffsetArray(reshape(1:6, 2, 3), 0:1, -1:1)
  1  3  5
  2  4  6
 
-julia> OffsetArray(reshape(1:6, 2, 3), :, -1:1) # : as a placeholder means no offset is applied at this dimension
+julia> OffsetArray(reshape(1:6, 2, 3), :, -1:1) # : as a placeholder to indicate that no offset is to be applied to this dimension
 2×3 OffsetArray(reshape(::UnitRange{$Int}, 2, 3), 1:2, -1:1) with eltype $Int with indices 1:2×-1:1:
  1  3  5
  2  4  6
@@ -123,54 +123,73 @@ function overflow_check(r, offset::T) where T
         throw_lower_overflow_error()
     end
 end
-## OffsetArray constructors
 
+# Tuples of integers are treated as offsets
+# Empty Tuples are handled here
+function OffsetArray(A::AbstractArray, offsets::Tuple{Vararg{Integer}})
+    _checkindices(A, offsets, "offsets")
+    OffsetArray{eltype(A), ndims(A), typeof(A)}(A, offsets)
+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 function $FT(A::AbstractArray{<:Any,$ND}, offsets::Tuple{Vararg{Integer}})
+        _checkindices(A, offsets, "offsets")
+        OffsetArray{eltype(A), $ND, typeof(A)}(A, offsets)
+    end
+    FTstr = string(FT)
+    @eval function $FT(A::AbstractArray, offsets::Tuple{Vararg{Integer}})
+        throw(ArgumentError($FTstr*" requires a "*string($ND)*"D array"))
+    end
+end
+
+## OffsetArray constructors
 for FT in (:OffsetArray, :OffsetVector, :OffsetMatrix)
-    # The only route out to inner constructor
-    @eval function $FT(A::AbstractArray{T, N}, offsets::NTuple{N, Integer}) where {T, N}
-        ndims(A) == N || throw(DimensionMismatch("The number of offsets $(N) should equal ndims(A) = $(ndims(A))"))
-        OffsetArray{T, ndims(A), typeof(A)}(A, offsets)
+    # Nested OffsetArrays may strip off the wrapper and collate the offsets
+    @eval function $FT(A::OffsetArray, offsets::Tuple{Vararg{Integer}})
+        _checkindices(A, offsets, "offsets")
+        $FT(parent(A), map(+, A.offsets, offsets))
+    end
+
+    # In general, indices get converted to AbstractUnitRanges.
+    # CartesianIndices{N} get converted to N ranges
+    @eval function $FT(A::AbstractArray, inds::Tuple{Any,Vararg{Any}})
+        $FT(A, _toAbstractUnitRanges(to_indices(A, axes(A), inds)))
     end
-    # nested OffsetArrays
-    @eval $FT(A::OffsetArray{T, N}, offsets::NTuple{N, Integer}) where {T,N} = $FT(parent(A), A.offsets .+ offsets)
+
     # convert ranges to offsets
-    @eval function $FT(A::AbstractArray{T}, inds::NTuple{N,OffsetAxisKnownLength}) where {T,N}
-        axparent = axes(A)
-        lA = map(length, axparent)
+    @eval function $FT(A::AbstractArray, inds::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}})
+        _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(DimensionMismatch("supplied axes do not agree with the size of the array (got size $lA for the array and $lI for the indices"))
-        $FT(A, map(_offset, axparent, inds)) 
-    end
-    # lower CartesianIndices and Colon
-    @eval function $FT(A::AbstractArray{T}, inds::NTuple{N, OffsetAxis}) where {T, N}
-        indsN = _uncolonindices(A, _expandCartesianIndices(inds))
-        $FT(A, indsN)
+        lA == lI || throw_dimerr(lA, lI)
+        $FT(A, map(_offset, axes(A), inds)) 
     end
-    @eval $FT(A::AbstractArray{T}, inds::Vararg{OffsetAxis,N}) where {T, N} = $FT(A, inds)
 
-    @eval $FT(A::AbstractArray, origin::Origin) = OffsetArray(A, origin(A))
+    @eval $FT(A::AbstractArray, inds::Vararg) = $FT(A, inds)
+
+    @eval $FT(A::AbstractArray, origin::Origin) = $FT(A, origin(A))
 end
 
 # array initialization
-function OffsetArray{T,N}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N}
+function OffsetArray{T,N}(init::ArrayInitializer, inds::Tuple{Vararg{OffsetAxisKnownLength}}) where {T,N}
+    _checkindices(N, inds, "indices")
     AA = Array{T,N}(init, map(_indexlength, inds))
     OffsetArray{T, N, typeof(AA)}(AA, map(_indexoffset, inds))
 end
-function OffsetArray{T, N}(init::ArrayInitializer, inds::NTuple{NT, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N, NT}
-    # NT is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
-    indsN = _expandCartesianIndices(inds)
-    length(indsN) == N || throw(DimensionMismatch("The number of offsets $(length(indsN)) should equal ndims(A) = $N"))
-    OffsetArray{T, N}(init, indsN)
+function OffsetArray{T, N}(init::ArrayInitializer, inds::Tuple) where {T, N}
+    OffsetArray{T, N}(init, _toAbstractUnitRanges(inds))
 end
-OffsetArray{T,N}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T,N} = OffsetArray{T,N}(init, inds)
+OffsetArray{T,N}(init::ArrayInitializer, inds::Vararg) where {T,N} = OffsetArray{T,N}(init, inds)
 
 OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, OffsetAxisKnownLength}) where {T,N} = OffsetArray{T,N}(init, inds)
-function OffsetArray{T}(init::ArrayInitializer, inds::NTuple{N, Union{OffsetAxisKnownLength, CartesianIndices}}) where {T, N}
-    # N is probably not the actual dimension of the array; CartesianIndices might contain multiple dimensions
-    indsN = _expandCartesianIndices(inds)
-    OffsetArray{T, length(indsN)}(init, indsN)
+function OffsetArray{T}(init::ArrayInitializer, inds::Tuple) where {T}
+    OffsetArray{T}(init, _toAbstractUnitRanges(inds))
 end
-OffsetArray{T}(init::ArrayInitializer, inds::Union{OffsetAxisKnownLength, CartesianIndices}...) where {T} = OffsetArray{T}(init, inds)
+OffsetArray{T}(init::ArrayInitializer, inds::Vararg) where {T} = OffsetArray{T}(init, inds)
 
 Base.IndexStyle(::Type{OA}) where {OA<:OffsetArray} = IndexStyle(parenttype(OA))
 parenttype(::Type{OffsetArray{T,N,AA}}) where {T,N,AA} = AA

src/axes.jl

@@ -132,6 +132,9 @@ offset_coerce(::Type{I}, r::AbstractUnitRange) where I<:AbstractUnitRange{T} whe
 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 Base.$f(r::IdOffsetRange) = $f(r.parent) .+ r.offset
+end
 
 @inline function Base.iterate(r::IdOffsetRange)
     ret = iterate(r.parent)
@@ -151,9 +154,12 @@ end
 @propagate_inbounds function Base.getindex(r::IdOffsetRange, s::AbstractUnitRange{<:Integer})
     return r.parent[s .- r.offset] .+ r.offset
 end
-@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdOffsetRange)
+@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdentityUnitRange)
     return IdOffsetRange(r.parent[s .- r.offset], r.offset)
 end
+@propagate_inbounds function Base.getindex(r::IdOffsetRange, s::IdOffsetRange)
+    return IdOffsetRange(r.parent[s.parent .+ (s.offset - r.offset)] .+ (r.offset - s.offset), s.offset)
+end
 
 # offset-preserve broadcasting
 Broadcast.broadcasted(::Base.Broadcast.DefaultArrayStyle{1}, ::typeof(-), r::IdOffsetRange{T}, x::Integer) where T =

src/utils.jl

@@ -8,14 +8,66 @@ _indexlength(i::Integer) = i
 _indexlength(i::Colon) = Colon()
 
 _offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
-_offset(axparent::AbstractUnitRange, ax::CartesianIndices) = _offset(axparent, first(ax.indices))
 _offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
 
-_uncolonindices(A::AbstractArray{<:Any,N}, inds::NTuple{N,Any}) where {N} = _uncolonindices(axes(A), inds)
-_uncolonindices(ax::Tuple, inds::Tuple) = (first(inds), _uncolonindices(tail(ax), tail(inds))...)
-_uncolonindices(ax::Tuple, inds::Tuple{Colon, Vararg{Any}}) = (first(ax), _uncolonindices(tail(ax), tail(inds))...)
-_uncolonindices(::Tuple{}, ::Tuple{}) = ()
+"""
+    OffsetArrays.AxisConversionStyle(typeof(indices))
 
-_expandCartesianIndices(inds::Tuple{<:CartesianIndices, Vararg{Any}}) = (convert(Tuple{Vararg{AbstractUnitRange{Int}}}, inds[1])..., _expandCartesianIndices(Base.tail(inds))...)
-_expandCartesianIndices(inds::Tuple{Any,Vararg{Any}}) = (inds[1], _expandCartesianIndices(Base.tail(inds))...)
-_expandCartesianIndices(::Tuple{}) = ()
+`AxisConversionStyle` declares if `indices` should be converted to a single `AbstractUnitRange{Int}` 
+or to a `Tuple{Vararg{AbstractUnitRange{Int}}}` while flattening custom types into indices. 
+This method is called after `to_indices(A::Array, axes(A), indices)` to provide 
+further information in case `to_indices` does not return a `Tuple` of `AbstractUnitRange{Int}`.
+
+Custom index types should extend `AxisConversionStyle` and return either `OffsetArray.SingleRange()`, 
+which is the default, or `OffsetArray.TupleOfRanges()`. In the former case, the type `T` should 
+define `Base.convert(::Type{AbstractUnitRange{Int}}, ::T)`, whereas in the latter it should define 
+`Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, ::T)`. 
+
+An example of the latter is `CartesianIndices`, which is converted to a `Tuple` of 
+`AbstractUnitRange{Int}` while flattening the indices.
+
+# Example
+```jldoctest; setup=:(using OffsetArrays)
+julia> struct NTupleOfUnitRanges{N}
+           x ::NTuple{N, UnitRange{Int}}
+       end
+
+julia> Base.to_indices(A, inds, t::Tuple{NTupleOfUnitRanges{N}}) where {N} = t;
+
+julia> OffsetArrays.AxisConversionStyle(::Type{NTupleOfUnitRanges{N}}) where {N} = OffsetArrays.TupleOfRanges();
+
+julia> Base.convert(::Type{Tuple{Vararg{AbstractUnitRange{Int}}}}, t::NTupleOfUnitRanges) = t.x;
+
+julia> a = zeros(3, 3);
+
+julia> inds = NTupleOfUnitRanges((3:5, 2:4));
+
+julia> oa = OffsetArray(a, inds);
+
+julia> axes(oa, 1) == 3:5
+true
+
+julia> axes(oa, 2) == 2:4
+true
+```
+"""
+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