Safe functional indexing (#316)

This allows common comparison operators be parsed into indices (`<`, `<=`, `>`, `>=`). So instead of doing `x[3:end]` you can do `x[>=(3)]` (if `x` uses `ArrayInterface.getindex`. The main utility right now is that we can ensure that the indices are always inbounds using this syntax. I'm still working out the best way to communicate this info to `checkbounds` so that bounds checking is passed over (similar to how `:` becomes `Slice` which is known to be inbounds); but that will probably be more involved so I'm going to make that a separate effort.

Author: Tokazama

src/ArrayInterface.jl

@@ -36,6 +36,9 @@ import Compat
 
 n_of_x(::StaticInt{N}, x::X) where {N,X} = ntuple(Compat.Returns(x), Val{N}())
 
+_add1(@nospecialize x) = x + oneunit(x)
+_sub1(@nospecialize x) = x - oneunit(x)
+
 @generated function merge_tuple_type(::Type{X}, ::Type{Y}) where {X<:Tuple,Y<:Tuple}
     Tuple{X.parameters...,Y.parameters...}
 end

src/indexing.jl

@@ -23,7 +23,7 @@ function Base.last(x::AbstractVector, n::StaticInt)
 end
 
 """
-    to_indices(A, I::Tuple) -> Tuple
+    ArrayInterface.to_indices(A, I::Tuple) -> Tuple
 
 Converts the tuple of indexing arguments, `I`, into an appropriate form for indexing into `A`.
 Typically, each index should be an `Int`, `StaticInt`, a collection with values of `Int`, or a collection with values of `CartesianIndex`
@@ -183,16 +183,50 @@ end
 end
 
 """
-    to_index([::IndexStyle, ]axis, arg) -> index
+    ArrayInterface.to_index([::IndexStyle, ]axis, arg) -> index
 
-Convert the argument `arg` that was originally passed to `getindex` for the dimension
-corresponding to `axis` into a form for native indexing (`Int`, Vector{Int}, etc.). New
-axis types with unique behavior should use an `IndexStyle` trait:
+Convert the argument `arg` that was originally passed to `ArrayInterface.getindex` for the
+dimension corresponding to `axis` into a form for native indexing (`Int`, Vector{Int}, etc.).
 
+`ArrayInterface.to_index` supports passing a function as an index. This function-index is
+transformed into a proper index.
+
+```julia
+julia> using ArrayInterface, Static
+
+julia> ArrayInterface.to_index(static(1):static(10), 5)
+5
+
+julia> ArrayInterface.to_index(static(1):static(10), <(5))
+static(1):4
+
+julia> ArrayInterface.to_index(static(1):static(10), <=(5))
+static(1):5
+
+julia> ArrayInterface.to_index(static(1):static(10), >(5))
+6:static(10)
+
+julia> ArrayInterface.to_index(static(1):static(10), >=(5))
+5:static(10)
+
+```
+
+Use of a function-index helps ensure that indices are inbounds
+
+```julia
+julia> ArrayInterface.to_index(static(1):static(10), <(12))
+static(1):10
+
+julia> ArrayInterface.to_index(static(1):static(10), >(-1))
+1:static(10)
+```
+
+New axis types with unique behavior should use an `IndexStyle` trait:
 ```julia
 to_index(axis::MyAxisType, arg) = to_index(IndexStyle(axis), axis, arg)
 to_index(::MyIndexStyle, axis, arg) = ...
 ```
+
 """
 to_index(x, i::Slice) = i
 to_index(x, ::Colon) = indices(x)
@@ -207,6 +241,18 @@ to_index(x::LinearIndices, i::AbstractArray{Bool}) = LogicalIndex{Int}(i)
 @inline to_index(x, i::CartesianIndex) = Tuple(i)
 @inline to_index(x, i::NDIndex) = Tuple(i)
 @inline to_index(x, i::AbstractArray{<:AbstractCartesianIndex}) = i
+@inline function to_index(x, i::Base.Fix2{<:Union{typeof(<),typeof(isless)},<:Union{Base.BitInteger,StaticInt}})
+    offset1(x):min(_sub1(canonicalize(i.x)), static_lastindex(x))
+end
+@inline function to_index(x, i::Base.Fix2{typeof(<=),<:Union{Base.BitInteger,StaticInt}})
+    offset1(x):min(canonicalize(i.x), static_lastindex(x))
+end
+@inline function to_index(x, i::Base.Fix2{typeof(>=),<:Union{Base.BitInteger,StaticInt}})
+    max(canonicalize(i.x), offset1(x)):static_lastindex(x)
+end
+@inline function to_index(x, i::Base.Fix2{typeof(>),<:Union{Base.BitInteger,StaticInt}})
+    max(_add1(canonicalize(i.x)), offset1(x)):static_lastindex(x)
+end
 # integer indexing
 to_index(x, i::AbstractArray{<:Integer}) = i
 to_index(x, @nospecialize(i::StaticInt)) = i

test/indexing.jl

@@ -202,6 +202,8 @@ end
     @test @inferred(ArrayInterface.getindex(cartesian, vec(cartesian))) == vec(cartesian)
     @test @inferred(ArrayInterface.getindex(linear, 2:3)) === 2:3
     @test @inferred(ArrayInterface.getindex(linear, 3:-1:1)) === 3:-1:1
+    @test @inferred(ArrayInterface.getindex(linear, >(1), <(3))) == linear[(begin+1):end, 1:(end-1)]
+    @test @inferred(ArrayInterface.getindex(linear, >=(1), <=(3))) == linear[begin:end, 1:end]
     @test_throws BoundsError ArrayInterface.getindex(linear, 4:13)
 end