Merge pull request #553 from phryneas/pr/trie-remove

Add `remove` method to `Trie`

.github/workflows/main.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        node_version: ['14', '16', '18', '19', '20']
+        node_version: ['16', '18', '19', '20', '21']
         os: [ubuntu-latest]
 
     steps:

packages/trie/src/index.ts

@@ -23,8 +23,9 @@ export class Trie<Data> {
     private makeData: (array: any[]) => Data = defaultMakeData,
   ) {}
 
-  public lookup<T extends any[]>(...array: T): Data {
-    return this.lookupArray(array);
+  public lookup<T extends any[]>(...array: T): Data;
+  public lookup(): Data {
+    return this.lookupArray(arguments);
   }
 
   public lookupArray<T extends IArguments | any[]>(array: T): Data {
@@ -35,31 +36,60 @@ export class Trie<Data> {
       : node.data = this.makeData(slice.call(array));
   }
 
-  public peek<T extends any[]>(...array: T): Data | undefined {
-    return this.peekArray(array);
+  public peek<T extends any[]>(...array: T): Data | undefined;
+  public peek(): Data | undefined {
+    return this.peekArray(arguments);
   }
 
   public peekArray<T extends IArguments | any[]>(array: T): Data | undefined {
     let node: Trie<Data> | undefined = this;
 
     for (let i = 0, len = array.length; node && i < len; ++i) {
-      const map: Trie<Data>["weak" | "strong"] =
-        this.weakness && isObjRef(array[i]) ? node.weak : node.strong;
-
+      const map = node.mapFor(array[i], false);
       node = map && map.get(array[i]);
     }
 
     return node && node.data;
   }
 
+  public remove(...array: any[]): Data | undefined;
+  public remove(): Data | undefined {
+    return this.removeArray(arguments);
+  }
+
+  public removeArray<T extends IArguments | any[]>(array: T): Data | undefined {
+    let data: Data | undefined;
+
+    if (array.length) {
+      const head = array[0];
+      const map = this.mapFor(head, false);
+      const child = map && map.get(head);
+      if (child) {
+        data = child.removeArray(slice.call(array, 1));
+        if (!child.data && !child.weak && !(child.strong && child.strong.size)) {
+          map.delete(head);
+        }
+      }
+    } else {
+      data = this.data;
+      delete this.data;
+    }
+
+    return data;
+  }
+
   private getChildTrie(key: any) {
-    const map = this.weakness && isObjRef(key)
-      ? this.weak || (this.weak = new WeakMap<any, Trie<Data>>())
-      : this.strong || (this.strong = new Map<any, Trie<Data>>());
+    const map = this.mapFor(key, true)!;
     let child = map.get(key);
     if (!child) map.set(key, child = new Trie<Data>(this.weakness, this.makeData));
     return child;
   }
+
+  private mapFor(key: any, create: boolean): Trie<Data>["weak" | "strong"] | undefined {
+    return this.weakness && isObjRef(key)
+      ? this.weak || (create ? this.weak = new WeakMap : void 0)
+      : this.strong || (create ? this.strong = new Map : void 0);
+  }
 }
 
 function isObjRef(value: any) {

packages/trie/src/tests/main.ts

@@ -109,4 +109,99 @@ describe("Trie", function () {
     assert.strictEqual(trie.peekArray([1, 2, 'x']), data1);
     assert.strictEqual(trie.peekArray([1, 2, obj]), data2);
   });
+
+  describe("can remove values", function () {
+    it("will remove values", () => {
+      const trie = new Trie(true, (args) => args);
+
+      trie.lookup(1, 2, "x");
+      trie.remove(1, 2, "x");
+      assert.strictEqual(trie.peek(1, 2, "x"), undefined);
+    });
+
+    it("removing will return the value", () => {
+      const trie = new Trie(true, (args) => args);
+
+      const data = trie.lookup(1, 2, "x");
+      assert.strictEqual(trie.remove(1, 2, "x"), data);
+    });
+
+    it("will remove empty parent nodes", () => {
+      const trie = new Trie(true, (args) => args);
+
+      const data = trie.lookup(1, 2, "x");
+      assert.strictEqual(trie.peek(1, 2, "x"), data);
+      assert.equal(pathExistsInTrie(trie, 1, 2, "x"), true);
+      assert.strictEqual(trie.remove(1, 2, "x"), data);
+      assert.equal(pathExistsInTrie(trie, 1), false);
+    });
+
+    it("will not remove parent nodes with other children", () => {
+      const trie = new Trie(true, (args) => args);
+
+      trie.lookup(1, 2, "x");
+      const data = trie.lookup(1, 2);
+      trie.remove(1, 2, "x");
+      assert.strictEqual(trie.peek(1, 2, "x"), undefined);
+      assert.strictEqual(trie.peek(1, 2), data);
+    });
+
+    it("will remove data, not the full node, if a node still has children", () => {
+      const trie = new Trie(true, (args) => args);
+
+      trie.lookup(1, 2);
+      const data = trie.lookup(1, 2, "x");
+      trie.remove(1, 2);
+      assert.strictEqual(trie.peek(1, 2), undefined);
+      assert.strictEqual(trie.peek(1, 2, "x"), data);
+    });
+
+    it("will remove direct children", () => {
+      const trie = new Trie(true, (args) => args);
+
+      trie.lookup(1);
+      trie.remove(1);
+      assert.strictEqual(trie.peek(1), undefined);
+    });
+
+    it("will remove nodes from WeakMaps", () => {
+      const trie = new Trie(true, (args) => args);
+      const obj = {};
+      const data = trie.lookup(1, obj, "x");
+      assert.equal(pathExistsInTrie(trie, 1), true);
+      assert.strictEqual(trie.remove(1, obj, "x"), data);
+      assert.strictEqual(trie.peek(1, obj, "x"), undefined);
+      assert.equal(pathExistsInTrie(trie, 1, obj), false);
+    });
+
+    it("will not remove nodes if they contain an (even empty) WeakMap", () => {
+      const trie = new Trie(true, (args) => args);
+      const obj = {};
+
+      const data = trie.lookup(1, 2, "x");
+      trie.lookup(1, obj);
+      trie.remove(1, obj);
+
+      assert.strictEqual(trie.peek(1, 2, "x"), data);
+      assert.equal(pathExistsInTrie(trie, 1), true);
+      assert.equal(pathExistsInTrie(trie, 1, 2), true);
+      assert.strictEqual(trie.remove(1, 2, "x"), data);
+      assert.equal(pathExistsInTrie(trie, 1), true);
+      assert.equal(pathExistsInTrie(trie, 1, 2), false);
+    });
+  });
+
+  function pathExistsInTrie(trie: Trie<unknown>, ...path: any[]) {
+    return (
+      path.reduce((node: Trie<unknown> | undefined, key: any) => {
+        const map: Trie<unknown>["weak" | "strong"] =
+          // not the full  implementation but enough for a test
+          trie["weakness"] && typeof key === "object"
+            ? node?.["weak"]
+            : node?.["strong"];
+
+        return map?.get(key);
+      }, trie) !== undefined
+    );
+  }
 });