perf(HAMT): flatten nested Branch structure

Author: FlyCloudC

immut/hashset/HAMT.mbt

@@ -38,6 +38,7 @@ pub fn[A : Eq + Hash] contains(self : T[A], key : A) -> Bool {
   guard self._ is Some(node) else { return false }
   loop (node, @path.of(key)) {
     (Leaf(key1, bucket), _) => key == key1 || bucket.contains(key)
+    (Flat(key1, path1), path) => path == path1 && key == key1
     (Branch(children), path) => {
       let idx = path.idx()
       if children.has(idx) {
@@ -50,13 +51,24 @@ pub fn[A : Eq + Hash] contains(self : T[A], key : A) -> Bool {
 }
 
 ///|
-fn[A] make_leaf(key : A, path : Path) -> Node[A] {
-  // make sure leaf nodes always appear at the bottom of the tree
-  if path.is_end() {
-    Leaf(key, @list.empty())
+/// require: key1 != key2, path1 and path2 has the same length
+fn[A] join_2(key1 : A, path1 : Path, key2 : A, path2 : Path) -> Node[A] {
+  let idx1 = path1.idx()
+  let idx2 = path2.idx()
+  if idx1 == idx2 {
+    let node = if path1.is_last() {
+      Leaf(key2, @list.singleton(key1))
+    } else {
+      join_2(key1, path1.next(), key2, path2.next())
+    }
+    Branch(@sparse_array.singleton(idx1, node))
   } else {
-    let child = make_leaf(key, path.next())
-    Branch(@sparse_array.singleton(path.idx(), child))
+    let (node1, node2) = if path1.is_last() {
+      (Leaf(key1, @list.empty()), Leaf(key2, @list.empty()))
+    } else {
+      (Flat(key1, path1.next()), Flat(key2, path2.next()))
+    }
+    Branch(@sparse_array.doubleton(idx1, node1, idx2, node2))
   }
 }
 
@@ -69,14 +81,20 @@ fn[A : Eq] add_with_path(self : Node[A], key : A, path : Path) -> Node[A] {
       } else {
         Leaf(key, bucket.add(key1))
       }
+    Flat(key1, path1) =>
+      if path == path1 && key == key1 {
+        self
+      } else {
+        join_2(key1, path1, key, path)
+      }
     Branch(children) => {
       let idx = path.idx()
       if children.has(idx) {
         let child = children[idx]
         let child = child.add_with_path(key, path.next())
         Branch(children.replace(idx, child))
       } else {
-        let child = make_leaf(key, path.next())
+        let child = Flat(key, path.next())
         Branch(children.add(idx, child))
       }
     }
@@ -87,7 +105,7 @@ fn[A : Eq] add_with_path(self : Node[A], key : A, path : Path) -> Node[A] {
 /// Add a key to the hashset.
 pub fn[A : Eq + Hash] add(self : T[A], key : A) -> T[A] {
   match self._ {
-    None => Some(make_leaf(key, @path.of(key)))
+    None => Some(Flat(key, @path.of(key)))
     Some(node) => Some(node.add_with_path(key, @path.of(key)))
   }
 }
@@ -115,15 +133,26 @@ fn[A : Eq] remove_with_path(self : Node[A], key : A, path : Path) -> Node[A]? {
       } else {
         Some(self)
       }
+    Flat(key1, path1) =>
+      if path == path1 && key == key1 {
+        None
+      } else {
+        Some(self)
+      }
     Branch(children) => {
       let idx = path.idx()
       if children.has(idx) {
         let child = children[idx]
         let new_child = child.remove_with_path(key, path.next())
-        match (children.size(), new_child) {
-          (1, None) => None
-          (_, None) => Some(Branch(children.remove(idx)))
-          (_, Some(new_child)) => Some(Branch(children.replace(idx, new_child)))
+        let new_children = match (children.size(), new_child) {
+          (1, None) => return None
+          (_, None) => children.remove(idx)
+          (_, Some(new_child)) => children.replace(idx, new_child)
+        }
+        match new_children.data {
+          [Flat(key1, path1)] =>
+            Some(Flat(key1, path1.push(new_children.elem_info.fitst_idx())))
+          _ => Some(Branch(new_children))
         }
       } else {
         Some(self)
@@ -140,6 +169,7 @@ pub fn[A] size(self : T[A]) -> Int {
   fn node_size(node) {
     match node {
       Leaf(_, bucket) => 1 + bucket.length()
+      Flat(_) => 1
       Branch(children) =>
         for i = 0, total_size = 0 {
           if i < children.data.length() {
@@ -194,6 +224,7 @@ pub fn[A] each(self : T[A], f : (A) -> Unit) -> Unit {
         f(k)
         bucket.each(f)
       }
+      Flat(k, _) => f(k)
       Branch(children) => children.each(go)
     }
   }
@@ -210,6 +241,7 @@ pub fn[A] iter(self : T[A]) -> Iter[A] {
   fn go(node) -> Iter[A] {
     match node {
       Leaf(k, bucket) => Iter::singleton(k) + bucket.iter()
+      Flat(k, _) => Iter::singleton(k)
       Branch(children) => children.data.iter().flat_map(go)
     }
   }
@@ -262,6 +294,7 @@ impl[A : Eq] Eq for Node[A] with op_equal(self, other) {
   match (self, other) {
     (Leaf(x, xs), Leaf(y, ys)) =>
       xs.length() == ys.length() && xs.add(x).iter().all(ys.add(y).contains(_))
+    (Flat(x, pathx), Flat(y, pathy)) => pathx == pathy && x == y
     (Branch(xs), Branch(ys)) => xs == ys
     _ => false
   }

immut/hashset/types.mbt

@@ -18,7 +18,9 @@ typealias @path.Path
 ///|
 /// An non-empty immutable hash set data structure
 priv enum Node[A] {
+  Flat(A, Path)
   Leaf(A, @list.T[A]) // use a list of buckets to resolve collision
+  /// number of all its leaf > 1. If equals 1, it should be represented as `Flat`
   Branch(@sparse_array.SparseArray[Node[A]])
 }
 

immut/internal/path/path.mbt

@@ -16,7 +16,7 @@
 /// A Path is a binary string divided into several segments, for example:
 /// 0b11_10000_10101_00100_11111_01010
 /// where the last segment is 10000, and the first segment is 01010.
-pub(all) type Path UInt
+pub(all) type Path UInt derive(Eq)
 
 ///|
 const SEGMENT_LENGTH : Int = 5
@@ -36,13 +36,13 @@ pub fn[A : Hash] of(key : A) -> Path {
 }
 
 ///|
-/// If SEGMENT_LENGTH == 5, END == 0b11
-const END : UInt = 0xffffffffU >> (SEGMENT_LENGTH * SEGMENT_NUM)
+/// If SEGMENT_LENGTH == 5, MAX_TAIL == 0b11_11111
+const MAX_TAIL : UInt = 0xffffffffU >> (SEGMENT_LENGTH * (SEGMENT_NUM - 1))
 
 ///|
-pub fn Path::is_end(self : Path) -> Bool {
+pub fn Path::is_last(self : Path) -> Bool {
   let Path(self) = self
-  self == END
+  self <= MAX_TAIL
 }
 
 ///|

immut/internal/path/path.mbti

@@ -6,9 +6,10 @@ fn[A : Hash] of(A) -> Path
 // Types and methods
 pub(all) type Path UInt
 fn Path::idx(Self) -> Int
-fn Path::is_end(Self) -> Bool
+fn Path::is_last(Self) -> Bool
 fn Path::next(Self) -> Self
 fn Path::push(Self, Int) -> Self
+impl Eq for Path
 
 // Type aliases
 

immut/internal/sparse_array/bitset.mbt

@@ -31,6 +31,11 @@ pub fn Bitset::index_of(self : Bitset, idx : Int) -> Int {
   (self._ & ((1U << idx) - 1)).popcnt()
 }
 
+///|
+pub fn Bitset::fitst_idx(self : Bitset) -> Int {
+  self._.ctz()
+}
+
 ///|
 pub fn Bitset::union(self : Bitset, other : Bitset) -> Bitset {
   Bitset(self._ | other._)

immut/internal/sparse_array/sparse_array.mbt

@@ -31,6 +31,23 @@ pub fn[X] singleton(idx : Int, value : X) -> SparseArray[X] {
   { elem_info: empty_bitset.add(idx), data: [value] }
 }
 
+///|
+pub fn[X] doubleton(
+  idx1 : Int,
+  value1 : X,
+  idx2 : Int,
+  value2 : X
+) -> SparseArray[X] {
+  {
+    elem_info: empty_bitset.add(idx1).add(idx2),
+    data: if idx1 < idx2 {
+      [value1, value2]
+    } else {
+      [value2, value1]
+    },
+  }
+}
+
 ///|
 pub fn[X] has(self : SparseArray[X], idx : Int) -> Bool {
   self.elem_info.has(idx)

immut/internal/sparse_array/sparse_array.mbti

@@ -3,6 +3,8 @@ package "moonbitlang/core/immut/internal/sparse_array"
 // Values
 fn[X] add(SparseArray[X], Int, X) -> SparseArray[X]
 
+fn[X] doubleton(Int, X, Int, X) -> SparseArray[X]
+
 fn[X] each(SparseArray[X], (X) -> Unit) -> Unit
 
 fn[X] empty() -> SparseArray[X]
@@ -23,6 +25,7 @@ fn[X] size(SparseArray[X]) -> Int
 pub(all) type Bitset UInt
 fn Bitset::add(Self, Int) -> Self
 fn Bitset::difference(Self, Self) -> Self
+fn Bitset::fitst_idx(Self) -> Int
 fn Bitset::has(Self, Int) -> Bool
 fn Bitset::index_of(Self, Int) -> Int
 fn Bitset::intersection(Self, Self) -> Self