Iterables and iterators :D

5 files changed, 144 insertions, 5 deletions

diff --git a/tests/features/test_drop.zc b/tests/features/test_drop.zc
new file mode 100644
index 0000000..8b34efe
--- /dev/null
+++ b/tests/features/test_drop.zc
@@ -0,0 +1,26 @@
+import "../../std/mem.zc"
+
+var DROP_CALLED = 0;
+
+struct MyResource {
+    id: int;
+}
+
+impl Drop for MyResource {
+    fn drop(self) {
+        println "Dropping MyResource {self.id}";
+        DROP_CALLED = 1;
+    }
+}
+
+test "drop_trait" {
+    {
+        var res = MyResource { id: 1 };
+        // Scope ends here, drop should be called
+    }
+    
+    if (DROP_CALLED != 1) {
+        println "Error: Drop was not called!";
+        exit(1);
+    }
+}
diff --git a/tests/features/test_iterator.zc b/tests/features/test_iterator.zc
new file mode 100644
index 0000000..1340a00
--- /dev/null
+++ b/tests/features/test_iterator.zc
@@ -0,0 +1,45 @@
+import "../../std/iter.zc"
+import "../../std/option.zc"
+
+struct RangeIter {
+    current: int;
+    stop: int;
+}
+
+impl RangeIter {
+    fn next(self) -> Option<int> {
+        if (self.current < self.stop) {
+            var v = self.current;
+            self.current += 1;
+            return Option<int>::Some(v);
+        }
+        return Option<int>::None();
+    }
+}
+
+struct MyRange {
+    start: int;
+    end: int;
+}
+
+impl MyRange {
+    fn iterator(self) -> RangeIter {
+        return RangeIter {
+            current: self.start,
+            stop: self.end
+        };
+    }
+}
+
+test "iterator_desugaring" {
+    var range = MyRange { start: 0, end: 5 };
+    var sum = 0;
+
+    // This loop should be desugared by the compiler
+    for i in range {
+        println "Got: {i}"; // Verification print
+        sum += i;
+    }
+
+    assert(sum == 10); // 0+1+2+3+4 = 10
+}
diff --git a/tests/features/test_iterator_drop.zc b/tests/features/test_iterator_drop.zc
new file mode 100644
index 0000000..43e22a7
--- /dev/null
+++ b/tests/features/test_iterator_drop.zc
@@ -0,0 +1,51 @@
+import "../../std/iter.zc"
+import "../../std/mem.zc"
+
+var DROP_CALLED = 0;
+
+struct DropIter {
+    count: int;
+}
+
+impl DropIter {
+    fn next(self) -> Option<int> {
+        if self.count > 0 {
+            self.count -= 1;
+            return Option<int>::Some(self.count);
+        }
+        return Option<int>::None();
+    }
+}
+
+impl Drop for DropIter {
+    fn drop(self) {
+        DROP_CALLED = 1;
+        println "Iterator dropped";
+    }
+}
+
+struct DropRange {
+    len: int;
+}
+
+impl DropRange {
+    fn iterator(self) -> DropIter {
+        return DropIter { count: self.len };
+    }
+}
+
+test "iterator_drop" {
+    {
+        var range = DropRange { len: 3 };
+        for i in range {
+            // Loop runs 3 times
+            var _ignore = i;
+        }
+        // Iterator should be dropped here
+    }
+
+    if (DROP_CALLED != 1) {
+        println "Error: Iterator was not dropped! (DROP_CALLED={DROP_CALLED})";
+        exit(1);
+    }
+}
diff --git a/tests/features/test_vec_iter.zc b/tests/features/test_vec_iter.zc
new file mode 100644
index 0000000..b7e8dcb
--- /dev/null
+++ b/tests/features/test_vec_iter.zc
@@ -0,0 +1,19 @@
+import "../../std/vec.zc"
+
+test "vec_int_iteration" {
+    var v = Vec<int>::new();
+    v.push(10);
+    v.push(20);
+    v.push(30);
+
+    var sum = 0;
+    for x in v {
+        sum = sum + x;
+    }
+    
+    if (sum != 60) {
+        println "Expected 60, got {sum}";
+        exit(1);
+    }
+    v.free();
+}
diff --git a/tests/memory/test_memory_safety.zc b/tests/memory/test_memory_safety.zc
index 9f9b649..652fe21 100644
--- a/tests/memory/test_memory_safety.zc
+++ b/tests/memory/test_memory_safety.zc
@@ -1,14 +1,12 @@
 
 import "std/mem.zc"
 
-// --- Globals ---
+// ** Globals **
 var DROP_COUNT = 0;
 
-// --- Structs & Helpers ---
+// ** Structs & Helpers **
 
-trait Drop {
-    fn drop(self);
-}
+// trait Drop defined in std/mem.zc
 
 struct Resource {
     id: int;