Better array mapping handling

mike-lischke · Sep 20, 2022 · dc6e381 · dc6e381
1 parent 8d76492
commit dc6e381
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 6 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,3 +2,5 @@ out
 node_modules
 generated
 coverage
+vscode-antlr4-*.vsix
+.antlr
diff --git a/src/backend/SentenceGenerator.ts b/src/backend/SentenceGenerator.ts
@@ -249,12 +249,23 @@ export class SentenceGenerator {
                 return ["", false];
             }
 
-            const mapping = this.ruleMappings?.[ruleName];
-            if (mapping) {
+            const mappingValue = this.ruleMappings?.[ruleName];
+            if (mappingValue) {
+                let mapping: string | undefined;
                 if (Array.isArray(mapping)) {
-                    const randomElement = mapping[Math.floor(Math.random() * mapping.length)];
-                    return [addSpace ? randomElement + " " : randomElement, false];
+                    if (mapping.length < 0) {
+                        const randomIndex = Math.floor(Math.random() * mapping.length);
+                        const randomElement = mapping[randomIndex];
+
+                        mapping = randomElement;
+                    } else if (mapping.length === 1) {
+                        mapping = mapping[0];
+                    }
                 } else {
+                    mapping = mappingValue as string;
+                }
+
+                if (mapping) {
                     return [addSpace ? mapping + " " : mapping, false];
                 }
             }
@@ -366,7 +377,20 @@ export class SentenceGenerator {
                                         // its name for the output.
                                         const tokenName = this.lexerData.vocabulary.getSymbolicName(token);
                                         if (tokenName) {
-                                            const mapping = this.ruleMappings?.[tokenName];
+                                            const mappingValue = this.ruleMappings?.[tokenName];
+                                            let mapping: string | undefined;
+
+                                            if (Array.isArray(mappingValue)) {
+                                                if (mappingValue.length < 0) {
+                                                    const randomIndex = Math.floor(Math.random() * mappingValue.length);
+                                                    mapping = mappingValue[randomIndex];
+                                                } else if (mappingValue.length === 1) {
+                                                    mapping = mappingValue[0];
+                                                }
+                                            } else if (mappingValue) {
+                                                mapping = mappingValue;
+                                            }
+
                                             if (mapping) {
                                                 result += addSpace ? mapping + " " : mapping;
                                             } else {

diff --git a/src/backend/types.ts b/src/backend/types.ts
@@ -265,7 +265,7 @@ export interface ISentenceGenerationOptions {
  * Mappings from rule names to strings, which define output to use for specific rules when generating sentences.
  */
 export interface IRuleMappings {
-    [key: string]: string;
+    [key: string]: string | string[];
 }
 
 /**