feat(core): improve priority queue performance (#11559)

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",

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",

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);
+    }
+  });
+});

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;
+  }
+}

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;
     }