Merge branch 'feat/web/edit-path-ranking' into fix/web/wordbreak-shift-at-head

web/src/engine/predictive-text/worker-thread/src/main/correction/context-tokenization.ts

@@ -167,11 +167,11 @@ export class ContextTokenization {
     const tokenDistribution = alignedTransformDistribution.map((entry) => {
       const remap: Map<number, ProbabilityMass<Transform>> = new Map();
 
-      for(let pair of entry.sample.entries()) {
+      for(const pair of entry.sample.entries()) {
         remap.set(pair[0], {
           sample: pair[1],
           p: entry.p
-        })
+        });
       }
 
       return remap;

web/src/test/auto/headless/engine/predictive-text/worker-thread/context/context-state.tests.ts

@@ -261,6 +261,38 @@ describe('ContextState', () => {
       assert.equal(newContextMatch.final.tokenization.alignment.tailTokenShift, 2);
     });
 
+    it("properly matches and aligns when whitespace before final empty token is extended", function() {
+      let existingContext = {
+        left: "an apple a day keeps the doctor ", startOfBuffer: true, endOfBuffer: true
+      };
+      let transform = {
+        insert: ' ',
+        deleteLeft: 0
+      }
+      let rawTokens = ["an", " ", "apple", " ", "a", " ", "day", " ", "keeps", " ", "the", " ", "doctor", "  ", ""];
+
+      let baseState = new ContextState(existingContext, plainModel);
+      let newContextMatch = baseState.analyzeTransition(existingContext, toWrapperDistribution(transform));
+      assert.isNotNull(newContextMatch?.final);
+      assert.deepEqual(newContextMatch.final.tokenization.tokens.map(token => token.exampleInput), rawTokens);
+      // We want to preserve the added whitespace when predicting a token that follows after it.
+      assert.deepEqual(newContextMatch.preservationTransform, { insert: ' ', deleteLeft: 0 });
+
+      // The 'wordbreak' transform
+      let state = newContextMatch?.final;
+      assert.isNotEmpty(state.tokenization.tokens[state.tokenization.tokens.length - 2].searchSpace.inputSequence);
+      assert.deepEqual(
+        state.tokenization.tokens[state.tokenization.tokens.length - 1].searchSpace.inputSequence,
+        [[{ sample: {insert: '', deleteLeft: 0}, p: 1 }]]
+      );
+
+      if(!newContextMatch.final.tokenization.alignment.canAlign) {
+        assert.fail("context alignment failed");
+      }
+      assert.equal(newContextMatch.final.tokenization.alignment.leadTokenShift, 0);
+      assert.equal(newContextMatch.final.tokenization.alignment.tailTokenShift, 0);
+    });
+
     it("properly matches and aligns when a 'wordbreak' is removed via backspace", function() {
       let existingContext = {
         left: "an apple a day keeps the doctor ", startOfBuffer: true, endOfBuffer: true

web/src/test/auto/headless/engine/predictive-text/worker-thread/context/context-tokenization.tests.ts

@@ -147,7 +147,38 @@ describe('ContextTokenization', function() {
           tailTokenShift: 0
         },
         plainModel,
-        [{ sample: inputTransformMap, p: 1}]
+        [{ sample: inputTransformMap, p: 1 }]
+      );
+
+      assert.isOk(tokenization);
+      assert.equal(tokenization.tokens.length, targetTokens.length);
+      assert.deepEqual(tokenization.tokens.map((t) => ({text: t.exampleInput, isWhitespace: t.isWhitespace})),
+        targetTokens
+      );
+    });
+
+    it('merges new whitespace character added to last whitespace token if tail is empty', () => {
+      const baseTokens = ['an', ' ', 'apple', ' ', 'a', ' ', 'day', ' ', ''];
+      const baseTokenization = new ContextTokenization(baseTokens.map(t => toToken(t)), null);
+
+      const targetTokens = ['an', ' ', 'apple', ' ', 'a', ' ', 'day', '  ', ''].map((t) => (
+        {text: t, isWhitespace: t != '' && t.trim() == ''}
+      ));
+      const inputTransformMap: Map<number, Transform> = new Map();
+      inputTransformMap.set(-1, { insert: ' ', deleteLeft: 0 });
+      inputTransformMap.set( 0, { insert: '',  deleteLeft: 0 });
+
+      const tokenization = baseTokenization.transitionTo(
+        targetTokens, {
+          canAlign: true,
+          leadTokenShift: 0,
+          leadEditLength: 0,
+          matchLength: 7,
+          tailEditLength: 2,
+          tailTokenShift: 0
+        },
+        plainModel,
+        [{ sample: inputTransformMap, p: 1 }]
       );
 
       assert.isOk(tokenization);

web/src/test/auto/headless/engine/predictive-text/worker-thread/transform-tokenization.tests.ts

@@ -353,6 +353,33 @@ describe('tokenizeTransform', () => {
       assert.equal(result.size, 3);
       assert.deepEqual(result, expectedMap);
     });
+
+    it('properly places extra whitespaces on preceding whitespace token', () => {
+      const context = {
+        left: 'do it properly ', // 'do', ' ', 'it', ' ', 'properly', ' ', ''
+        right: '',
+        startOfBuffer: true,
+        endOfBuffer: true
+      };
+
+      // Adjacent whitespace entries are generally merged into a single blob.
+      const editTransform = {
+        insert: ' ',  // Should be combined with the final ' ', not the tail ''.
+        deleteLeft: 0
+      };
+
+      const result = tokenizeTransform(
+        defaultTokenize,
+        context,
+        editTransform
+      );
+
+      const expectedMap = new Map<number, Transform>();
+      expectedMap.set(-1, { insert: ' ', deleteLeft: 0 });
+      expectedMap.set(0, { insert: '', deleteLeft: 0 });
+      assert.equal(result.size, 2);
+      assert.deepEqual(result, expectedMap);
+    });
   });
 
   describe('with mocked dictionary-based wordbreaking', () => {