fix getindex and iterate

Author: ogauthe

NDTensors/src/lib/GradedAxes/src/labelledunitrangedual.jl

@@ -25,6 +25,15 @@ end
 
 # fix ambiguities
 Base.getindex(a::LabelledUnitRangeDual, i::Integer) = dual(nondual(a)[i])
+function Base.getindex(a::LabelledUnitRangeDual, indices::AbstractUnitRange{<:Integer})
+  return dual(nondual(a)[indices])
+end
+
+function Base.iterate(a::LabelledUnitRangeDual, i)
+  i == last(a) && return nothing
+  next = convert(eltype(a), labelled(i + step(a), label(a)))
+  return (next, next)
+end
 
 function Base.show(io::IO, ::MIME"text/plain", a::LabelledUnitRangeDual)
   println(io, typeof(a))

NDTensors/src/lib/GradedAxes/test/test_dual.jl

@@ -72,6 +72,11 @@ end
   @test space_isequal(la, la)
   @test label_type(la) == U1
 
+  @test iterate(la) == (1, 1)
+  @test iterate(la) == (1, 1)
+  @test iterate(la, 1) == (2, 2)
+  @test isnothing(iterate(la, 2))
+
   lad = dual(la)
   @test lad isa LabelledUnitRangeDual
   @test label(lad) == U1(-1)
@@ -87,10 +92,22 @@ end
   @test dual(lad) === la
   @test label_type(lad) == U1
 
+  @test iterate(lad) == (1, 1)
+  @test iterate(lad) == (1, 1)
+  @test iterate(lad, 1) == (2, 2)
+  @test isnothing(iterate(lad, 2))
+
+  lad2 = lad[1:1]
+  @test lad2 isa LabelledUnitRangeDual
+  @test label(lad2) == U1(-1)
+  @test unlabel(lad2) == 1:1
+
   laf = flip(la)
   @test laf isa LabelledUnitRangeDual
   @test label(laf) == U1(1)
   @test unlabel(laf) == 1:2
+  @test labelled_isequal(la, laf)
+  @test !space_isequal(la, laf)
 
   ladf = flip(dual(la))
   @test ladf isa LabelledUnitRange