Improve LazyAxis test coverage (#311)

* Improve LazyAxis test coverage

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.15"
+version = "6.0.16"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

src/axes.jl

@@ -191,9 +191,7 @@ struct LazyAxis{N,P} <: AbstractUnitRange{Int}
 end
 
 @inline Base.parent(x::LazyAxis{N,P}) where {N,P} = axes(getfield(x, :parent), static(N))
-@inline function Base.parent(x::LazyAxis{:,P}) where {P}
-    return eachindex(IndexLinear(), getfield(x, :parent))
-end
+@inline Base.parent(x::LazyAxis{:,P}) where {P} = eachindex(IndexLinear(), getfield(x, :parent))
 
 @inline parent_type(::Type{LazyAxis{N,P}}) where {N,P} = axes_types(P, static(N))
 # TODO this approach to parent_type(::Type{LazyAxis{:}}) is a bit hacky. Something like
@@ -229,7 +227,8 @@ end
 ArrayInterfaceCore.known_last(::Type{LazyAxis{N,P}}) where {N,P} = known_last(axes_types(P, static(N)))
 ArrayInterfaceCore.known_last(::Type{LazyAxis{:,P}}) where {P} = known_length(P)
 Base.last(x::LazyAxis) = _last(known_last(x), x)
-_last(::Nothing, x) = last(parent(x))
+_last(::Nothing, x::LazyAxis{:}) = lastindex(getfield(x, :parent))
+_last(::Nothing, x::LazyAxis{N}) where {N} = lastindex(getfield(x, :parent), N)
 _last(N::Int, x) = N
 
 known_length(::Type{<:LazyAxis{:,P}}) where {P} = known_length(P)
@@ -242,7 +241,8 @@ Base.axes1(x::LazyAxis) = x
 Base.axes(x::Slice{<:LazyAxis}) = (Base.axes1(x),)
 # assuming that lazy loaded params like dynamic length from `size(::Array, dim)` are going
 # be used again later with `Slice{LazyAxis}`, we quickly load indices
-Base.axes1(x::Slice{<:LazyAxis}) = indices(parent(x.indices))
+Base.axes1(x::Slice{LazyAxis{N,A}}) where {N,A} = indices(getfield(x.indices, :parent), StaticInt(N))
+Base.axes1(x::Slice{LazyAxis{:,A}}) where {A} = indices(getfield(x.indices, :parent))
 Base.to_shape(x::LazyAxis) = length(x)
 
 @propagate_inbounds function Base.getindex(x::LazyAxis, i::CanonicalInt)

test/axes.jl

@@ -11,20 +11,27 @@ m = Array{Float64}(undef, 4, 3)
 @testset "LazyAxis" begin
     A = zeros(3,4,5);
     SA = MArray(zeros(3,4,5))
+    DA = MArray(zeros(3,4,5), LinearIndices((Base.Slice(1:3), 1:4, 1:5)))
     lz1 = ArrayInterface.LazyAxis{1}(A)
     slz1 = ArrayInterface.LazyAxis{1}(SA)
+    dlz1 = ArrayInterface.LazyAxis{1}(DA)
     lzc = ArrayInterface.LazyAxis{:}(A)
     slzc = ArrayInterface.LazyAxis{:}(SA)
-
-    @test @inferred(first(lz1)) === @inferred(first(slz1))
-    @test @inferred(first(lzc)) === @inferred(first(slzc))
-    @test @inferred(last(lz1)) === @inferred(last(slz1))
-    @test @inferred(last(lzc)) === @inferred(last(slzc))
-    @test @inferred(length(lz1)) === @inferred(length(slz1))
-    @test @inferred(length(lzc)) === @inferred(length(slzc))
-    @test @inferred(Base.to_shape(lzc)) == length(slzc)
+    dlzc = ArrayInterface.LazyAxis{:}(DA)
+
+    @test @inferred(first(lz1)) === @inferred(first(slz1)) === @inferred(first(dlz1))
+    @test @inferred(first(lzc)) === @inferred(first(slzc)) === @inferred(first(dlzc))
+    @test @inferred(last(lz1)) === @inferred(last(slz1)) === @inferred(last(dlz1))
+    @test @inferred(last(lzc)) === @inferred(last(slzc)) === @inferred(last(dlzc))
+    @test @inferred(length(lz1)) === @inferred(length(slz1)) === @inferred(length(dlz1))
+    @test @inferred(length(lzc)) === @inferred(length(slzc)) === @inferred(length(dlzc))
+    @test @inferred(Base.to_shape(lzc)) == length(slzc) == length(dlzc)
     @test @inferred(Base.checkindex(Bool, lzc, 1)) & @inferred(Base.checkindex(Bool, slzc, 1))
     @test axes(lzc)[1] == Base.axes1(lzc) == axes(Base.Slice(lzc))[1] == Base.axes1(Base.Slice(lzc))
+    @test keys(axes(A, 1)) == @inferred(keys(lz1))
+
+    @test @inferred(ArrayInterface.known_first(slzc)) === 1
+    @test @inferred(ArrayInterface.known_length(slz1)) === 3
 
     @test @inferred(getindex(lz1, 2)) == 2
     @test @inferred(getindex(lz1, 1:2)) == 1:2
@@ -33,6 +40,8 @@ m = Array{Float64}(undef, 4, 3)
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{:}(A))) <: Base.OneTo{Int}
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{4}(SA))) <: ArrayInterface.SOneTo{1}
     @test @inferred(ArrayInterface.parent_type(ArrayInterface.LazyAxis{:}(SA))) <: ArrayInterface.SOneTo{60}
+    @test @inferred(IndexStyle(SA)) isa IndexLinear
+    @test @inferred(IndexStyle(DA)) isa IndexLinear
     @test ArrayInterface.can_change_size(ArrayInterface.LazyAxis{1,Vector{Any}})
 
     Aperm = PermutedDimsArray(A, (3,1,2))