impr(polyglot): Add new structure (wordsetMap) to PolyglotWordset to get words from any language by uniform distribution (@IliyaZinoviev)

frontend/src/ts/test/funbox/funbox-functions.ts

@@ -1,4 +1,4 @@
-import { FunboxWordsFrequency, Wordset } from "../wordset";
+import { FunboxWordsFrequency, WordsetPick, Wordset } from "../wordset";
 import * as GetText from "../../utils/generate";
 import Config, { setConfig, toggleFunbox } from "../../config";
 import * as Misc from "../../utils/misc";
@@ -122,7 +122,7 @@ class PseudolangWordGenerator extends Wordset {
     }
   }
 
-  public override randomWord(): string {
+  public override randomWord(): WordsetPick {
     let word = "";
     for (;;) {
       const prefix = word.slice(-prefixSize);
@@ -141,24 +141,57 @@ class PseudolangWordGenerator extends Wordset {
       }
       word += nextChar;
     }
-    return word;
+    return { word };
   }
 }
 
 export class PolyglotWordset extends Wordset {
-  public wordsWithLanguage: Map<string, Language>;
-  public languageProperties: Map<Language, JSONData.LanguageProperties>;
+  readonly wordsetMap: Map<Language, Wordset>;
+  readonly languageProperties: Map<Language, JSONData.LanguageProperties>;
+  readonly langs: Language[];
 
   constructor(
-    wordsWithLanguage: Map<string, Language>,
+    words: string[],
+    wordsetMap: Map<Language, Wordset>,
     languageProperties: Map<Language, JSONData.LanguageProperties>,
   ) {
-    // build and shuffle the word array
-    const wordArray = Array.from(wordsWithLanguage.keys());
-    Arrays.shuffle(wordArray);
-    super(wordArray);
-    this.wordsWithLanguage = wordsWithLanguage;
+    super(words);
     this.languageProperties = languageProperties;
+    this.langs = Array.from(languageProperties.keys());
+    this.wordsetMap = wordsetMap;
+    this.resetIndexes();
+    this.length = words.length;
+  }
+
+  override resetIndexes(): void {
+    this.wordsetMap.forEach((ws) => {
+      ws.resetIndexes();
+    });
+  }
+
+  private pickLang(): Language {
+    const index = Math.floor(Math.random() * this.langs.length);
+    return this.langs[index] as Language;
+  }
+
+  private getWordsetAndLang(): { wordset: Wordset; language: Language } {
+    const language = this.pickLang();
+    return { wordset: this.wordsetMap.get(language) as Wordset, language };
+  }
+
+  override randomWord(mode: FunboxWordsFrequency): WordsetPick {
+    const { wordset, language } = this.getWordsetAndLang();
+    return { word: wordset.randomWord(mode).word, language };
+  }
+
+  override shuffledWord(): WordsetPick {
+    const { wordset, language } = this.getWordsetAndLang();
+    return { word: wordset.shuffledWord().word, language };
+  }
+
+  override nextWord(): WordsetPick {
+    const { wordset, language } = this.getWordsetAndLang();
+    return { word: wordset.nextWord().word, language };
   }
 }
 
@@ -700,7 +733,7 @@ const list: Partial<Record<FunboxName, FunboxFunctions>> = {
         }),
       );
 
-      const languages = (await Promise.all(promises)).filter(
+      const languages: LanguageObject[] = (await Promise.all(promises)).filter(
         (lang): lang is LanguageObject => lang !== null,
       );
 
@@ -752,25 +785,29 @@ const list: Partial<Record<FunboxName, FunboxFunctions>> = {
       }
 
       // build languageProperties
-      const languageProperties = new Map(
-        languages.map((lang) => [
-          lang.name,
-          {
-            noLazyMode: lang.noLazyMode,
-            ligatures: lang.ligatures,
-            rightToLeft: lang.rightToLeft,
-            additionalAccents: lang.additionalAccents,
-          },
-        ]),
-      );
-
-      const wordsWithLanguage = new Map(
-        languages.flatMap((lang) =>
-          lang.words.map((word) => [word, lang.name]),
-        ),
-      );
-
-      return new PolyglotWordset(wordsWithLanguage, languageProperties);
+      const languageProperties: Map<Language, JSONData.LanguageProperties> =
+        new Map(
+          languages.map((lang) => [
+            lang.name,
+            {
+              noLazyMode: lang.noLazyMode,
+              ligatures: lang.ligatures,
+              rightToLeft: lang.rightToLeft,
+              additionalAccents: lang.additionalAccents,
+            },
+          ]),
+        );
+      // build wordsetMap and words
+      const wordsetMap = new Map<Language, Wordset>();
+      let end = 0;
+      const words: string[] = [];
+      for (const lang of languages) {
+        const start = end;
+        end += lang.words.length;
+        words.push(...lang.words);
+        wordsetMap.set(lang.name, new Wordset(words.slice(start, end)));
+      }
+      return new PolyglotWordset(words, wordsetMap, languageProperties);
     },
   },
 };

