Trie optimizations

scripts/parser-helper.ts

@@ -106,12 +106,27 @@ export function validateTrie(root: TrieNode, prefixes: [string, number][]): void
 }
 
 export function minimizeIds(root: TrieNode): void {
-	let i = 0;
+	const counter = new Map<TrieNode, number>();
+
 	for (const node of root.getAllNodes()) {
-		node.id = i++;
+		counter.set(node, counter.get(node) ?? 0);
+		for (const c of [...node.children.values(), ...node.shortcuts.values()]) {
+			counter.set(node, (counter.get(c) ?? 0) + 1);
+		}
 	}
 
-	console.log('Minimized ids with', root.getAllNodes().size, 'nodes');
+	const nodes = [...counter];
+	nodes.sort((a, b) => {
+		if (a[1] !== b[1]) return b[1] - a[1];
+		return a[0].id - b[0].id;
+	});
+
+	let i = 0;
+	for (const node of nodes) {
+		node[0].id = i++;
+	}
+
+	console.log('Minimized ids with', nodes.length, 'nodes');
 }
 
 interface IEntity {

src/lib/models/trie.test.ts

@@ -52,3 +52,16 @@ describe('parseString', () => {
 		expect(xx?.overrides).toEqual({});
 	});
 });
+
+describe('getAllNodes', () => {
+	test('Basic test', () => {
+		const root = new TrieNode({ id: 0 });
+		const a = new TrieNode({ id: 1 });
+		const b = new TrieNode({ id: 2 });
+
+		root.children.set('A', a);
+		a.shortcuts.set('CB', b);
+
+		expect([...root.getAllNodes()]).toEqual([root, a, b]);
+	});
+});

src/lib/models/trie.ts

@@ -168,13 +168,18 @@ export class TrieNode {
 
 		const newEntity = this.entity ?? currentEntity;
 		const newOverrides = currentOverrides.merge(this.overrides);
-		for (const [k, child] of this.children.entries()) {
+		for (const [k, child] of [...this.children.entries(), ...this.shortcuts.entries()]) {
 			if (child.collapseNodes(newEntity, newOverrides)) {
 				this.children.delete(k);
 			}
 		}
 
-		return this.children.size == 0 && !this.entity && !this.overrides.toString();
+		return (
+			this.children.size === 0 &&
+			this.shortcuts.size === 0 &&
+			!this.entity &&
+			!this.overrides.toString()
+		);
 	}
 
 	buildShortcuts(): void {
@@ -205,6 +210,8 @@ export class TrieNode {
 				stack += key;
 			}
 
+			if (stack.length < 2) continue;
+
 			this.shortcuts.set(stack, curr);
 			this.children.delete(k);
 		}