From ebb6fa6782abfe4fe90578e0d49ebd32629dd6c5 Mon Sep 17 00:00:00 2001 From: c-degni Date: Sat, 14 Mar 2026 19:05:18 -0400 Subject: [PATCH] Frequency based suggestions --- src/constants.ts | 13 +++++++ src/dictionary/dictionary.ts | 14 +++++--- src/dictionary/trie.ts | 51 ++++++++++++++++++-------- src/dictionary/words.ts | 1 - src/main.ts | 70 +++++++++++++++++++++++++++++------- src/settings/settings.ts | 7 ++-- src/settings/ui.ts | 16 +++++---- src/utils.ts | 7 ---- 8 files changed, 131 insertions(+), 48 deletions(-) create mode 100644 src/constants.ts diff --git a/src/constants.ts b/src/constants.ts new file mode 100644 index 0000000..2ab5f4e --- /dev/null +++ b/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'}; \ No newline at end of file diff --git a/src/dictionary/dictionary.ts b/src/dictionary/dictionary.ts index 0b8de1f..b0dd98b 100644 --- a/src/dictionary/dictionary.ts +++ b/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 }; \ No newline at end of file +export const DEFAULT_DICTIONARIES: Dictionary[] = [ + { name: 'default', words: DEFAULT_WORDS, baseScore: 1 }, + { name: 'computer-science', words: CS_WORDS, baseScore: 1 }, +]; \ No newline at end of file diff --git a/src/dictionary/trie.ts b/src/dictionary/trie.ts index c01cbe9..6ccefb0 100644 --- a/src/dictionary/trie.ts +++ b/src/dictionary/trie.ts @@ -58,16 +58,7 @@ export default class Trie { } bumpWordScore(word: string, delta: number = 1): void { - if (!this.wordScores.has(word)) return; - - this.wordScores.set(word, (this.getScore(word) ?? 0) + delta); - - const path: TrieNode[] | null = this.getPathNodes(word); - if (!path) return; - - for (const node of path) { - this.refreshCache(node); - } + this.setWordScore(word, (this.getScore(word) ?? 0) + delta); } setWordScore(word: string, score: number): void { @@ -96,6 +87,7 @@ export default class Trie { 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; @@ -157,6 +149,7 @@ export default class Trie { 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; @@ -167,8 +160,8 @@ export default class Trie { const childEntries: Array<[string, TrieNode]> = Object.entries(node.children); childEntries.sort(([, childA], [, childB]) => { - const bestA: string | undefined = childA.cache[0]; - const bestB: string | undefined = childB.cache[0]; + const bestA: string | undefined = this.getBestCandidate(childA); + const bestB: string | undefined = this.getBestCandidate(childB); if (!bestA && !bestB) return 0; if (!bestA) return 1; @@ -220,8 +213,14 @@ export default class Trie { } } - for (const char in node.children) { - for (const word of node.children[char].cache) { + 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); } @@ -233,10 +232,34 @@ export default class Trie { .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; diff --git a/src/dictionary/words.ts b/src/dictionary/words.ts index 691bc44..c51b10f 100644 --- a/src/dictionary/words.ts +++ b/src/dictionary/words.ts @@ -7470,7 +7470,6 @@ const DEFAULT_WORDS : string[] = [ "cheers", "idiot", "morbid", - "e-mail", "outcome", "gilt", "coldness", diff --git a/src/main.ts b/src/main.ts index d68866b..b024762 100644 --- a/src/main.ts +++ b/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 { + this.wordTrie.bumpWordScore(word, ACCEPTED_SUGGESTION_BUMP); + this.settings.wordScores[word] = this.wordTrie.getScore(word); + await this.saveSettings(); + } + async addCustomWord(word: string): Promise { 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); diff --git a/src/settings/settings.ts b/src/settings/settings.ts index 23d17e1..0f5afcb 100644 --- a/src/settings/settings.ts +++ b/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; } 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 { diff --git a/src/settings/ui.ts b/src/settings/ui.ts index 87f9c47..1a2bbaf 100644 --- a/src/settings/ui.ts +++ b/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); diff --git a/src/utils.ts b/src/utils.ts index 1e22189..cf77064 100644 --- a/src/utils.ts +++ b/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;