From 1269d891d27e55d0e71777d3d3a22985f3ddc42e Mon Sep 17 00:00:00 2001
From: Matt Fishman <mtfishman@users.noreply.github.com>
Date: Sat, 25 May 2024 17:17:19 -0400
Subject: [PATCH] [GradedAxes] Fix iterate (#1468)

* [GradedAxes] Fix iterate

* Bump to v0.3.9
---
 NDTensors/Project.toml                        |  2 +-
 .../src/lib/GradedAxes/src/gradedunitrange.jl |  7 ++++++
 .../src/lib/GradedAxes/test/test_basics.jl    | 25 ++++++++++++++++++-
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml
index 80d7a38a09..b1a56aa106 100644
--- a/NDTensors/Project.toml
+++ b/NDTensors/Project.toml
@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org>"]
-version = "0.3.8"
+version = "0.3.9"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
diff --git a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
index c10d5052e3..8a19b6e04f 100644
--- a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
+++ b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
@@ -161,6 +161,13 @@ function Base.first(a::GradedUnitRange)
   return labelled(first(unlabel_blocks(a)), label(a[Block(1)]))
 end
 
+Base.iterate(a::GradedUnitRange) = isempty(a) ? nothing : (first(a), first(a))
+function Base.iterate(a::GradedUnitRange, i)
+  i == last(a) && return nothing
+  next = a[i + step(a)]
+  return (next, next)
+end
+
 function firstblockindices(a::GradedUnitRange)
   return labelled.(firstblockindices(unlabel_blocks(a)), blocklabels(a))
 end
diff --git a/NDTensors/src/lib/GradedAxes/test/test_basics.jl b/NDTensors/src/lib/GradedAxes/test/test_basics.jl
index aa3a9d7560..41815db4d1 100644
--- a/NDTensors/src/lib/GradedAxes/test/test_basics.jl
+++ b/NDTensors/src/lib/GradedAxes/test/test_basics.jl
@@ -9,7 +9,7 @@ using BlockArrays:
   blocklengths,
   blocks
 using NDTensors.GradedAxes: GradedUnitRange, blocklabels, gradedrange
-using NDTensors.LabelledNumbers: LabelledUnitRange, label, labelled, unlabel
+using NDTensors.LabelledNumbers: LabelledUnitRange, islabelled, label, labelled, unlabel
 using Test: @test, @test_broken, @testset
 @testset "GradedAxes basics" begin
   for a in (
@@ -18,7 +18,30 @@ using Test: @test, @test_broken, @testset
     gradedrange(["x" => 2, "y" => 3]),
   )
     @test a isa GradedUnitRange
+    for x in iterate(a)
+      @test x == 1
+      @test label(x) == "x"
+    end
+    for x in iterate(a, labelled(1, "x"))
+      @test x == 2
+      @test label(x) == "x"
+    end
+    for x in iterate(a, labelled(2, "x"))
+      @test x == 3
+      @test label(x) == "y"
+    end
+    for x in iterate(a, labelled(3, "y"))
+      @test x == 4
+      @test label(x) == "y"
+    end
+    for x in iterate(a, labelled(4, "y"))
+      @test x == 5
+      @test label(x) == "y"
+    end
+    @test isnothing(iterate(a, labelled(5, "y")))
     @test length(a) == 5
+    @test step(a) == 1
+    @test !islabelled(step(a))
     @test length(blocks(a)) == 2
     @test blocks(a)[1] == 1:2
     @test label(blocks(a)[1]) == "x"