From 82ba375fb55bde4e474f68cd91343d314ddc5dce Mon Sep 17 00:00:00 2001
From: Eddie Antonio Santos <easantos@ualberta.ca>
Date: Fri, 18 Jan 2019 14:23:16 -0700
Subject: [PATCH] Enable word list model in the top-level API.

---
 common/predictive-text/message.d.ts                  | 12 ++++++++++--
 .../in_browser/cases/worker-wordlist-integration.js  |  2 +-
 common/predictive-text/worker/index.ts               |  2 ++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/common/predictive-text/message.d.ts b/common/predictive-text/message.d.ts
index d128ac8f09..188bb78efa 100644
--- a/common/predictive-text/message.d.ts
+++ b/common/predictive-text/message.d.ts
@@ -100,12 +100,20 @@ interface Capabilities {
 }
 
 /**
- * TODO: discriminated union of different model types.
+ * Describes what kind of model to instantiate.
  */
-interface ModelDescription {
+type ModelDescription = DummyModelDescription | WordListModelDescription;
+interface DummyModelDescription {
   type: 'dummy';
   futureSuggestions?: Suggestion[][];
 }
+interface WordListModelDescription {
+  type: 'wordlist';
+  /**
+   * Words to predict, one word per entry in the array.
+   */
+  wordlist: string[];
+}
 
 /**
  * Configuration of the LMLayer, sent back to the keyboard.
diff --git a/common/predictive-text/unit_tests/in_browser/cases/worker-wordlist-integration.js b/common/predictive-text/unit_tests/in_browser/cases/worker-wordlist-integration.js
index e3ab03ff6c..3e3a86ba1a 100644
--- a/common/predictive-text/unit_tests/in_browser/cases/worker-wordlist-integration.js
+++ b/common/predictive-text/unit_tests/in_browser/cases/worker-wordlist-integration.js
@@ -28,7 +28,7 @@ describe('LMLayer using the word list model', function () {
       }).then(function (suggestions) {
         // TODO: not super sure what to assert here!
         assert.isAtLeast(suggestions.length, EXPECTED_SUGGESTIONS);
-        return lmLayer.predict(type('t'), emptyContext());
+        return lmLayer.predict(type('t'), atEndOfBuffer(''));
       }).then(function (suggestions) {
         assert.isAtLeast(suggestions.length, EXPECTED_SUGGESTIONS);
         return lmLayer.predict(type('h'), atEndOfBuffer('t'));
diff --git a/common/predictive-text/worker/index.ts b/common/predictive-text/worker/index.ts
index 5bd3b03837..fc4d00b830 100644
--- a/common/predictive-text/worker/index.ts
+++ b/common/predictive-text/worker/index.ts
@@ -145,6 +145,8 @@ class LMLayerWorker {
       model = new LMLayerWorker.models.DummyModel(capabilities, {
         futureSuggestions: desc.futureSuggestions
       });
+    } else if (desc.type === 'wordlist') {
+      model = new LMLayerWorker.models.WordListModel(capabilities, desc.wordlist);
     } else {
       throw new Error('Invalid model');
     }