frontend/src/ts/test/weak-spot.ts

@@ -63,7 +63,7 @@ export function getWord(wordset: Wordset): string {
   let highScore;
   let randomWord = "";
   for (let i = 0; i < wordSamples; i++) {
-    const newWord = wordset.randomWord("normal");
+    const newWord = wordset.randomWord("normal").word;
     const newScore = score(newWord);
     if (i === 0 || highScore === undefined || newScore > highScore) {
       randomWord = newWord;

frontend/src/ts/test/words-generator.ts

@@ -26,7 +26,7 @@ import { WordGenError } from "../utils/word-gen-error";
 
 import { showLoaderBar, hideLoaderBar } from "../signals/loader-bar";
 import { PolyglotWordset } from "./funbox/funbox-functions";
-import { LanguageObject } from "@monkeytype/schemas/languages";
+import { Language, LanguageObject } from "@monkeytype/schemas/languages";
 
 //pin implementation
 const random = Math.random;
@@ -385,25 +385,21 @@ async function applyBritishEnglishToWord(
   return await BritishEnglish.replace(word, previousWord);
 }
 
-function applyLazyModeToWord(word: string, language: LanguageObject): string {
-  // polyglot mode, use the word's actual language
-  if (currentWordset && currentWordset instanceof PolyglotWordset) {
-    const langName = currentWordset.wordsWithLanguage.get(word);
-    const langProps = langName
-      ? currentWordset.languageProperties.get(langName)
-      : undefined;
-    const allowLazyMode =
-      (langProps && !langProps.noLazyMode) === true || Config.mode === "custom";
-    if (Config.lazyMode && allowLazyMode && langProps) {
-      word = LazyMode.replaceAccents(word, langProps.additionalAccents);
-    }
-    return word;
-  }
+function applyLazyModeToWord(
+  word: string,
+  language: LanguageObject,
+  polyglotLang?: Language,
+): string {
+  const langProps =
+    polyglotLang && currentWordset instanceof PolyglotWordset
+      ? currentWordset.languageProperties.get(polyglotLang)
+      : language;
+
+  if (!langProps) return word;
 
-  // normal mode
-  const allowLazyMode = !language.noLazyMode || Config.mode === "custom";
+  const allowLazyMode = !langProps.noLazyMode || Config.mode === "custom";
   if (Config.lazyMode && allowLazyMode) {
-    word = LazyMode.replaceAccents(word, language.additionalAccents);
+    word = LazyMode.replaceAccents(word, langProps.additionalAccents);
   }
   return word;
 }
@@ -508,9 +504,8 @@ async function getQuoteWordList(
     // because it will be reversed again in the generateWords function
     if (wordOrder === "reverse") {
       return currentWordset.words.reverse();
-    } else {
-      return currentWordset.words;
     }
+    return currentWordset.words;
   }
   const languageToGet = language.name.startsWith("swiss_german")
     ? "german"
@@ -655,17 +650,13 @@ export async function generateWords(
 
   const funbox = findSingleActiveFunboxWithFunction("withWords");
   if (funbox) {
-    const result = await funbox.functions.withWords(wordList);
+    currentWordset = await funbox.functions.withWords(wordList);
     // PolyglotWordset if polyglot otherwise Wordset
-    if (result instanceof PolyglotWordset) {
-      const polyglotResult = result;
-      currentWordset = polyglotResult;
+    if (currentWordset instanceof PolyglotWordset) {
       // set allLigatures if any language in languageProperties has ligatures true
       ret.allLigatures = Array.from(
-        polyglotResult.languageProperties.values(),
+        currentWordset.languageProperties.values(),
       ).some((props) => !!props.ligatures);
-    } else {
-      currentWordset = result;
     }
   } else {
     currentWordset = await withWords(wordList);
@@ -803,7 +794,7 @@ export async function getNextWord(
   }
 
   const funboxFrequency = getFunboxWordsFrequency() ?? "normal";
-  let randomWord = currentWordset.randomWord(funboxFrequency);
+  let pick = currentWordset.randomWord(funboxFrequency);
   const previousWordRaw = previousWord.replace(/[.?!":\-,]/g, "").toLowerCase();
   const previousWord2Raw = previousWord2
     ?.replace(/[.?!":\-,']/g, "")
@@ -813,80 +804,85 @@ export async function getNextWord(
     const funboxSection = await getFunboxSection();
 
     if (Config.mode === "quote") {
-      randomWord = currentWordset.nextWord();
+      pick = currentWordset.nextWord();
     } else if (Config.mode === "custom" && CustomText.getMode() === "repeat") {
-      randomWord = currentWordset.nextWord();
+      pick = currentWordset.nextWord();
     } else if (
       Config.mode === "custom" &&
       CustomText.getMode() === "random" &&
       (currentWordset.length < 4 || PractiseWords.before.mode !== null)
     ) {
-      randomWord = currentWordset.randomWord(funboxFrequency);
+      pick = currentWordset.randomWord(funboxFrequency);
     } else if (Config.mode === "custom" && CustomText.getMode() === "shuffle") {
-      randomWord = currentWordset.shuffledWord();
+      pick = currentWordset.shuffledWord();
     } else if (
       Config.mode === "custom" &&
       CustomText.getLimitMode() === "section"
     ) {
-      randomWord = currentWordset.randomWord(funboxFrequency);
+      pick = currentWordset.randomWord(funboxFrequency);
 
       const previousSection = Arrays.nthElementFromArray(sectionHistory, -1);
       const previousSection2 = Arrays.nthElementFromArray(sectionHistory, -2);
 
       let regenerationCount = 0;
       while (
         regenerationCount < 100 &&
-        (previousSection === randomWord || previousSection2 === randomWord)
+        (previousSection === pick.word || previousSection2 === pick.word)
       ) {
         regenerationCount++;
-        randomWord = currentWordset.randomWord(funboxFrequency);
+        pick = currentWordset.randomWord(funboxFrequency);
       }
     } else if (isCurrentlyUsingFunboxSection) {
-      randomWord = funboxSection.join(" ");
+      pick.word = funboxSection.join(" ");
     } else {
       let regenarationCount = 0; //infinite loop emergency stop button
-      let firstAfterSplit = (randomWord.split(" ")[0] as string).toLowerCase();
+      let firstAfterSplit = (pick.word.split(" ")[0] as string).toLowerCase();
       let firstAfterSplitLazy = applyLazyModeToWord(
         firstAfterSplit,
         currentLanguage,
+        pick.language,
       );
       while (
         regenarationCount < 100 &&
         (previousWordRaw === firstAfterSplitLazy ||
           previousWord2Raw === firstAfterSplitLazy ||
           (Config.mode !== "custom" &&
             !Config.punctuation &&
-            randomWord === "I") ||
+            pick.word === "I") ||
           (Config.mode !== "custom" &&
             !Config.punctuation &&
             !Config.language.startsWith("code") &&
-            /[-=_+[\]{};'\\:"|,./<>?]/i.test(randomWord)) ||
+            /[-=_+[\]{};'\\:"|,./<>?]/i.test(pick.word)) ||
           (Config.mode !== "custom" &&
             !Config.numbers &&
-            /[0-9]/i.test(randomWord)))
+            /[0-9]/i.test(pick.word)))
       ) {
         regenarationCount++;
-        randomWord = currentWordset.randomWord(funboxFrequency);
-        firstAfterSplit = randomWord.split(" ")[0] as string;
+        pick = currentWordset.randomWord(funboxFrequency);
+        firstAfterSplit = pick.word.split(" ")[0] as string;
         firstAfterSplitLazy = applyLazyModeToWord(
           firstAfterSplit,
           currentLanguage,
+          pick.language,
         );
       }
     }
-    randomWord = randomWord.replace(/ +/g, " ");
-    randomWord = randomWord.replace(/(^ )|( $)/g, "");
+    pick.word = pick.word.replace(/ +/g, " ");
+    pick.word = pick.word.replace(/(^ )|( $)/g, "");
 
-    randomWord = getFunboxWord(randomWord, wordIndex, currentWordset);
+    pick.word = getFunboxWord(pick.word, wordIndex, currentWordset);
 
-    currentSection = [...randomWord.split(" ")];
-    sectionHistory.push(randomWord);
-    randomWord = currentSection.shift() as string;
+    currentSection = [...pick.word.split(" ")];
+    sectionHistory.push(pick.word);
+    pick.word = currentSection.shift() as string;
     sectionIndex++;
   } else {
-    randomWord = currentSection.shift() as string;
+    pick.word = currentSection.shift() as string;
   }
 
+  let randomWord = pick.word;
+  const randomWordLanguage = pick.language ?? Config.language;
+
   if (randomWord === undefined) {
     throw new WordGenError("Random word is undefined");
   }
@@ -900,10 +896,6 @@ export async function getNextWord(
   }
 
   const usingFunboxWithGetWord = isFunboxActiveWithFunction("getWord");
-  const randomWordLanguage =
-    (currentWordset instanceof PolyglotWordset
-      ? currentWordset.wordsWithLanguage.get(randomWord)
-      : Config.language) ?? Config.language; // Fall back to Config language if per-word language is unavailable
 
   if (
     Config.mode !== "custom" &&
@@ -922,9 +914,9 @@ export async function getNextWord(
 
   randomWord = randomWord.replace(/ +/gm, " ");
   randomWord = randomWord.replace(/(^ )|( $)/gm, "");
-  randomWord = applyLazyModeToWord(randomWord, currentLanguage);
+  randomWord = applyLazyModeToWord(randomWord, currentLanguage, pick.language);
 
-  if (Config.language.startsWith("swiss_german")) {
+  if (randomWordLanguage.startsWith("swiss_german")) {
     randomWord = randomWord.replace(/ß/g, "ss");
   }