1.2.2

Improved suggestion engine

manifest.json

@@ -1,11 +1,11 @@
 {
 	"id": "text-autocomplete",
 	"name": "Text Autocomplete",
-	"version": "1.2.1",
+	"version": "1.2.2",
 	"minAppVersion": "1.8.10",
 	"description": "Autocomplete text to type more efficiently.",
 	"author": "Christ Degni",
 	"authorUrl": "https://github.com/c-degni",
 	"fundingUrl": "https://buymeacoffee.com/c.degni",
 	"isDesktopOnly": true
-}
+}

package-lock.json

@@ -1,12 +1,12 @@
 {
 	"name": "text-autocomplete",
-	"version": "1.0.0",
+	"version": "1.2.2",
 	"lockfileVersion": 3,
 	"requires": true,
 	"packages": {
 		"": {
 			"name": "text-autocomplete",
-			"version": "1.0.2",
+			"version": "1.2.2",
 			"license": "MIT",
 			"devDependencies": {
 				"@eslint/eslintrc": "^3.3.1",

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "text-autocomplete",
-	"version": "1.2.1",
+	"version": "1.2.2",
 	"description": "Autocomplete text to type more efficiently.",
 	"main": "main.js",
 	"scripts": {

src/constants.ts

@@ -0,0 +1,13 @@
+import { Result } from "./utils";
+
+export const DEFAULT_WORD_SCORE: number = 1;
+export const IMPORT_WORD_SCORE: number = 2;
+export const CUSTOM_WORD_SCORE: number = 5;
+export const ACCEPTED_SUGGESTION_BUMP: number = 3;
+
+export const OK: Result = { ok: true };
+export const NOT_OK: Result = { ok: false };
+export const EMPTY_FILE: Result = { ok: false, reason: 'emptyFile' };
+export const DUPLICATE_FILE: Result = { ok: false, reason: 'duplicateFile' };
+export const EMPTY_WORD: Result = { ok: false, reason: 'emptyWord' };
+export const DUPLICATE_WORD: Result = { ok: false, reason: 'duplicateWord'};

src/dictionary/dictionary.ts

@@ -1,8 +1,12 @@
-import { Trie } from './trie';
 import { DEFAULT_WORDS, CS_WORDS } from './words'
 
-let DEFAULT_TRIE = new Trie();
-DEFAULT_WORDS.forEach(word => DEFAULT_TRIE.insert(word));
-CS_WORDS.forEach(word => DEFAULT_TRIE.insert(word));
+export type Dictionary = {
+    name: string,
+    words: string[],
+    baseScore: number
+};
 
-export { DEFAULT_TRIE };
+export const DEFAULT_DICTIONARIES: Dictionary[] = [
+    { name: 'default', words: DEFAULT_WORDS, baseScore: 1 },
+    { name: 'computer-science', words: CS_WORDS, baseScore: 1 },
+];

src/dictionary/trie.ts

@@ -1,30 +1,55 @@
 class TrieNode {
     children: Record<string, TrieNode>;
     endOfWord: boolean;
+    prefix: string;
     words: Set<string>;
+    cache: string[];
 
-    constructor() {
+    constructor(prefix: string = '') {
         this.children = {};
         this.endOfWord = false;
+        this.prefix = prefix;
         this.words = new Set();
+        this.cache = [];
     }
 }
 
 export default class Trie {
     root: TrieNode;
-    constructor() {
+    wordScores: Map<string, number>;
+    maxCacheSize: number;
+
+    constructor(maxCacheSize: number = 10) {
         this.root = new TrieNode();
+        this.wordScores = new Map();
+        this.maxCacheSize = maxCacheSize;
     }
 
-    insert(word: string): void {
-        let tmp: TrieNode = this.root;
-        const normalized: string = word.toLowerCase();
-        for (const char of normalized) { 
-            if (!tmp.children[char]) tmp.children[char] = new TrieNode();
-            tmp = tmp.children[char];
+    insert(word: string, score: number = 0): void {
+        if (word.length === 0) return;
+        
+        if (!this.wordScores.has(word)) {
+            this.wordScores.set(word, score);
         }
+        
+        let tmp: TrieNode = this.root;
+        let path: TrieNode[] = [this.root];
+        const normalized: string = word.toLowerCase();
+
+        for (const char of normalized) { 
+            if (!tmp.children[char]) { 
+                tmp.children[char] = new TrieNode(tmp.prefix + char);
+            }
+            tmp = tmp.children[char];
+            path.push(tmp);
+        }
+
         tmp.endOfWord = true;
         tmp.words.add(word);
+
+        for (const node of path) {
+            this.updateCache(node, word);
+        }
     }
 
     remove(word: string): void {
@@ -32,6 +57,23 @@ export default class Trie {
         this.removeWord(this.root, word, 0);
     }
 
+    bumpWordScore(word: string, delta: number = 1): void {
+        this.setWordScore(word, (this.getScore(word) ?? 0) + delta);
+    }
+
+    setWordScore(word: string, score: number): void {
+        if (!this.wordScores.has(word)) return;
+
+        this.wordScores.set(word, score);
+
+        const path: TrieNode[] | null = this.getPathNodes(word);
+        if (!path) return;
+
+        for (const node of path) {
+            this.refreshCache(node);
+        }
+    }
+
     findWordsWithPrefix(prefix: string, limit: number = Infinity, caseSensitive: boolean) : string[] {
         let tmp: TrieNode = this.root;
         const normalized: string = prefix.toLowerCase();
@@ -41,52 +83,195 @@ export default class Trie {
         }
 
         const results: string[] = [];
-        this.collectWords(tmp, results, limit);
+        const seen: Set<string> = new Set();
 
-        if (caseSensitive) {
-            return results.filter(word => word.startsWith(prefix));
+        for (const word of tmp.cache) {
+            if (results.length >= limit) break;
+            if (!this.shouldSuggestWord(prefix, word)) continue;
+            if (caseSensitive && !word.startsWith(prefix)) continue;
+            if (seen.has(word)) continue;
+
+            seen.add(word);
+            results.push(word);
+        }
+
+        
+        if (results.length < limit) {
+            this.collectWords(tmp, results, seen, prefix, limit, caseSensitive);
         }
 
         return results;
     }
 
+    getScore(word: string): number | undefined {
+        return this.wordScores.get(word);
+    }
+
     private removeWord(node: TrieNode, word: string, index: number): boolean {
         // Base Case: currently at the end of word
         if (index === word.length) {
             if (!node.words.has(word)) return false;
 
             node.words.delete(word);
+            this.wordScores.delete(word);
             if (node.words.size === 0) node.endOfWord = false;
 
-            return !node.endOfWord && Object.keys(node.children).length === 0;
+            this.refreshCache(node);
+
+            return !node.endOfWord && this.hasNoChildren(node);
         }
 
         const char: string = word[index].toLowerCase();
-        if (!node.children[char]) return false;
+        const child: TrieNode | undefined = node.children[char];
+        if (!child) return false;
 
-        const canDeleteChild: boolean = this.removeWord(node.children[char], word, index + 1);
-        if (canDeleteChild) {
-            delete node.children[char];
-            return !node.endOfWord && Object.keys(node.children).length === 0;
-        }
-        return false;
+        const canDeleteChild: boolean = this.removeWord(child, word, index + 1);
+        if (canDeleteChild) delete node.children[char];
+
+        this.refreshCache(node);
+        return !node.endOfWord && this.hasNoChildren(node);
     }
 
-    private collectWords(node: TrieNode, results: string[], limit: number): void {
+    private collectWords(
+        node: TrieNode, 
+        results: string[], 
+        seen: Set<string>, 
+        prefix: string, 
+        limit: number, 
+        caseSensitive: boolean
+    ): void {
         // Base Case: result list has reached search limit
         if (results.length >= limit) return;
 
-        if (node.endOfWord) { 
-            for (const word of node.words) {
-                results.push(word);
+        if (node.endOfWord) {
+            const terminalWords: string[] = Array.from(node.words)
+                .sort((a: string, b: string) => this.compareWords(a, b));
+
+            for (const word of terminalWords) {
                 if (results.length >= limit) return;
+                if (!this.shouldSuggestWord(prefix, word)) continue;
+                if (seen.has(word)) continue;
+                if (caseSensitive && !word.startsWith(prefix)) continue;
+
+                seen.add(word);
+                results.push(word);
             }
         }
 
-        for (const char in node.children) {
-            this.collectWords(node.children[char], results, limit);
+        const childEntries: Array<[string, TrieNode]> = Object.entries(node.children);
+        childEntries.sort(([, childA], [, childB]) => {
+            const bestA: string | undefined = this.getBestCandidate(childA);
+            const bestB: string | undefined = this.getBestCandidate(childB);
+
+            if (!bestA && !bestB) return 0;
+            if (!bestA) return 1;
+            if (!bestB) return -1;
+
+            return this.compareWords(bestA, bestB);
+        });
+
+        for (const [, child] of childEntries) {
+            if (results.length >= limit) return;
+            this.collectWords(child, results, seen, prefix, limit, caseSensitive);
         }
     }
+
+    private getPathNodes(word: string): TrieNode[] | null {
+        let tmp: TrieNode = this.root;
+        const normalized: string = word.toLowerCase();
+        const path: TrieNode[] = [this.root];
+
+        for (const char of normalized) {
+            if (!tmp.children[char]) return null;
+            tmp = tmp.children[char];
+            path.push(tmp);
+        }
+
+        return path;
+    }
+
+    private updateCache(node: TrieNode, word: string): void {
+        const existingIndex: number = node.cache.indexOf(word);
+        if (existingIndex !== -1) {
+            node.cache.splice(existingIndex, 1);
+        }
+
+        node.cache.push(word);
+        node.cache.sort((a: string, b: string) => this.compareWords(a, b));
+
+        if (node.cache.length > this.maxCacheSize) {
+            node.cache.length = this.maxCacheSize;
+        }
+    }
+
+    private refreshCache(node: TrieNode): void {
+        const candidates: Set<string> = new Set();
+
+        for (const word of node.words) {
+            if (this.wordScores.has(word) && this.shouldCacheWord(node, word)) {
+                candidates.add(word);
+            }
+        }
+
+        for (const child of Object.values(node.children)) {
+            for (const word of child.words) {
+                if (this.wordScores.has(word) && this.shouldCacheWord(node, word)) {
+                    candidates.add(word);
+                }
+            }
+
+            for (const word of child.cache) {
+                if (this.wordScores.has(word) && this.shouldCacheWord(node, word)) {
+                    candidates.add(word);
+                }
+            }
+        }
+
+        node.cache = Array.from(candidates)
+            .sort((a: string, b: string) => this.compareWords(a, b)) 
+            .slice(0, this.maxCacheSize);
+    }
+
+    private getBestCandidate(node: TrieNode): string | undefined {
+        let best: string | undefined;
+
+        for (const word of node.words) {
+            if (!this.wordScores.has(word)) continue;
+            if (!best || this.compareWords(word, best) < 0) {
+                best = word;
+            }
+        }
+
+        for (const word of node.cache) {
+            if (!this.wordScores.has(word)) continue;
+            if (!best || this.compareWords(word, best) < 0) {
+                best = word;
+            }
+        }
+
+        return best;
+    }
+
+    private shouldCacheWord(node: TrieNode, word: string): boolean {
+        return word !== node.prefix;
+    }
+
+    private shouldSuggestWord(prefix: string, word: string): boolean {
+        return word !== prefix;
+    }
+
+    private compareWords(a: string, b: string): number {
+        const scoreA: number = this.getScore(a) ?? 0;
+        const scoreB: number = this.getScore(b) ?? 0;
+
+        if (scoreA !== scoreB) return scoreB - scoreA;
+        return a.localeCompare(b);
+    }
+
+    private hasNoChildren(node: TrieNode): boolean {
+        for (const _ in node.children) return false;
+        return true;
+    }
 }
 
 export { Trie };

src/dictionary/words.ts

@@ -7470,7 +7470,6 @@ const DEFAULT_WORDS : string[] = [
     "cheers",
     "idiot",
     "morbid",
-    "e-mail",
     "outcome",
     "gilt",
     "coldness",

src/main.ts

@@ -1,9 +1,10 @@
 import { Editor, EditorPosition, Plugin, Menu, MenuItem, Notice } from 'obsidian';
 import { TASettingsTab, DEFAULT_SETTINGS, TASettings, DictionaryFile } from './settings/settings';
 import { createTAUI, destroyTAUI, updateSuggestions } from './settings/ui';
-import { DUPLICATE_FILE, DUPLICATE_WORD, EMPTY_FILE, inCodeBlock, NOT_OK, OK, Result, stringInWordOrBeforePunctuation, wordBeforeString } from './utils';
-import { DEFAULT_TRIE } from './dictionary/dictionary';
+import { inCodeBlock, Result, stringInWordOrBeforePunctuation, wordBeforeString } from './utils';
+import { DEFAULT_DICTIONARIES, Dictionary } from './dictionary/dictionary';
 import { Trie } from './dictionary/trie';
+import { ACCEPTED_SUGGESTION_BUMP, CUSTOM_WORD_SCORE, DEFAULT_WORD_SCORE, DUPLICATE_FILE, DUPLICATE_WORD, EMPTY_FILE, IMPORT_WORD_SCORE, NOT_OK, OK } from './constants';
 
 export default class TAPlugin extends Plugin {
 	settings: TASettings;
@@ -33,11 +34,31 @@ export default class TAPlugin extends Plugin {
 	}
 
 	async loadWordTrie() {
-		this.wordTrie = DEFAULT_TRIE;
-		this.settings.customDict.forEach(word => this.wordTrie.insert(word));
-		this.settings.dictFiles.forEach(file => 
-			file.words.forEach(word => this.wordTrie.insert(word))
-		);
+		this.wordTrie = new Trie(this.settings.maxSuggestions);
+
+		DEFAULT_DICTIONARIES.forEach((dict: Dictionary) => {
+			dict.words.forEach((word: string) => {
+				const score: number = this.settings.wordScores[word] ?? DEFAULT_WORD_SCORE;
+				this.settings.wordScores[word] = score;
+				this.wordTrie.insert(word, score);
+			});
+		});
+
+		this.settings.customDict.forEach((word: string) => {
+			const score: number = this.settings.wordScores[word] ?? CUSTOM_WORD_SCORE;
+			this.settings.wordScores[word] = score;
+			this.wordTrie.insert(word, score);
+		});
+
+		this.settings.dictFiles.forEach((file: DictionaryFile) => {
+			file.words.forEach((word: string) => {
+				const score: number = this.settings.wordScores[word] ?? IMPORT_WORD_SCORE;
+				this.settings.wordScores[word] = score;
+				this.wordTrie.insert(word)
+			});
+		});
+
+		await this.saveSettings();
 	}
 
 	async saveSettings() {
@@ -73,11 +94,14 @@ export default class TAPlugin extends Plugin {
 		}
 
 		const word: string = match[1];
-		const suggestions: string[] = this.wordTrie.findWordsWithPrefix(word, this.settings.maxSuggestions, this.settings.caseSensitive)
-			.filter((w: string) => w !== word);
+		const suggestions: string[] = this.wordTrie.findWordsWithPrefix(
+			word, 
+			this.settings.maxSuggestions, 
+			this.settings.caseSensitive
+		);
 
 		if (suggestions.length > 0) {
-			updateSuggestions(suggestions, editor, this.settings);
+			updateSuggestions(suggestions, editor, this);
 		} else {
 			destroyTAUI();
 		}
@@ -122,7 +146,9 @@ export default class TAPlugin extends Plugin {
 
 			if (evt.key === 'Enter' || evt.key === 'Tab') {
 				const selected = active || items[0];
-				if (selected) selected.dispatchEvent(new Event('mousedown'));
+				if (selected) {
+					selected.dispatchEvent(new Event('mousedown'));
+				}
 				destroyTAUI();
 				return;
 			}
@@ -137,12 +163,22 @@ export default class TAPlugin extends Plugin {
 		if (this.settingsTab) this.settingsTab.display();
 	}
 
+	async bumpAcceptedWord(word: string): Promise<void> {
+		this.wordTrie.bumpWordScore(word, ACCEPTED_SUGGESTION_BUMP);
+		this.settings.wordScores[word] = this.wordTrie.getScore(word);
+		await this.saveSettings();
+	}
+
 	async addCustomWord(word: string): Promise<Result> {
 		if (!word) return EMPTY_FILE;
 		if (this.settings.customDict.includes(word)) return DUPLICATE_WORD;
 
 		this.settings.customDict.push(word);
-		this.wordTrie.insert(word);
+
+		const score: number = this.settings.wordScores[word] ?? CUSTOM_WORD_SCORE;
+		this.settings.wordScores[word] = score;
+
+		this.wordTrie.insert(word, score);
 		await this.saveSettings();
 		return OK;
 	}
@@ -152,8 +188,10 @@ export default class TAPlugin extends Plugin {
 		if (!word) return NOT_OK;
 
 		this.settings.customDict.splice(index, 1);
+
 		if (!this.wordExistsInAnotherDict(word, -1)) {
 			this.wordTrie.remove(word);
+			delete this.settings.wordScores[word];
 		}
 		
 		await this.saveSettings();
@@ -164,6 +202,7 @@ export default class TAPlugin extends Plugin {
 		this.settings.customDict.forEach((word: string) => {
 			if (!this.wordExistsInAnotherDict(word, -1)) {
 				this.wordTrie.remove(word);
+				delete this.settings.wordScores[word];
 			}
 		});
 		
@@ -185,7 +224,11 @@ export default class TAPlugin extends Plugin {
 
 		if (words.length === 0) return { result: EMPTY_FILE };
 
-		words.forEach((word: string) => this.wordTrie.insert(word));
+		words.forEach((word: string) => {
+			const score: number = this.settings.wordScores[word] ?? IMPORT_WORD_SCORE;
+			this.settings.wordScores[word] = score;
+			this.wordTrie.insert(word, score);
+		});
 
 		const dictFile: DictionaryFile = {
 			filename: file.name,
@@ -204,6 +247,7 @@ export default class TAPlugin extends Plugin {
 		dictFile.words.forEach((word: string) => {
 			if (!this.wordExistsInAnotherDict(word, index)) {
 				this.wordTrie.remove(word);
+				delete this.settings.wordScores[word];
 			}
 		});
 		this.settings.dictFiles.splice(index, 1);

src/settings/settings.ts

@@ -1,7 +1,8 @@
 import { PluginSettingTab, App, Setting, Notice, DropdownComponent, ToggleComponent, SliderComponent, TextComponent, ButtonComponent } from 'obsidian';
 import type TAPlugin from 'src/main';
 import { destroyTAUI } from './ui';
-import { DUPLICATE_FILE, DUPLICATE_WORD, EMPTY_FILE, OK, Result } from 'src/utils';
+import { DUPLICATE_FILE, DUPLICATE_WORD, EMPTY_FILE } from 'src/constants';
+import { Result } from 'src/utils';
 
 export interface DictionaryFile {
     filename: string;
@@ -16,6 +17,7 @@ export interface TASettings {
     caseSensitive: boolean;
     customDict: string[];
     dictFiles: DictionaryFile[];
+    wordScores: Record<string, number | undefined>;
 }
 
 export const DEFAULT_SETTINGS: TASettings = {
@@ -25,7 +27,8 @@ export const DEFAULT_SETTINGS: TASettings = {
     addSpace: false,
     caseSensitive: true,
     customDict: [],
-    dictFiles: []
+    dictFiles: [],
+    wordScores: {}
 }
 
 export class TASettingsTab extends PluginSettingTab {

src/settings/ui.ts

@@ -1,7 +1,7 @@
 import { ViewPlugin, ViewUpdate, PluginValue } from '@codemirror/view';
 import { Editor } from 'obsidian';
-import { TASettings } from './settings';
 import { stringInWordOrBeforePunctuation, wordBeforeString } from 'src/utils';
+import TAPlugin from 'src/main';
 
 let dropdownEl: HTMLUListElement | null = null;
 
@@ -57,7 +57,7 @@ export function destroyTAUI() {
     dropdownEl = null;
 }
 
-export function updateSuggestions(suggestions: string[], editor: Editor, settings: TASettings) {
+export function updateSuggestions(suggestions: string[], editor: Editor, plugin: TAPlugin) {
     destroyTAUI();
 
     const cm = (editor as CodeMirrorEditor).cm;
@@ -68,22 +68,26 @@ export function updateSuggestions(suggestions: string[], editor: Editor, setting
     dropdownEl = createEl('ul');
     dropdownEl!.className = 'autocomplete-dropdown';
 
-    suggestions.forEach(suggestion => {
+    suggestions.forEach((suggestion: string) => {
         const li = createEl('li');
         li.textContent = suggestion;
-        li.addEventListener('mousedown', () => {
+        li.addEventListener('mousedown', async () => {
             const cursor = editor.getCursor();
             const line = editor.getLine(cursor.line);
             const beforeCursor = line.substring(0, cursor.ch);
             const match = wordBeforeString(beforeCursor);
+
             if (match) {
-                if (settings.addSpace) suggestion += " ";
+                let insert: string = `${suggestion}`; 
+                if (plugin.settings.addSpace) insert += " ";
                 editor.replaceRange(
-                    suggestion,
+                    insert,
                     { line: cursor.line, ch: cursor.ch - match[1].length }, // start position (line, position in line)
                     cursor // end position (line, position in line)
                 );
+                await plugin.bumpAcceptedWord(suggestion);
             }
+
             destroyTAUI();
         });
         dropdownEl?.appendChild(li);

src/utils.ts

@@ -5,13 +5,6 @@ export interface Result {
     reason?: string;
 };
 
-export const OK: Result = { ok: true };
-export const NOT_OK: Result = { ok: false };
-export const EMPTY_FILE: Result = { ok: false, reason: 'emptyFile' };
-export const DUPLICATE_FILE: Result = { ok: false, reason: 'duplicateFile' };
-export const EMPTY_WORD: Result = { ok: false, reason: 'emptyWord' };
-export const DUPLICATE_WORD: Result = { ok: false, reason: 'duplicateWord'};
-
 export const inCodeBlock = (editor: Editor, cursor: EditorPosition): boolean => {
     const lines = editor.getValue().split('\n');
     let inCodeBlock = false;

versions.json

@@ -1,3 +1,4 @@
 {
-	"1.0.0": "0.15.0"
-}
+	"1.0.0": "0.15.0",
+	"1.2.2": "1.8.10"
+}