adapt iter API for set

moonbitlang · Jan 12, 2025 · 147e136 · 147e136
1 parent ffb701d
commit 147e136
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 deletions.
diff --git a/sorted_map/map.mbt b/sorted_map/map.mbt
@@ -258,18 +258,19 @@ fn range_aux[K : Compare, V](
       Some(node) => {
         let cmp_key_low = node.key.compare(low)
         let cmp_key_high = node.key.compare(high)
-
         if cmp_key_low > 0 {
           // we should go left
           guard let IterContinue = range_aux(node.left, low, high).run(yield_) else {
             x => return x
           }
+
         }
         if cmp_key_low >= 0 && cmp_key_high <= 0 {
           // we should yield this node
           guard let IterContinue = yield_(node.key, node.value) else {
             x => return x
           }
+
         }
         if cmp_key_high < 0 {
           guard let IterContinue = range_aux(node.right, low, high).run(yield_) else {

diff --git a/sorted_set/set.mbt b/sorted_set/set.mbt
@@ -429,10 +429,38 @@ fn to_array[T](self : Node[T]) -> Array[T] {
 }
 
 ///|
-pub fn range[V : Compare](self : T[V], low : V, high : V) -> Array[V] {
-  let result = []
-  self.each(fn(x) { if x >= low && x <= high { result.push(x) } })
-  result
+pub fn range[V : Compare](self : T[V], low : V, high : V) -> Iter[V] {
+  range_aux(self.root, low, high)
+}
+
+///|
+fn range_aux[V : Compare](root : Node[V]?, low : V, high : V) -> Iter[V] {
+  Iter::new(fn(yield_) {
+    match root {
+      None => IterContinue
+      Some(node) => {
+        let cmp_key_low = node.value.compare(low)
+        let cmp_key_high = node.value.compare(high)
+        if cmp_key_low > 0 {
+          guard let IterContinue = range_aux(node.left, low, high).run(yield_) else {
+            x => return x
+          }
+
+        }
+        if cmp_key_low >= 0 && cmp_key_high <= 0 {
+          guard let IterContinue = yield_(node.value) else { x => return x }
+
+        }
+        if cmp_key_high < 0 {
+          guard let IterContinue = range_aux(node.right, low, high).run(yield_) else {
+            x => return x
+          }
+
+        }
+        IterContinue
+      }
+    }
+  })
 }
 
 // AVL tree operations

diff --git a/sorted_set/sorted_set.mbti b/sorted_set/sorted_set.mbti
@@ -26,7 +26,7 @@ impl T {
   is_empty[V : Compare](Self[V]) -> Bool
   iter[V](Self[V]) -> Iter[V]
   op_equal[V : Compare](Self[V], Self[V]) -> Bool
-  range[V : Compare](Self[V], V, V) -> Array[V]
+  range[V : Compare](Self[V], V, V) -> Iter[V]
   remove[V : Compare](Self[V], V) -> Unit
   size[V : Compare](Self[V]) -> Int64
   subset[V : Compare](Self[V], Self[V]) -> Bool