Merge pull request #108 from Tokazama/indexing-tests

Indexing tests

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.14.14"
+version = "2.14.15"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/indexing.jl

@@ -322,40 +322,25 @@ end
 Reconstruct `A` given the values in `data`. New methods using `unsafe_reconstruct`
 should only dispatch on `A`.
 """
-function unsafe_reconstruct(A::OneTo, data; kwargs...)
-    if can_change_size(A)
-        return typeof(A)(data)
+function unsafe_reconstruct(axis::OneTo, data; kwargs...)
+    if axis === data
+        return axis
     else
-        if data isa Slice ||
-           !(known_length(A) === nothing || known_length(A) !== known_length(data))
-            return A
-        else
-            return OneTo(data)
-        end
+        return OneTo(data)
     end
 end
-function unsafe_reconstruct(A::UnitRange, data; kwargs...)
-    if can_change_size(A)
-        return typeof(A)(data)
+function unsafe_reconstruct(axis::UnitRange, data; kwargs...)
+    if axis === data
+        return axis
     else
-        if data isa Slice ||
-           !(known_length(A) === nothing || known_length(A) !== known_length(data))
-            return A
-        else
-            return UnitRange(data)
-        end
+        return UnitRange(first(data), last(data))
     end
 end
-function unsafe_reconstruct(A::OptionallyStaticUnitRange, data; kwargs...)
-    if can_change_size(A)
-        return typeof(A)(data)
+function unsafe_reconstruct(axis::OptionallyStaticUnitRange, data; kwargs...)
+    if axis === data
+        return axis
     else
-        if data isa Slice ||
-           !(known_length(A) === nothing || known_length(A) !== known_length(data))
-            return A
-        else
-            return OptionallyStaticUnitRange(data)
-        end
+        return OptionallyStaticUnitRange(static_first(data), static_last(data))
     end
 end
 function unsafe_reconstruct(A::AbstractUnitRange, data; kwargs...)
@@ -420,6 +405,15 @@ previously executed `to_index(old_axis, arg) -> index`. `to_axis` assumes that
         return to_axis(IndexStyle(axis), axis, inds)
     end
 end
+
+# don't need to check size b/c slice means it's the entire axis
+@inline function to_axis(axis, inds::Slice)
+    if can_change_size(axis)
+        return copy(axis)
+    else
+        return axis
+    end
+end
 @inline function to_axis(S::IndexStyle, axis, inds)
     return unsafe_reconstruct(axis, StaticInt(1):static_length(inds))
 end

test/indexing.jl

@@ -8,7 +8,7 @@
 end
 
 @testset "UnsafeIndex" begin
-    @test @inferred(ArrayInterface.UnsafeIndex(ones(2,2,2), typeof((1,[1,2],1)))) == ArrayInterface.UnsafeGetCollection() 
+    @test @inferred(ArrayInterface.UnsafeIndex(ones(2,2,2), typeof((1,[1,2],1)))) == ArrayInterface.UnsafeGetCollection()
     @test @inferred(ArrayInterface.UnsafeIndex(ones(2,2,2), typeof((1,1,1)))) == ArrayInterface.UnsafeGetElement() 
 end
 
@@ -19,10 +19,24 @@ end
     @test @inferred(ArrayInterface.to_index(axis, [1, 2])) == [1, 2]
     @test @inferred(ArrayInterface.to_index(axis, [true, false, false])) == [1]
 
-    @test_throws BoundsError  ArrayInterface.to_index(axis, 4)
-    @test_throws BoundsError  ArrayInterface.to_index(axis, 1:4)
+    @test_throws BoundsError ArrayInterface.to_index(axis, 4)
+    @test_throws BoundsError ArrayInterface.to_index(axis, 1:4)
     @test_throws BoundsError ArrayInterface.to_index(axis, [1, 2, 5])
-    @test_throws BoundsError  ArrayInterface.to_index(axis, [true, false, false, true])
+    @test_throws BoundsError ArrayInterface.to_index(axis, [true, false, false, true])
+end
+
+@testset "unsafe_reconstruct" begin
+    one_to = Base.OneTo(10)
+    opt_ur = StaticInt(1):10
+    ur = 1:10
+    @test @inferred(ArrayInterface.unsafe_reconstruct(one_to, opt_ur)) === one_to
+    @test @inferred(ArrayInterface.unsafe_reconstruct(one_to, one_to)) === one_to
+
+    @test @inferred(ArrayInterface.unsafe_reconstruct(opt_ur, opt_ur)) === opt_ur
+    @test @inferred(ArrayInterface.unsafe_reconstruct(opt_ur, one_to)) === opt_ur
+
+    @test @inferred(ArrayInterface.unsafe_reconstruct(ur, ur)) === ur
+    @test @inferred(ArrayInterface.unsafe_reconstruct(ur, one_to)) === ur
 end
 
 @testset "to_indices" begin
@@ -69,6 +83,9 @@ end
     @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),)
+
+    @test ArrayInterface.to_axis(axis, axis) === axis
+    @test ArrayInterface.to_axis(axis, ArrayInterface.indices(axis)) === axis
 end
 
 @testset "0-dimensional" begin
@@ -96,6 +113,12 @@ end
     # TODO should this be implemented in ArrayInterface with vectorization?
     #@test_throws ArgumentError Base._sub2ind((1:3,), 2)
     #@test_throws ArgumentError Base._ind2sub((1:3,), 2)
+    x = Array{Int,2}(undef, (2, 2))
+    ArrayInterface.unsafe_set_element!(x, 1, (2, 2))
+    @test ArrayInterface.unsafe_get_element(x, (2, 2)) === 1
+
+    @test_throws MethodError ArrayInterface.unsafe_set_element!(x, 1, (:x, :x))
+    @test_throws MethodError ArrayInterface.unsafe_get_element(x, (:x, :x))
 end
 
 @testset "2-dimensional" begin
@@ -154,3 +177,4 @@ end
         @test @inferred(ArrayInterface.getindex(LinearIndices(A),ArrayInterface.getindex(CartesianIndices(A),i))) == i
     end
 end
+