From fe6e86e1a472d1613ec2ecb1efd7f340ac09c91f Mon Sep 17 00:00:00 2001
From: Joeri Exelmans <joeri.exelmans@gmail.com>
Date: Sat, 10 May 2025 20:28:40 +0200
Subject: [PATCH] fix bug in trie

---
 lib/util/trie.js | 54 +++++++++++++++++++++++--------
 tests/trie.js    | 82 +++++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 112 insertions(+), 24 deletions(-)

diff --git a/lib/util/trie.js b/lib/util/trie.js
index 8bc5966..433c630 100644
--- a/lib/util/trie.js
+++ b/lib/util/trie.js
@@ -6,6 +6,21 @@ export const emptyTrie = {
   children: [],
 };
 
+// invariant: the keys of a trie node should always be sorted
+export const isProperlySorted = trie => {
+  for (let i=0; i<trie.children.length-1; i++) {
+    if (trie.children[i] >= trie.children[i+1]) {
+      return false; // not properly sorted!
+    }
+  }
+  for (const [_, node] of trie.children) {
+    if (!isProperlySorted(node)) {
+      return false;
+    }
+  }
+  return true;
+}
+
 // find maximal common prefix, and whether string A is smaller than B
 const commonPrefix = (strA, strB) => {
   let i=0
@@ -41,6 +56,14 @@ const __binarySearch = (ls, key, min, max) => {
   return __binarySearch(ls, key, middle+1, max);
 }
 
+const check = trie => {
+  // uncomment if you think shit is broken:
+  // if (!isProperlySorted(trie)) {
+  //   throw new Error('not properly sorted!')
+  // }
+  return trie;
+}
+
 // insert (key,value) into trie.
 export const insert = trie => key => value => {
   if (key.length === 0) {
@@ -54,54 +77,59 @@ export const insert = trie => key => value => {
   const [insertPos, prefix] = binarySearch(trie.children, key);
   if (insertPos === trie.children.length) {
     // insert node at end
-    return {
+    return check({
       value: trie.value,
       children: [
         ...trie.children,
         [key, {value, children:[]}],
       ],
-    };
+    });
   }
   if (prefix.length === 0) {
     // nothing in common...
     // insert new node into children
-    return {
+    return check({
       value: trie.value,
       children: trie.children.toSpliced(
         insertPos, // insert position
         0, // delete nothing
         [key, {value, children:[]}],
       ),
-    };
+    });
   }
 
   const [haveKey, haveChildNode] = trie.children[insertPos];
   if (prefix.length === haveKey.length) {
     // recurse
-    return {
+    return check({
       value: trie.value,
       children: trie.children.with(
         insertPos, // position to update
         [haveKey, insert(haveChildNode)(key.slice(prefix.length))(value)],
       )
-    }
+    })
   }
 
   // otherwise, split entry:
   const havePostFix = haveKey.slice(prefix.length);
   const postFix = key.slice(prefix.length);
-  return {
+  return check({
     value: trie.value,
     children: trie.children.with(
       insertPos, // position to update
       [prefix, {
-        children: [
-          [havePostFix, haveChildNode],
-          [postFix, {value, children: []}
-        ],
-      ]}],
+        children: (havePostFix < postFix)
+          ? [
+              [havePostFix, haveChildNode],
+              [postFix, {value, children: []}],
+            ]
+          : [
+              [postFix, {value, children: []}],
+              [havePostFix, haveChildNode],
+            ],
+      }],
     ),
-  };
+  });
 };
 
 // given a prefix, return a string X such that prefix+X is a possibly larger prefix for the same entries as the original prefix.
diff --git a/tests/trie.js b/tests/trie.js
index 7016b97..2c6eaf8 100644
--- a/tests/trie.js
+++ b/tests/trie.js
@@ -1,7 +1,14 @@
-import { pretty } from "../lib/util/pretty.js";
-import { insert, emptyTrie, growPrefix, suggest } from "../lib/util/trie.js";
+import * as assert from "node:assert";
+
+import { module2Env } from "../lib/environment/env.js";
+import { ModuleStd } from "../lib/stdlib.js";
+import { pretty } from "../lib/util/pretty.js";
+import { emptyTrie, growPrefix, insert, isProperlySorted, suggest } from "../lib/util/trie.js";
+
+
+///////////////////
+// Setup TRIE 1 ...
 
-// insertion
 const with1Item = insert(emptyTrie)('abba')('dancing queen');
 console.log(pretty(with1Item));
 const with2Items = insert(with1Item)('aboriginal')('australia');
@@ -19,14 +26,67 @@ console.log(pretty(with7Items));
 const with8Items = insert(with7Items)('')('hi!');
 console.log(pretty(with8Items));
 
+///////////////////
+// Setup TRIE 2 ...
+
+const bigTrie = module2Env(ModuleStd).name2dyn;
+
+///////////////////
+// Test...
+
 // grow key (for auto-complete)
-console.log(growPrefix(with6Items)("a")); // b
-console.log(growPrefix(with6Items)("ab")); // (empty string)
-console.log(growPrefix(with6Items)("abb")); // a
-console.log(growPrefix(with6Items)("f")); // ood
-console.log(growPrefix(with6Items)("abo")); // riginal
+assert.equal(growPrefix(with6Items)("a"  ), "b"      );
+assert.equal(growPrefix(with6Items)("ab" ), ""       );
+assert.equal(growPrefix(with6Items)("abb"), "a"      );
+assert.equal(growPrefix(with6Items)("f"  ), "ood"    );
+assert.equal(growPrefix(with6Items)("abo"), "riginal");
 
 // suggest (also for auto-complete)
-console.log(suggest(with8Items)("a")(3)); // 'ab', 'abba', 'aboriginal'
-console.log(suggest(with8Items)("a")(4)); // 'ab', 'abba', 'aboriginal', 'aboriginally'
-console.log(suggest(with8Items)("a")(5)); // 'ab', 'abba', 'aboriginal', 'aboriginally', 'absent'
+assert.deepEqual(
+  suggest(with8Items)("a")(3),
+  [
+    [ 'ab', 'yup' ],
+    [ 'abba', 'dancing queen' ],
+    [ 'aboriginal', 'australia' ]
+  ]);
+
+assert.deepEqual(
+  suggest(with8Items)("a")(4),
+  [
+    [ 'ab', 'yup' ],
+    [ 'abba', 'dancing queen' ],
+    [ 'aboriginal', 'australia' ],
+    [ 'aboriginally', '??' ]
+  ]);
+
+assert.deepEqual(
+  suggest(with8Items)("a")(5),
+  [
+    [ 'ab', 'yup' ],
+    [ 'abba', 'dancing queen' ],
+    [ 'aboriginal', 'australia' ],
+    [ 'aboriginally', '??' ],
+    [ 'absent', 'not here' ]
+  ]);
+
+assert.deepEqual(
+  suggest(with8Items)("a")(15),
+  [
+    [ 'ab', 'yup' ],
+    [ 'abba', 'dancing queen' ],
+    [ 'aboriginal', 'australia' ],
+    [ 'aboriginally', '??' ],
+    [ 'absent', 'not here' ]
+  ]);
+
+assert.equal(
+  isProperlySorted(with8Items),
+  true,
+  "trie should always be properly sorted!"
+);
+
+assert.equal(
+  isProperlySorted(bigTrie),
+  true,
+  "trie should always be properly sorted!"
+);