From 0a4e9dde95b6ea2283bf9219bec96572ac16460a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=8C=97=E9=9C=B2?=
 <69190413+illusory0x0@users.noreply.github.com>
Date: Mon, 13 Jan 2025 15:15:03 +0800
Subject: [PATCH] add compare for sorted_set

---
 sorted_set/set.mbt         | 35 ++++++++++++++++++++++++++++++++++-
 sorted_set/set_test.mbt    | 17 +++++++++++++++++
 sorted_set/sorted_set.mbti |  1 +
 3 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/sorted_set/set.mbt b/sorted_set/set.mbt
index 3caabf9eb..e12c35599 100644
--- a/sorted_set/set.mbt
+++ b/sorted_set/set.mbt
@@ -299,9 +299,42 @@ pub fn disjoint[V : Compare](self : T[V], src : T[V]) -> Bool {
 
 // General collection operations
 
+///|
+pub impl[V : Compare] Compare for T[V] with compare(self : T[V], other : T[V]) -> Int {
+  let rhs = other.to_array()
+  let mut i = 0
+  let mut result = 0
+  fn dfs(root : Node[V]?) -> Unit {
+    match root {
+      Some(root) => {
+        if i >= rhs.length() {
+          return
+        }
+        dfs(root.left)
+        let cmp = root.value.compare(rhs[i])
+        if cmp != 0 {
+          result = cmp
+          return
+        }
+        i = i + 1
+        dfs(root.right)
+      }
+      None => ()
+    }
+    ()
+  }
+
+  dfs(self.root)
+  if result != 0 {
+    result
+  } else {
+    (self.size - other.size).to_int()
+  }
+}
+
 ///|
 pub fn op_equal[V : Compare](self : T[V], other : T[V]) -> Bool {
-  self.to_array() == other.to_array()
+  self.compare(other) == 0
 }
 
 ///|
diff --git a/sorted_set/set_test.mbt b/sorted_set/set_test.mbt
index 422a2e1a9..7d6699a7c 100644
--- a/sorted_set/set_test.mbt
+++ b/sorted_set/set_test.mbt
@@ -228,3 +228,20 @@ test "range" {
   let result = set.range(2, 4)
   inspect!(result, content="[2, 3, 4]")
 }
+
+test "compare" {
+  let xs = [1, 2, 3]
+  let ys = [1, 2]
+  let xs = @sorted_set.of(xs)
+  let ys = @sorted_set.of(ys)
+  inspect!(xs < ys, content="false")
+  inspect!(xs > ys, content="true")
+  inspect!(xs == ys, content="false")
+  let xs = [2, 7, 4]
+  let ys = [1]
+  let xs = @sorted_set.of(xs)
+  let ys = @sorted_set.of(ys)
+  inspect!(xs < ys, content="false")
+  inspect!(xs > ys, content="true")
+  inspect!(xs == ys, content="false")
+}
diff --git a/sorted_set/sorted_set.mbti b/sorted_set/sorted_set.mbti
index 731ffc22b..2329f6987 100644
--- a/sorted_set/sorted_set.mbti
+++ b/sorted_set/sorted_set.mbti
@@ -25,6 +25,7 @@ impl T {
   intersection[V : Compare](Self[V], Self[V]) -> Self[V]
   is_empty[V : Compare](Self[V]) -> Bool
   iter[V](Self[V]) -> Iter[V]
+  compare[V : Compare](Self[V], Self[V]) -> Int
   op_equal[V : Compare](Self[V], Self[V]) -> Bool
   range[V : Compare](Self[V], V, V) -> Iter[V]
   remove[V : Compare](Self[V], V) -> Unit