From 1bd31b67cd45c388aed2fea92ade567224410af9 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Wed, 9 Apr 2025 14:56:32 +0000
Subject: [PATCH] feat(core): improve priority queue performance (#11559)

---
 packages/common/infra/package.json            |   1 -
 packages/common/nbstore/package.json          |   1 -
 .../__tests__/binary-search-tree.spec.ts      | 233 ++++++++++++++
 .../nbstore/src/utils/binary-search-tree.ts   | 296 ++++++++++++++++++
 .../nbstore/src/utils/priority-queue.ts       |   6 +-
 yarn.lock                                     |   9 -
 6 files changed, 532 insertions(+), 14 deletions(-)
 create mode 100644 packages/common/nbstore/src/utils/__tests__/binary-search-tree.spec.ts
 create mode 100644 packages/common/nbstore/src/utils/binary-search-tree.ts

diff --git a/packages/common/infra/package.json b/packages/common/infra/package.json
index 492d200ee4..e248824f6f 100644
--- a/packages/common/infra/package.json
+++ b/packages/common/infra/package.json
@@ -16,7 +16,6 @@
     "@affine/env": "workspace:*",
     "@affine/error": "workspace:*",
     "@affine/templates": "workspace:*",
-    "@datastructures-js/binary-search-tree": "^5.3.2",
     "@preact/signals-core": "^1.8.0",
     "eventemitter2": "^6.4.9",
     "foxact": "^0.2.43",
diff --git a/packages/common/nbstore/package.json b/packages/common/nbstore/package.json
index 3790a907c5..9483dc511c 100644
--- a/packages/common/nbstore/package.json
+++ b/packages/common/nbstore/package.json
@@ -18,7 +18,6 @@
     "./frontend": "./src/frontend/index.ts"
   },
   "dependencies": {
-    "@datastructures-js/binary-search-tree": "^5.3.2",
     "@toeverything/infra": "workspace:*",
     "eventemitter2": "^6.4.9",
     "graphemer": "^1.4.0",
diff --git a/packages/common/nbstore/src/utils/__tests__/binary-search-tree.spec.ts b/packages/common/nbstore/src/utils/__tests__/binary-search-tree.spec.ts
new file mode 100644
index 0000000000..dcd39fc07e
--- /dev/null
+++ b/packages/common/nbstore/src/utils/__tests__/binary-search-tree.spec.ts
@@ -0,0 +1,233 @@
+import { describe, expect, test } from 'vitest';
+
+import { BinarySearchTree } from '../binary-search-tree';
+
+describe('Binary Search Tree', () => {
+  // Helper function to create a tree with numbers
+  function createNumberTree() {
+    const tree = new BinarySearchTree<number>((a, b) => a - b);
+    return tree;
+  }
+
+  // Helper function to create a tree with objects
+  function createObjectTree() {
+    const tree = new BinarySearchTree<{ id: string; value: number }>((a, b) => {
+      return a.value === b.value
+        ? a.id === b.id
+          ? 0
+          : a.id > b.id
+            ? 1
+            : -1
+        : a.value - b.value;
+    });
+    return tree;
+  }
+
+  test('insert and find', () => {
+    const tree = createNumberTree();
+
+    // Insert values
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+    tree.insert(3);
+    tree.insert(7);
+
+    // Check if values can be found
+    expect(tree.find(10)).not.toBeNull();
+    expect(tree.find(5)).not.toBeNull();
+    expect(tree.find(15)).not.toBeNull();
+    expect(tree.find(3)).not.toBeNull();
+    expect(tree.find(7)).not.toBeNull();
+
+    // Check value that doesn't exist
+    expect(tree.find(99)).toBeNull();
+
+    // Check tree structure
+    expect(tree.root?.value).toBe(10);
+    expect(tree.root?.left?.value).toBe(5);
+    expect(tree.root?.right?.value).toBe(15);
+    expect(tree.root?.left?.left?.value).toBe(3);
+    expect(tree.root?.left?.right?.value).toBe(7);
+  });
+
+  test('min and max', () => {
+    const tree = createNumberTree();
+
+    // Empty tree
+    expect(tree.min()).toBeNull();
+    expect(tree.max()).toBeNull();
+
+    // Insert values
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+    tree.insert(3);
+    tree.insert(7);
+    tree.insert(20);
+    tree.insert(1);
+
+    // Check min and max
+    expect(tree.min()?.getValue()).toBe(1);
+    expect(tree.max()?.getValue()).toBe(20);
+  });
+
+  test('remove leaf node', () => {
+    const tree = createNumberTree();
+
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+
+    // Remove a leaf node
+    expect(tree.remove(5)).toBe(true);
+    expect(tree.find(5)).toBeNull();
+    expect(tree.root?.left).toBeNull();
+    expect(tree.count()).toBe(2);
+  });
+
+  test('remove node with one child', () => {
+    const tree = createNumberTree();
+
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+    tree.insert(3);
+
+    // Remove a node with one child
+    expect(tree.remove(5)).toBe(true);
+    expect(tree.find(5)).toBeNull();
+    expect(tree.root?.left?.value).toBe(3);
+    expect(tree.count()).toBe(3);
+  });
+
+  test('remove node with two children', () => {
+    const tree = createNumberTree();
+
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+    tree.insert(3);
+    tree.insert(7);
+
+    // Remove a node with two children
+    expect(tree.remove(5)).toBe(true);
+    expect(tree.find(5)).toBeNull();
+
+    // The 7 should replace 5
+    expect(tree.root?.left?.value).toBe(7);
+    expect(tree.root?.left?.left?.value).toBe(3);
+    expect(tree.root?.left?.right).toBeNull();
+    expect(tree.count()).toBe(4);
+  });
+
+  test('remove root node', () => {
+    const tree = createNumberTree();
+
+    // Single node tree
+    tree.insert(10);
+    expect(tree.remove(10)).toBe(true);
+    expect(tree.root).toBeNull();
+    expect(tree.count()).toBe(0);
+
+    // Root with children
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+
+    expect(tree.remove(10)).toBe(true);
+    expect(tree.find(10)).toBeNull();
+    // 15 should become the new root
+    expect(tree.root?.value).toBe(15);
+    expect(tree.root?.left?.value).toBe(5);
+    expect(tree.count()).toBe(2);
+  });
+
+  test('clear tree', () => {
+    const tree = createNumberTree();
+
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+
+    tree.clear();
+    expect(tree.root).toBeNull();
+    expect(tree.count()).toBe(0);
+  });
+
+  test('count', () => {
+    const tree = createNumberTree();
+
+    expect(tree.count()).toBe(0);
+
+    tree.insert(10);
+    expect(tree.count()).toBe(1);
+
+    tree.insert(5);
+    tree.insert(15);
+    expect(tree.count()).toBe(3);
+
+    tree.remove(15);
+    expect(tree.count()).toBe(2);
+
+    tree.clear();
+    expect(tree.count()).toBe(0);
+  });
+
+  test('object comparisons', () => {
+    const tree = createObjectTree();
+
+    tree.insert({ id: 'a', value: 10 });
+    tree.insert({ id: 'b', value: 5 });
+    tree.insert({ id: 'c', value: 15 });
+
+    // Same value, different id
+    tree.insert({ id: 'd', value: 5 });
+
+    // Find by value and id
+    expect(tree.find({ id: 'b', value: 5 })?.getValue()).toEqual({
+      id: 'b',
+      value: 5,
+    });
+    expect(tree.find({ id: 'd', value: 5 })?.getValue()).toEqual({
+      id: 'd',
+      value: 5,
+    });
+
+    // The min should be the one with the lowest id
+    expect(tree.min()?.getValue()).toEqual({ id: 'b', value: 5 });
+
+    // Remove one of the objects with value 5
+    expect(tree.remove({ id: 'b', value: 5 })).toBe(true);
+
+    // The min should now be the other object with value 5
+    expect(tree.min()?.getValue()).toEqual({ id: 'd', value: 5 });
+  });
+
+  test('replace node with same value', () => {
+    const tree = createNumberTree();
+
+    const node1 = tree.insert(10);
+    const node2 = tree.insert(10); // This should replace the previous node
+
+    expect(node1).toBe(node2);
+    expect(tree.count()).toBe(1);
+  });
+
+  test('removeNode', () => {
+    const tree = createNumberTree();
+
+    tree.insert(10);
+    tree.insert(5);
+    tree.insert(15);
+
+    const node = tree.find(5);
+    expect(node).not.toBeNull();
+
+    if (node) {
+      tree.removeNode(node);
+      expect(tree.find(5)).toBeNull();
+      expect(tree.count()).toBe(2);
+    }
+  });
+});
diff --git a/packages/common/nbstore/src/utils/binary-search-tree.ts b/packages/common/nbstore/src/utils/binary-search-tree.ts
new file mode 100644
index 0000000000..70966b1972
--- /dev/null
+++ b/packages/common/nbstore/src/utils/binary-search-tree.ts
@@ -0,0 +1,296 @@
+/**
+ * Represents a node in the binary search tree.
+ */
+export class TreeNode<T> {
+  value: T;
+  left: TreeNode<T> | null = null;
+  right: TreeNode<T> | null = null;
+  parent: TreeNode<T> | null = null;
+
+  constructor(value: T) {
+    this.value = value;
+  }
+
+  /**
+   * Gets the value stored in this node.
+   */
+  getValue(): T {
+    return this.value;
+  }
+}
+
+/**
+ * Binary Search Tree implementation using iterative (non-recursive) algorithms.
+ */
+export class BinarySearchTree<T> {
+  root: TreeNode<T> | null = null;
+  private size = 0;
+  private readonly compareFunction: (a: T, b: T) => number;
+
+  /**
+   * Creates a new binary search tree with a custom comparison function.
+   *
+   * @param compareFunction - Function that compares two elements.
+   *   Should return negative if a < b, zero if a = b, positive if a > b.
+   */
+  constructor(compareFunction: (a: T, b: T) => number) {
+    this.compareFunction = compareFunction;
+  }
+
+  /**
+   * Inserts a value into the tree.
+   *
+   * @param value - The value to insert
+   * @returns The newly created node
+   */
+  insert(value: T): TreeNode<T> {
+    const newNode = new TreeNode<T>(value);
+
+    if (this.root === null) {
+      this.root = newNode;
+      this.size++;
+      return newNode;
+    }
+
+    let current = this.root;
+
+    while (true) {
+      const compareResult = this.compareFunction(value, current.value);
+
+      if (compareResult < 0) {
+        // Go left
+        if (current.left === null) {
+          current.left = newNode;
+          newNode.parent = current;
+          this.size++;
+          return newNode;
+        }
+        current = current.left;
+      } else if (compareResult > 0) {
+        // Go right
+        if (current.right === null) {
+          current.right = newNode;
+          newNode.parent = current;
+          this.size++;
+          return newNode;
+        }
+        current = current.right;
+      } else {
+        // Value already exists, replace it
+        current.value = value;
+        return current;
+      }
+    }
+  }
+
+  /**
+   * Finds a node with the given value.
+   *
+   * @param value - The value to find
+   * @returns The node containing the value, or null if not found
+   */
+  find(value: T): TreeNode<T> | null {
+    let current = this.root;
+
+    while (current !== null) {
+      const compareResult = this.compareFunction(value, current.value);
+
+      if (compareResult === 0) {
+        return current;
+      } else if (compareResult < 0) {
+        current = current.left;
+      } else {
+        current = current.right;
+      }
+    }
+
+    return null;
+  }
+
+  /**
+   * Removes a value from the tree.
+   *
+   * @param value - The value to remove
+   * @returns True if the value was removed, false if it wasn't found
+   */
+  remove(value: T): boolean {
+    const nodeToRemove = this.find(value);
+
+    if (nodeToRemove === null) {
+      return false;
+    }
+
+    this.removeNode(nodeToRemove);
+    return true;
+  }
+
+  /**
+   * Removes a specific node from the tree.
+   *
+   * @param node - The node to remove
+   */
+  removeNode(node: TreeNode<T>): void {
+    // Case 1: Node has no children
+    if (node.left === null && node.right === null) {
+      if (node.parent === null) {
+        // Node is root
+        this.root = null;
+      } else if (node === node.parent.left) {
+        node.parent.left = null;
+      } else {
+        node.parent.right = null;
+      }
+      this.size--;
+    }
+    // Case 2: Node has one child (right)
+    else if (node.left === null) {
+      if (node.parent === null) {
+        // Node is root
+        this.root = node.right;
+        if (this.root) {
+          this.root.parent = null;
+        }
+      } else if (node === node.parent.left) {
+        node.parent.left = node.right;
+        if (node.right) {
+          node.right.parent = node.parent;
+        }
+      } else {
+        node.parent.right = node.right;
+        if (node.right) {
+          node.right.parent = node.parent;
+        }
+      }
+      this.size--;
+    }
+    // Case 3: Node has one child (left)
+    else if (node.right === null) {
+      if (node.parent === null) {
+        // Node is root
+        this.root = node.left;
+        if (this.root) {
+          this.root.parent = null;
+        }
+      } else if (node === node.parent.left) {
+        node.parent.left = node.left;
+        if (node.left) {
+          node.left.parent = node.parent;
+        }
+      } else {
+        node.parent.right = node.left;
+        if (node.left) {
+          node.left.parent = node.parent;
+        }
+      }
+      this.size--;
+    }
+    // Case 4: Node has two children
+    else {
+      // Find the successor (minimum value in right subtree)
+      const successor = this.findMin(node.right);
+
+      // Save successor value
+      const successorValue = successor.value;
+
+      // Instead of recursively calling removeNode, we'll handle successor removal directly
+      // The successor must have at most one child (the right child)
+
+      // Remove the successor - it can't have a left child by definition
+      if (successor.parent === node) {
+        // Successor is direct child of node we're removing
+        node.right = successor.right;
+        if (successor.right) {
+          successor.right.parent = node;
+        }
+      } else {
+        // Successor is further down the tree
+        // oxlint-disable-next-line no-non-null-assertion
+        successor.parent!.left = successor.right;
+        if (successor.right) {
+          successor.right.parent = successor.parent;
+        }
+      }
+
+      // Copy successor value to the node we're removing
+      node.value = successorValue;
+
+      // Decrement size counter
+      this.size--;
+    }
+  }
+
+  /**
+   * Finds the node with the minimum value in the subtree rooted at the given node.
+   *
+   * @param subtreeRoot - The root of the subtree to search
+   * @returns The node with the minimum value
+   */
+  private findMin(subtreeRoot: TreeNode<T>): TreeNode<T> {
+    let current = subtreeRoot;
+
+    while (current.left !== null) {
+      current = current.left;
+    }
+
+    return current;
+  }
+
+  /**
+   * Finds the node with the maximum value in the subtree rooted at the given node.
+   *
+   * @param subtreeRoot - The root of the subtree to search
+   * @returns The node with the maximum value
+   */
+  private findMax(subtreeRoot: TreeNode<T>): TreeNode<T> {
+    let current = subtreeRoot;
+
+    while (current.right !== null) {
+      current = current.right;
+    }
+
+    return current;
+  }
+
+  /**
+   * Returns the node with the maximum value in the tree.
+   *
+   * @returns The node with the maximum value, or null if the tree is empty
+   */
+  max(): TreeNode<T> | null {
+    if (this.root === null) {
+      return null;
+    }
+
+    return this.findMax(this.root);
+  }
+
+  /**
+   * Returns the node with the minimum value in the tree.
+   *
+   * @returns The node with the minimum value, or null if the tree is empty
+   */
+  min(): TreeNode<T> | null {
+    if (this.root === null) {
+      return null;
+    }
+
+    return this.findMin(this.root);
+  }
+
+  /**
+   * Clears all nodes from the tree.
+   */
+  clear(): void {
+    this.root = null;
+    this.size = 0;
+  }
+
+  /**
+   * Returns the number of nodes in the tree.
+   *
+   * @returns The number of nodes
+   */
+  count(): number {
+    return this.size;
+  }
+}
diff --git a/packages/common/nbstore/src/utils/priority-queue.ts b/packages/common/nbstore/src/utils/priority-queue.ts
index 9cee608840..f27ae7321b 100644
--- a/packages/common/nbstore/src/utils/priority-queue.ts
+++ b/packages/common/nbstore/src/utils/priority-queue.ts
@@ -1,4 +1,4 @@
-import { BinarySearchTree } from '@datastructures-js/binary-search-tree';
+import { BinarySearchTree } from './binary-search-tree';
 
 export class PriorityQueue {
   tree = new BinarySearchTree<{ id: string; priority: number }>((a, b) => {
@@ -39,8 +39,8 @@ export class PriorityQueue {
     return id;
   }
 
-  remove(id: string, priority?: number) {
-    priority ??= this.priorityMap.get(id);
+  remove(id: string) {
+    const priority = this.priorityMap.get(id);
     if (priority === undefined) {
       return false;
     }
diff --git a/yarn.lock b/yarn.lock
index 4cfffbe785..e798909748 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -813,7 +813,6 @@ __metadata:
     "@affine/error": "workspace:*"
     "@affine/graphql": "workspace:*"
     "@blocksuite/affine": "workspace:*"
-    "@datastructures-js/binary-search-tree": "npm:^5.3.2"
     "@toeverything/infra": "workspace:*"
     eventemitter2: "npm:^6.4.9"
     fake-indexeddb: "npm:^6.0.0"
@@ -4377,13 +4376,6 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@datastructures-js/binary-search-tree@npm:^5.3.2":
-  version: 5.3.2
-  resolution: "@datastructures-js/binary-search-tree@npm:5.3.2"
-  checksum: 10/05936b8710e7db5f6e748ffbcd857ca692ba5de8edb87988d5f093085df9e9a15af5caa1a9de020682703ba15f425e2fd513c540a1cb1da293bf0c67793cdae7
-  languageName: node
-  linkType: hard
-
 "@date-fns/tz@npm:^1.2.0":
   version: 1.2.0
   resolution: "@date-fns/tz@npm:1.2.0"
@@ -14134,7 +14126,6 @@ __metadata:
     "@affine/env": "workspace:*"
     "@affine/error": "workspace:*"
     "@affine/templates": "workspace:*"
-    "@datastructures-js/binary-search-tree": "npm:^5.3.2"
     "@emotion/react": "npm:^11.14.0"
     "@preact/signals-core": "npm:^1.8.0"
     "@swc/core": "npm:^1.10.1"