Add tests and fix bugs for indexing

* Fixed a typo in flatten_args
* Fixed bug that missed boundschecks on multidim logical indexing
* to_axes wasn't accounting for linear indexing on multidim arrays

Author: Tokazama

src/indexing.jl

@@ -131,7 +131,7 @@ end
 end
 @inline function flatten_args(A, args::Tuple{Arg,Vararg{Any}}) where {N,Arg<:AbstractArray{Bool,N}}
     if length(args) === 1
-        if s isa IndexLinear
+        if IndexStyle(A) isa IndexLinear
             return (LogicalIndex{Int}(first(args)),)
         else
             return (LogicalIndex(first(args)),)
@@ -218,8 +218,13 @@ to_indices(A, axs::Tuple{}, args::Tuple{}) = ()
 
 
 _multi_check_index(axs::Tuple, arg) = _multi_check_index(axs, axes(arg))
+_multi_check_index(axs::Tuple, arg::LogicalIndex) = prod(map(length, axs)) == length(arg)
 function _multi_check_index(axs::Tuple, arg::AbstractArray{T}) where {T<:CartesianIndex}
-    return checkindex(Bool, axs, arg)
+    b = true
+    for i in arg
+        b &= Base.checkbounds_indices(Bool, axs, (i,))
+    end
+    return b
 end
 _multi_check_index(::Tuple{}, ::Tuple{}) = true
 function _multi_check_index(axs::Tuple, args::Tuple)
@@ -331,6 +336,8 @@ pair of axes and indices calling [`to_axis`](@ref).
 @inline function to_axes(A, inds::Tuple)
     if ndims(A) === 1
         return (to_axis(axes(A, 1), first(inds)),)
+    elseif is_linear_indexing(A, inds)
+        return (to_axis(eachindex(IndexLinear(), A), first(inds)),)
     else
         return to_axes(A, axes(A), inds)
     end

test/indexing.jl

@@ -50,9 +50,27 @@ end
     @test @inferred(ArrayInterface.to_indices(a, (1, 1, 1))) == (1,1, 1)
     @test @inferred ArrayInterface.to_indices(a, ([CartesianIndex(1,1,1), CartesianIndex(1,2,1)],)) == (CartesianIndex{3}[CartesianIndex(1, 1, 1), CartesianIndex(1, 2, 1)],)
     @test @inferred ArrayInterface.to_indices(a, ([CartesianIndex(1,1), CartesianIndex(1,2)],1:1)) == (CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(1, 2)], 1:1)
+    @test @inferred(first(ArrayInterface.to_indices(a, (fill(true, 2, 2, 1),)))) isa Base.LogicalIndex
+
+    @test_throws BoundsError ArrayInterface.to_indices(a, (fill(true, 2, 2, 2),))
     @test_throws ErrorException ArrayInterface.to_indices(ones(2,2,2), (1, 1))
 end
 
+@testset "to_axes" begin
+    A = ones(3, 3)
+    axis = StaticInt(1):StaticInt(3)
+    inds = StaticInt(1):StaticInt(2)
+    multi_inds = [CartesianIndex(1, 1), CartesianIndex(1, 2)]
+
+    @test @inferred(ArrayInterface.to_axes(A, (axis, axis), (inds, inds))) === (inds, inds)
+    # vector indexing
+    @test @inferred(ArrayInterface.to_axes(ones(3), (axis,), (inds,))) === (inds,)
+    # linear indexing
+    @test @inferred(ArrayInterface.to_axes(A, (axis, axis), (inds,))) === (inds,)
+    # multidim arg
+    @test @inferred(ArrayInterface.to_axes(A, (axis, axis), (multi_inds,))) === (Base.OneTo(2),)
+end
+
 @testset "0-dimensional" begin
     x = Array{Int,0}(undef)
     ArrayInterface.setindex!(x, 1)