From b01c99f2c0d62bba7ad3696dbf347dd1c3a8af2b Mon Sep 17 00:00:00 2001 From: Michael Hendricks Date: Tue, 16 Dec 2025 12:50:08 -0700 Subject: [PATCH] feat: iterate SortedMap and SortedSet in reverse order These methods follow the other pattern used by the other rev_* methods in the standard library. --- immut/sorted_map/README.mbt.md | 24 ++++++ immut/sorted_map/pkg.generated.mbti | 2 + immut/sorted_map/utils.mbt | 60 +++++++++++++++ immut/sorted_map/utils_test.mbt | 115 ++++++++++++++++++++++++++++ immut/sorted_set/README.mbt.md | 12 +++ immut/sorted_set/generic.mbt | 30 ++++++++ immut/sorted_set/generic_test.mbt | 86 +++++++++++++++++++++ immut/sorted_set/pkg.generated.mbti | 1 + 8 files changed, 330 insertions(+) diff --git a/immut/sorted_map/README.mbt.md b/immut/sorted_map/README.mbt.md index a423cab07..84898bb19 100644 --- a/immut/sorted_map/README.mbt.md +++ b/immut/sorted_map/README.mbt.md @@ -181,3 +181,27 @@ test { assert_eq(keys.collect(), ["a", "b", "c"]) } ``` + +### Reverse Iteration + +Use `rev_keys()` to get all keys in descending order. + +```mbt check +///| +test { + let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) + let keys = map.rev_keys().collect() + assert_eq(keys, ["c", "b", "a"]) +} +``` + +Use `rev_values()` to get all values in descending order of their keys. + +```mbt check +///| +test { + let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) + let values = map.rev_values().collect() + assert_eq(values, [3, 2, 1]) +} +``` diff --git a/immut/sorted_map/pkg.generated.mbti b/immut/sorted_map/pkg.generated.mbti index 34ef3eaa7..529ee847e 100644 --- a/immut/sorted_map/pkg.generated.mbti +++ b/immut/sorted_map/pkg.generated.mbti @@ -61,6 +61,8 @@ pub fn[K, V] SortedMap::new() -> Self[K, V] pub fn[K : Compare, V] SortedMap::remove(Self[K, V], K) -> Self[K, V] #alias(foldr_with_key) pub fn[K, V, A] SortedMap::rev_fold(Self[K, V], (A, K, V) -> A, init~ : A) -> A +pub fn[K, V] SortedMap::rev_keys(Self[K, V]) -> Iter[K] +pub fn[K, V] SortedMap::rev_values(Self[K, V]) -> Iter[V] #as_free_fn pub fn[K, V] SortedMap::singleton(K, V) -> Self[K, V] pub fn[K, V] SortedMap::to_array(Self[K, V]) -> Array[(K, V)] diff --git a/immut/sorted_map/utils.mbt b/immut/sorted_map/utils.mbt index f9bf9466e..f3e39ff09 100644 --- a/immut/sorted_map/utils.mbt +++ b/immut/sorted_map/utils.mbt @@ -328,6 +328,66 @@ pub fn[K, V] SortedMap::values(self : SortedMap[K, V]) -> Iter[V] { self.iter().map(p => p.1) } +///| +/// Return all keys of the map in descending order. +/// +/// # Example +/// +/// ```mbt +/// let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) +/// let keys = map.rev_keys().collect() +/// assert_eq(keys, ["c", "b", "a"]) +/// ``` +pub fn[K, V] SortedMap::rev_keys(self : SortedMap[K, V]) -> Iter[K] { + Iter::new(yield_ => { + fn go(t) { + match t { + Empty => IterContinue + Tree(k, value=_, l, r, ..) => + if go(r) is IterEnd { + IterEnd + } else if yield_(k) is IterEnd { + IterEnd + } else { + go(l) + } + } + } + + go(self) + }) +} + +///| +/// Return all values of the map in descending order of their keys. +/// +/// # Example +/// +/// ```mbt +/// let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) +/// let values = map.rev_values().collect() +/// assert_eq(values, [3, 2, 1]) +/// ``` +pub fn[K, V] SortedMap::rev_values(self : SortedMap[K, V]) -> Iter[V] { + Iter::new(yield_ => { + fn go(t) { + match t { + Empty => IterContinue + Tree(_k, value~, l, r, ..) => + if go(r) is IterEnd { + IterEnd + } else if yield_(value) is IterEnd { + IterEnd + } else { + go(l) + } + } + } + + go(self) + }) +} + ///| pub fn[K : Show, V : ToJson] SortedMap::to_json(self : SortedMap[K, V]) -> Json { ToJson::to_json(self) diff --git a/immut/sorted_map/utils_test.mbt b/immut/sorted_map/utils_test.mbt index d991d0af9..531cbe353 100644 --- a/immut/sorted_map/utils_test.mbt +++ b/immut/sorted_map/utils_test.mbt @@ -196,3 +196,118 @@ test "from_json" { assert_eq(xs, @json.from_json(xs.to_json())) } } + +///| +test "rev_keys returns keys in descending order" { + let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) + let keys = map.rev_keys().collect() + inspect(keys, content="[\"c\", \"b\", \"a\"]") +} + +///| +test "rev_values returns values in descending key order" { + let map = @sorted_map.from_array([("a", 1), ("b", 2), ("c", 3)]) + let values = map.rev_values().collect() + inspect(values, content="[3, 2, 1]") +} + +///| +test "rev_keys on empty map" { + let map : @sorted_map.SortedMap[String, Int] = @sorted_map.new() + inspect(map.rev_keys().collect(), content="[]") +} + +///| +test "rev_values on empty map" { + let map : @sorted_map.SortedMap[String, Int] = @sorted_map.new() + inspect(map.rev_values().collect(), content="[]") +} + +///| +test "rev_keys on single element map" { + let map = @sorted_map.singleton("a", 1) + let keys = map.rev_keys().collect() + inspect(keys, content="[\"a\"]") +} + +///| +test "rev_values on single element map" { + let map = @sorted_map.singleton("a", 1) + let values = map.rev_values().collect() + inspect(values, content="[1]") +} + +///| +test "rev_keys early termination" { + let map = @sorted_map.from_array([(1, "one"), (2, "two"), (3, "three")]) + let result = map.rev_keys().take(2).collect() + inspect(result, content="[3, 2]") +} + +///| +test "rev_values early termination" { + let map = @sorted_map.from_array([(1, "one"), (2, "two"), (3, "three")]) + let result = map.rev_values().take(2).collect() + inspect(result, content="[\"three\", \"two\"]") +} + +///| +test "rev_keys is reverse of keys" { + let map = @sorted_map.from_array([ + (3, "three"), + (8, "eight"), + (1, "one"), + (2, "two"), + (0, "zero"), + ]) + let forward = map.keys_as_iter().collect() + let reverse = map.rev_keys().collect() + assert_eq(reverse, forward.rev()) +} + +///| +test "rev_values is reverse of values" { + let map = @sorted_map.from_array([ + (3, "three"), + (8, "eight"), + (1, "one"), + (2, "two"), + (0, "zero"), + ]) + let forward = map.values().collect() + let reverse = map.rev_values().collect() + assert_eq(reverse, forward.rev()) +} + +///| +test "rev_keys with complex tree" { + let map = @sorted_map.from_array([ + (15, "15"), + (10, "10"), + (20, "20"), + (5, "5"), + (12, "12"), + (18, "18"), + (25, "25"), + ]) + let keys = map.rev_keys().collect() + inspect(keys, content="[25, 20, 18, 15, 12, 10, 5]") +} + +///| +test "rev_values with complex tree" { + let map = @sorted_map.from_array([ + (15, "15"), + (10, "10"), + (20, "20"), + (5, "5"), + (12, "12"), + (18, "18"), + (25, "25"), + ]) + let values = map.rev_values().collect() + inspect( + values, + content="[\"25\", \"20\", \"18\", \"15\", \"12\", \"10\", \"5\"]", + ) +} diff --git a/immut/sorted_set/README.mbt.md b/immut/sorted_set/README.mbt.md index b019a6483..ddeba6a23 100644 --- a/immut/sorted_set/README.mbt.md +++ b/immut/sorted_set/README.mbt.md @@ -170,6 +170,18 @@ test { } ``` +You can also use `rev_iter()` to iterate in descending order. + +```mbt check +///| +test { + let set = @sorted_set.from_array([1, 2, 3, 4, 5]) + let arr = [] + set.rev_iter().each(v => arr.push(v)) + assert_eq(arr, [5, 4, 3, 2, 1]) +} +``` + ## All & Any `all` and `any` can detect whether all elements in the set match or if there are elements that match. diff --git a/immut/sorted_set/generic.mbt b/immut/sorted_set/generic.mbt index 7326385d5..c301a12e6 100644 --- a/immut/sorted_set/generic.mbt +++ b/immut/sorted_set/generic.mbt @@ -37,6 +37,36 @@ pub fn[A] SortedSet::iter(self : SortedSet[A]) -> Iter[A] { }) } +///| +/// Iterate over the elements in the set in descending order. +/// +/// # Example +/// +/// ```mbt +/// let set = @sorted_set.from_array([1, 2, 3, 4, 5]) +/// let result = set.rev_iter().collect() +/// assert_eq(result, [5, 4, 3, 2, 1]) +/// ``` +pub fn[A] SortedSet::rev_iter(self : SortedSet[A]) -> Iter[A] { + Iter::new(yield_ => { + fn go(t) { + match t { + Empty => IterContinue + Node(left~, right~, value~, ..) => + if go(right) is IterEnd { + IterEnd + } else if yield_(value) is IterEnd { + IterEnd + } else { + go(left) + } + } + } + + go(self) + }) +} + ///| pub fn[A] SortedSet::iterator(self : SortedSet[A]) -> Iterator[A] { let mut curr_node = self diff --git a/immut/sorted_set/generic_test.mbt b/immut/sorted_set/generic_test.mbt index da2eb3653..f52b6ffbd 100644 --- a/immut/sorted_set/generic_test.mbt +++ b/immut/sorted_set/generic_test.mbt @@ -78,3 +78,89 @@ test "iterator" { let result = set.iterator().collect() inspect(result, content="[1, 2, 3]") } + +///| +test "rev_iter returns elements in descending order" { + let set = @sorted_set.from_array([1, 2, 3, 4, 5]) + let result = set.rev_iter().collect() + inspect(result, content="[5, 4, 3, 2, 1]") +} + +///| +test "rev_iter on empty set" { + let set : @sorted_set.SortedSet[Int] = @sorted_set.new() + let result = set.rev_iter().collect() + inspect(result, content="[]") +} + +///| +test "rev_iter on single element set" { + let set = @sorted_set.singleton(42) + let result = set.rev_iter().collect() + inspect(result, content="[42]") +} + +///| +test "rev_iter early termination" { + let set = @sorted_set.from_array([1, 2, 3, 4, 5]) + let result = set.rev_iter().take(3).collect() + inspect(result, content="[5, 4, 3]") +} + +///| +test "rev_iter on unbalanced insertion" { + let set = @sorted_set.from_array([7, 2, 9, 4, 5, 6, 3, 8, 1]) + let result = set.rev_iter().collect() + inspect(result, content="[9, 8, 7, 6, 5, 4, 3, 2, 1]") +} + +///| +test "rev_iter is reverse of iter" { + let set = @sorted_set.from_array([3, 1, 4, 1, 5, 9, 2, 6]) + let forward = set.iter().collect() + let reverse = set.rev_iter().collect() + assert_eq(reverse, forward.rev()) +} + +///| +test "rev_iter with early termination on right subtree" { + // Create a binary search tree with both left and right subtrees + let set = @sorted_set.new() + .add(2) // root + .add(1) // left child + .add(3) // right child + + // Create an iterator that stops after processing the first element (from right) + let mut count = 0 + let iter = set.rev_iter() + let _ = iter.run((_ : Int) => { + count = count + 1 + if count == 1 { + IterEnd + } else { + IterContinue + } + }) + + // Verify that we stopped after processing the first element + inspect(count, content="1") +} + +///| +test "rev_iter with complex tree structure" { + let set = @sorted_set.from_array([ + 15, 10, 20, 5, 12, 18, 25, 3, 7, 11, 14, 17, 19, 23, 27, + ]) + let result = set.rev_iter().collect() + inspect( + result, + content="[27, 25, 23, 20, 19, 18, 17, 15, 14, 12, 11, 10, 7, 5, 3]", + ) +} + +///| +test "rev_iter terminates on left" { + let set = @sorted_set.from_array([1, 2, 3]) + let result = set.rev_iter().take(2).to_array() + inspect(result, content="[3, 2]") +} diff --git a/immut/sorted_set/pkg.generated.mbti b/immut/sorted_set/pkg.generated.mbti index eb7263aca..5f097e99e 100644 --- a/immut/sorted_set/pkg.generated.mbti +++ b/immut/sorted_set/pkg.generated.mbti @@ -50,6 +50,7 @@ pub fn[A] SortedSet::min_option(Self[A]) -> A? pub fn[A] SortedSet::new() -> Self[A] pub fn[A : Compare] SortedSet::remove(Self[A], A) -> Self[A] pub fn[A] SortedSet::remove_min(Self[A]) -> Self[A] +pub fn[A] SortedSet::rev_iter(Self[A]) -> Iter[A] #as_free_fn pub fn[A] SortedSet::singleton(A) -> Self[A] pub fn[A : Compare] SortedSet::split(Self[A], A) -> (Self[A], Bool, Self[A])