From 768435365cab644f4e677e2b132c006e0c201f5f Mon Sep 17 00:00:00 2001
From: "Joshua A. Horton" <joshua_horton@sil.org>
Date: Thu, 4 Jul 2024 15:14:51 +0700
Subject: [PATCH] refactor(common/models): refactor migrated trie-construction
 code for improved encapsulation

---
 common/models/templates/src/index.ts        |  2 +-
 common/models/templates/src/trie-builder.ts | 20 ++++++++++++-
 developer/src/kmc-model/src/build-trie.ts   | 32 ++-------------------
 3 files changed, 23 insertions(+), 31 deletions(-)

diff --git a/common/models/templates/src/index.ts b/common/models/templates/src/index.ts
index b32c37d45a..b67528a355 100644
--- a/common/models/templates/src/index.ts
+++ b/common/models/templates/src/index.ts
@@ -7,5 +7,5 @@ export { Tokenization, tokenize, getLastPreCaretToken, wordbreak } from "./token
 export {
   Entry, InternalNode, Leaf, Node
 } from './trie.js';
-export { addUnsorted, TrieBuilder } from './trie-builder.js';
+export { TrieBuilder } from './trie-builder.js';
 export { default as TrieModel, TrieModelOptions } from "./trie-model.js";
\ No newline at end of file
diff --git a/common/models/templates/src/trie-builder.ts b/common/models/templates/src/trie-builder.ts
index 8b90b5f28d..0ac5cfa575 100644
--- a/common/models/templates/src/trie-builder.ts
+++ b/common/models/templates/src/trie-builder.ts
@@ -20,7 +20,7 @@ function createRootNode(): Node {
  * @param index the index in the key and also the trie depth. Should be set to
  *              zero when adding onto the root node of the trie.
  */
-export function addUnsorted(node: Node, entry: Entry, index: number = 0) {
+function addUnsorted(node: Node, entry: Entry, index: number = 0) {
   // Each node stores the MAXIMUM weight out of all of its decesdents, to
   // enable a greedy search through the trie.
   node.weight = Math.max(node.weight, entry.weight);
@@ -145,6 +145,20 @@ export class TrieBuilder extends Trie {
     this.totalWeight = 0;
   }
 
+  addEntry(word: string, weight?: number) {
+    weight = (isNaN(weight ?? NaN) || !weight) ? 1 : weight;
+    this.totalWeight += weight;
+
+    addUnsorted(
+      this.root, {
+        key: this.toKey(word),
+        content: word,
+        weight: weight
+      },
+      0
+    );
+  }
+
   sort() {
     sortTrie(this.root);
   }
@@ -152,4 +166,8 @@ export class TrieBuilder extends Trie {
   getRoot(): Node {
     return this.root;
   }
+
+  getTotalWeight(): number {
+    return this.totalWeight;
+  }
 }
\ No newline at end of file
diff --git a/developer/src/kmc-model/src/build-trie.ts b/developer/src/kmc-model/src/build-trie.ts
index efcda3d38c..9e34cd5469 100644
--- a/developer/src/kmc-model/src/build-trie.ts
+++ b/developer/src/kmc-model/src/build-trie.ts
@@ -1,7 +1,7 @@
 import { ModelCompilerError, ModelCompilerMessageContext, ModelCompilerMessages } from "./model-compiler-messages.js";
 import { callbacks } from "./compiler-callbacks.js";
 
-import { addUnsorted, Node, TrieBuilder } from '@keymanapp/models-templates';
+import { TrieBuilder } from '@keymanapp/models-templates';
 
 // Supports LF or CRLF line terminators.
 const NEWLINE_SEPARATOR = /\u000d?\u000a/;
@@ -217,7 +217,7 @@ export function buildTrie(wordlist: WordList, keyFunction: SearchTermToKey): obj
 
   buildFromWordList(collater, wordlist);
   return {
-    totalWeight: sumWeights(collater.getRoot()),
+    totalWeight: collater.getTotalWeight(),
     root: collater.getRoot()
   }
 }
@@ -228,38 +228,12 @@ export function buildTrie(wordlist: WordList, keyFunction: SearchTermToKey): obj
  */
 function buildFromWordList(trieCollator: TrieBuilder, words: WordList): TrieBuilder {
   for (let [wordform, weight] of Object.entries(words)) {
-    let key = trieCollator.toKey(wordform);
-    addUnsorted(trieCollator.getRoot(), { content: wordform, key, weight });
+    trieCollator.addEntry(wordform, weight);
   }
   trieCollator.sort();
   return trieCollator;
 }
 
-/**
- * O(n) recursive traversal to sum the total weight of all leaves in the
- * trie, starting at the provided node.
- *
- * @param node The node to start summing weights.
- */
-function sumWeights(node: Node): number {
-  let val: number;
-  if (node.type === 'leaf') {
-    val = node.entries
-      .map(entry => entry.weight)
-      //.map(entry => isNaN(entry.weight) ? 1 : entry.weight)
-      .reduce((acc, count) => acc + count, 0);
-  } else {
-    val = Object.keys(node.children)
-      .map((key) => sumWeights(node.children[key]))
-      .reduce((acc, count) => acc + count, 0);
-  }
-
-  if(isNaN(val)) {
-    throw new Error("Unexpected NaN has appeared!");
-  }
-  return val;
-}
-
 /**
  * Detects the encoding of a text file.
  *