From feb97c2cd2f3791a92a2f4d7967d1e1cee2d0232 Mon Sep 17 00:00:00 2001
From: tsctx <91457664+tsctx@users.noreply.github.com>
Date: Thu, 7 Dec 2023 16:15:42 +0900
Subject: [PATCH] perf: rewrite tree

---
 lib/core/tree.js | 143 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 110 insertions(+), 33 deletions(-)

diff --git a/lib/core/tree.js b/lib/core/tree.js
index f130179b1a8..081b9a379a9 100644
--- a/lib/core/tree.js
+++ b/lib/core/tree.js
@@ -1,56 +1,133 @@
 const { wellknownHeaderNames } = require('./constants')
 
-class Tree {
-  #node = {}
-  #depth = 0
+class Node {
+  /** @type {any} */
+  value
+  /** @type {null | Node} */
+  left
+  /** @type {null | Node} */
+  middle
+  /** @type {null | Node} */
+  right
+  /** @type {number} */
+  code
+  /**
+   * @param {Uint8Array} key
+   * @param {any} value
+   */
+  constructor (key, value) {
+    const length = key.length
+    if (length === 0) {
+      throw new TypeError('Unreachable')
+    }
+    this.value = null
+    this.left = null
+    this.middle = null
+    this.right = null
+    this.code = key[0]
+    if (length > 1) {
+      this.middle = new Node(key.subarray(1), value)
+    } else {
+      this.value = value
+    }
+  }
 
   /**
-   * Lowercases key and inserts its value into the node.
-   * @param {string} value
+   * @param {Uint8Array} key
+   * @param {any} value
    */
-  insert (value) {
-    const target = Buffer.from((value = value.toLowerCase()))
-    const length = target.length
-    let node = this.#node
-    for (let i = 0; i < length; ++i) {
-      const key = target[i]
-      node[key] ??= {}
-      // a-z
-      if (key >= 0x61 && key <= 0x7a) {
-        // Uppercase letters preserve references to lowercase ones.
-        node[key & ~32] ??= node[key]
+  add (key, value) {
+    const code = key[0]
+    if (this.code === code) {
+      if (key.length === 1) {
+        this.value = value
+      } else if (this.middle !== null) {
+        this.middle.add(key.subarray(1), value)
+      } else {
+        this.middle = new Node(key.subarray(1), value)
+      }
+    } else if (this.code < code) {
+      if (this.left !== null) {
+        this.left.add(key, value)
+      } else {
+        this.left = new Node(key, value)
+      }
+    } else {
+      if (this.right !== null) {
+        this.right.add(key, value)
+      } else {
+        this.right = new Node(key, value)
       }
-      node = node[key]
     }
-    node[256] = value
-    if (length > this.#depth) {
-      this.#depth = length
+  }
+
+  /**
+   * @param {Uint8Array} key
+   * @return {Node | null}
+   */
+  search (key) {
+    const keylength = key.length
+    if (keylength === 0) {
+      return null
     }
+    let index = 0
+    let node = this
+    while (node !== null && index < keylength) {
+      let code = key[index]
+      // A-Z
+      if (code >= 0x41 && code <= 0x5a) {
+        // Lowercase for uppercase.
+        code |= 32
+      }
+      while (node !== null) {
+        if (code === node.code) {
+          if (keylength === ++index) {
+            // Returns Node since it is the last key.
+            return node
+          }
+          node = node.middle
+          break
+        }
+        node = node.code < code ? node.left : node.right
+      }
+    }
+    return null
   }
+}
+
+class TernarySearchTree {
+  /** @type {Node | null} */
+  node = null
 
   /**
-   * Retrieves values from a node.
-   * @param {Uint8Array} key Node Key
-   * @returns {string | null} Value
+   * @param {Uint8Array} key
+   * @param {any} value
+   * */
+  insert (key, value) {
+    if (this.node === null) {
+      this.node = new Node(key, value)
+    } else {
+      this.node.add(key, value)
+    }
+  }
+
+  /**
+   * @param {Uint8Array} key
    */
   lookup (key) {
-    const length = key.length
-    if (length > this.#depth) return null
-    let node = this.#node
-    for (let i = 0; i < length; ++i) {
-      if ((node = node?.[key[i]]) === undefined) return null
-    }
-    return node?.[256] ?? null
+    return this.node?.search(key)?.value ?? null
   }
 }
 
-const tree = new Tree()
+const tree = new TernarySearchTree()
 
 for (let i = 0; i < wellknownHeaderNames.length; ++i) {
-  tree.insert(wellknownHeaderNames[i])
+  const key = wellknownHeaderNames[i]
+  const lowerCasedKey = key.toLowerCase()
+  tree.insert(Buffer.from(lowerCasedKey), lowerCasedKey)
 }
 
 module.exports = {
-  Tree,
+  TernarySearchTree,
   tree
 }