Интерфейс Iterable для MyTreeSet

lab2/src/main/java/ru/spbstu/telematics/java/App.java

@@ -8,5 +8,10 @@ public class App
         System.out.println(tree.add(10));
         System.out.println(tree.add(20));
         System.out.println(tree.add(10));
+
+        System.out.println("Содержимое дерева:");
+        for (Integer val: tree) {
+            System.out.println(val);
+        }
     }
 }

lab2/src/main/java/ru/spbstu/telematics/java/MyTreeSet.java

@@ -1,7 +1,10 @@
 package ru.spbstu.telematics.java;
 
+import java.util.Iterator;
+import java.util.NoSuchElementException;
+import java.util.Stack;
 
-public class MyTreeSet<E extends Comparable<E>> {
+public class MyTreeSet<E extends Comparable<E>> implements Iterable<E> {
     // Класс для представления узла дерева
     private class Node {
         E value;
@@ -119,4 +122,42 @@ public class MyTreeSet<E extends Comparable<E>> {
         }
         return current;
     }
+
+	@Override
+	public Iterator<E> iterator() {
+        return new MyTreeSetIterator();
+	}
+
+    private class MyTreeSetIterator implements Iterator<E> {
+        // На вершине стека всегда будет лежать наименьший элемент в дереве
+        Stack<Node> nodeStack = new Stack<>();
+
+        MyTreeSetIterator() {
+            pushAllLeftToStack(root);
+        }
+
+        void pushAllLeftToStack(Node current) {
+            if (current == null) return;
+            nodeStack.push(current);
+            pushAllLeftToStack(current.left);
+        }
+        
+		@Override
+		public boolean hasNext() {
+            return nodeStack.size() != 0;
+		}
+
+		@Override
+		public E next() {
+            if (!hasNext()) throw new NoSuchElementException();
+
+            Node minNode = nodeStack.pop();
+
+            if (minNode.right != null) {
+                pushAllLeftToStack(minNode.right);
+            }
+
+            return minNode.value;
+		}
+    }
 }

lab2/src/test/java/ru/spbstu/telematics/java/AppTest.java

@@ -4,6 +4,8 @@ import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import static org.junit.jupiter.api.Assertions.*;
 
+import java.util.Iterator;
+import java.util.NoSuchElementException;
 import java.util.TreeSet;
 
 class MyTreeSetTests {
@@ -71,4 +73,27 @@ class MyTreeSetTests {
         assertEquals(myTreeSet.remove(50), treeSet.remove(50));
         assertEquals(myTreeSet.size(), treeSet.size());
     }
+
+    @Test
+    void testIterator() {
+        treeSet.add(15);
+        myTreeSet.add(15);
+        treeSet.add(250);
+        myTreeSet.add(250);
+        treeSet.add(50);
+        myTreeSet.add(50);
+        Iterator<Integer> itMyTreeSet = myTreeSet.iterator();
+        Iterator<Integer> itTreeSet = treeSet.iterator();
+
+        assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext());
+        assertEquals(itMyTreeSet.next(), itTreeSet.next());
+        assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext());
+        assertEquals(itMyTreeSet.next(), itTreeSet.next());
+        assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext());
+        assertEquals(itMyTreeSet.next(), itTreeSet.next());
+        
+        assertEquals(itMyTreeSet.hasNext(), itTreeSet.hasNext());
+        assertThrows(NoSuchElementException.class, () -> itMyTreeSet.next());
+        assertThrows(NoSuchElementException.class, () -> itTreeSet.next());
+    }
 }