fix(web): fix construction for suggestions emitted after left-deleting whitespace 🚂

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

@@ -730,7 +730,7 @@ export class ContextTokenization {
     return new ContextTokenization(
       tokenSequence,
       null,
-      determineTaillessTrueKeystroke(transitionEdge)
+      determineTaillessTrueKeystroke(transitionEdge.inputs[0].sample)
     );
   }
 }
@@ -1244,20 +1244,19 @@ export function assembleTransforms(stackedInserts: string[], stackedDeletes: num
  * Used to construct and represent the part of the incoming transform that does
  * not land as part of the final token in the resulting context.  This component
  * should be preserved by any suggestions that get applied.
- * @param tokenizationAnalysis Precomputed metadata about a potential transition
- * from a pre-transition context tokenization to a post-transition context
- * tokenization
+ * @param tokenizedInputs The precomputed tokenization for incoming inputs
+ * involved in a pre-transition context tokenization to a post-transition
+ * context tokenization.
  * @returns
  */
-export function determineTaillessTrueKeystroke(tokenizationAnalysis: TransitionEdge) {
+export function determineTaillessTrueKeystroke(tokenizedInput: Map<number, Transform>) {
   // undefined by default; we haven't yet determined if we're still affecting
   // the same token that was the tail in the previous tokenization state.
   let taillessTrueKeystroke: Transform;
 
   // If tokens were inserted, emit an empty transform; this prevents
   // suggestions from replacing the "current" token.
-  const bestTokenizedInput = tokenizationAnalysis.inputs[0].sample;
-  if(bestTokenizedInput.has(1)) {
+  if(tokenizedInput.has(1)) {
     // Sets a default transform that will be returned even if the main
     // transform body lies entirely within a new token.
     taillessTrueKeystroke = { insert: '', deleteLeft: 0 };
@@ -1266,26 +1265,45 @@ export function determineTaillessTrueKeystroke(tokenizationAnalysis: TransitionE
     // by the loop that follows, without fail.
   }
 
-  const transformKeys = [...tokenizationAnalysis.inputs[0].sample.keys()];
+  const transformKeys = [...tokenizedInput.keys()];
+  do {
+    const penultimateKey = transformKeys[transformKeys.length - 2];
+    const tailKey = transformKeys[transformKeys.length - 1];
+
+    const penultimateTransform = tokenizedInput.get(penultimateKey);
+    const tailTransform = tokenizedInput.get(tailKey);
+
+    // Do not treat pure-backspace transforms at the tail end of context as the
+    // transform applied to the suggestion if the input was tokenized; this
+    // scenario implies that a prior token is being edited instead.
+    if(TransformUtils.isBackspace(tailTransform) && transformKeys.length > 1) {
+      transformKeys.pop();
+      continue;
+    } else if(!penultimateTransform) {
+      break;
+    } else if(
+      // Erasing a single-char whitespace requires deletion of two tokens, the
+      // last of which is empty.  Check for this case and handle it accordingly
+      // as well.
+      TransformUtils.isEmpty(tailTransform) && TransformUtils.isBackspace(penultimateTransform)
+    ) {
+      transformKeys.pop();
+      continue;
+    } else {
+      break;
+    }
+  } while(true);
+
+  // Ignore the transform that applies to the suggestion-root token - it should
+  // contribute to the suggestion, rather than be a fixed, universally-applied
+  // constant.
   transformKeys.pop();
 
+  // If no inputs remain, that's fine - that means an empty transform applies to
+  // whatever token exists to the token indexed before the first input-key
+  // entry.
   for(let i of transformKeys) {
-    /*
-      * Thinking ahead to multitokenization:
-      *
-      * If what we have is not on the "true" tokenization, then... we need to
-      * do multitoken effects, right?  We're basing new suggestions based on a
-      * state that does not currently exist!  We'd need to enforce THAT state,
-      * *then* do the suggestion!
-      * - Which gets fun if we auto-apply such a case, as the new "true" tokenization
-      *   no longer results directly from the true input.
-      *
-      * If we give tokens unique IDs on first creation, we could backtrace to
-      * find the most recent common ancestor.
-      * - simple cases (same 'token', but different input transform lengths/effects)
-      *   will have the same prior token ID
-      */
-    const primaryInput = tokenizationAnalysis.inputs[0].sample.get(i);
+    const primaryInput = tokenizedInput.get(i);
     if(!taillessTrueKeystroke) {
       taillessTrueKeystroke = {...primaryInput};
     } else {

web/src/engine/predictive-text/worker-thread/src/main/predict-helpers.ts

@@ -336,10 +336,13 @@ export function determineSuggestionAlignment(
    */
   predictionContext: Context,
   /**
-   * The total number of characters to delete for generated suggestions
-   * in order to replace the prediction root token entirely.
+   * The total number of characters to delete from the token to be corrected.
    */
-  deleteLeft: number
+  correctionDeleteLeft: number
+  /**
+   * The number of characters deleted from tokens aside from the one being corrected.
+   */
+  committedDeleteLeft: number
 } {
   const transitionEdits = tokenization.transitionEdits;
   const context = transition.base.context;
@@ -361,18 +364,24 @@ export function determineSuggestionAlignment(
       // As the word/token being corrected/predicted didn't originally exist,
       // there's no part of it to 'replace'.  (Suggestions are applied to the
       // pre-transform state.)
-      deleteLeft: 0
+      correctionDeleteLeft: 0,
+      committedDeleteLeft: 0
     };
     // If the tokenized context length is shorter... sounds like a backspace (or similar).
-  } else if (transitionEdits?.removedOldTokens) {
+  } else if (transitionEdits?.removedOldTokens || TransformUtils.isBackspace(inputTransform)) {
     /* Ooh, we've dropped context here.  Almost certainly from a backspace or
      * similar effect.  Even if we drop multiple tokens... well, we know exactly
      * how many chars were actually deleted - `inputTransform.deleteLeft`. Since
      * we replace a word being corrected/predicted, we take length of the
      * remaining context's tail token in addition to however far was deleted to
      * reach that state.
      */
-    deleteLeft = KMWString.length(wordbreak(postContext)) + inputTransform.deleteLeft;
+    return {
+      predictionContext: models.applyTransform({...inputTransform, insert: ''}, context),
+      // Pre-apply delete-lefts, but do not include any inserted portion.
+      correctionDeleteLeft: KMWString.length(wordbreak(postContext)) - KMWString.length(inputTransform.insert),
+      committedDeleteLeft: inputTransform.deleteLeft
+    };
   } else {
     // Suggestions are applied to the pre-input context, so get the token's original length.
     // We're on the same token, so just delete its text for the replacement op.
@@ -385,7 +394,11 @@ export function determineSuggestionAlignment(
     deleteLeft = 0;
   }
 
-  return { predictionContext: context, deleteLeft };
+  return {
+    predictionContext: context,
+    correctionDeleteLeft: deleteLeft,
+    committedDeleteLeft: 0
+  };
 }
 
 /**
@@ -467,15 +480,15 @@ export function buildAndMapPredictions(
   // No matter the prediction, once we know the root of the prediction, we'll
   // always 'replace' the same amount of text.  We can handle this before the
   // big 'prediction root' loop.
-  const { predictionContext, deleteLeft } = determineSuggestionAlignment(transition, tokenization, model);
+  const { predictionContext, correctionDeleteLeft, committedDeleteLeft } = determineSuggestionAlignment(transition, tokenization, model);
 
   let correction = match.matchString;
   let rootCost = match.totalCost;
 
   // Replace the existing context with the correction.
   const correctionTransform: Transform = {
     insert: correction,  // insert correction string
-    deleteLeft: deleteLeft,
+    deleteLeft: correctionDeleteLeft,
     id: transition.transitionId // The correction should always be based on the most recent external transform/transcription ID.
   }
 
@@ -489,17 +502,9 @@ export function buildAndMapPredictions(
   let predictions = predictFromCorrections(model, [predictionRoot], predictionContext);
   predictions.forEach((entry) => {
     entry.preservationTransform = tokenization.taillessTrueKeystroke;
-    // // Will need an extra lookup layer if the suggestion is generated from within a cluster.
-    // entry.baseTokenization = transition.final.tokenizationSourceMap.get(tokenization);
+    entry.prediction.sample.transform.deleteLeft += committedDeleteLeft;
   });
 
-  // Backspaces that shorten a multi-codepoint whitespace token are not handled well by default.
-  // As a new empty token is placed at the end for such cases, we can detect and handle such cases.
-  const inputTransform = transition.inputDistribution?.[0].sample ?? { insert: '', deleteLeft: 0 };
-  if(tokenization.tokens.length > 1 && tokenization.tail.searchModule.codepointLength == 0 && inputTransform.deleteLeft > 0) {
-    predictions.forEach((p) => p.prediction.sample.transform.deleteLeft += inputTransform.deleteLeft);
-  }
-
   return predictions;
 }
 

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

@@ -28,7 +28,8 @@ import {
   models,
   TransitionEdge,
   SearchQuotientSpur,
-  traceInsertEdits
+  traceInsertEdits,
+  determineTaillessTrueKeystroke
 } from '@keymanapp/lm-worker/test-index';
 
 import Transform = LexicalModelTypes.Transform;
@@ -2538,4 +2539,111 @@ describe('ContextTokenization', function() {
       assert.deepEqual(results, expectedMap);
     });
   });
+
+  describe('determineTaillessTrueKeystroke', () => {
+    it('handles simple tail-token extensions correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(0, { insert: '', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.isNotOk(preservedTransform);
+    });
+
+    it('handles simple tail-terminating whitespace inputs correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(1, { insert: ' ', deleteLeft: 0 });
+      tokenizedInput.set(2, { insert: '', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.deepEqual(preservedTransform, {
+        insert: ' ',
+        deleteLeft: 0
+      });
+    });
+
+    it('handles simple tail-token char deletions correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(0, { insert: '', deleteLeft: 1 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.isNotOk(preservedTransform);
+    });
+
+    it('handles tail whitespace-token deletions correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(-1, { insert: '', deleteLeft: 1 });
+      tokenizedInput.set(0, { insert: '', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.isNotOk(preservedTransform);
+    });
+
+    it('handles multi-token insert with small delete correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(0, { insert: 'a', deleteLeft: 1 });
+      tokenizedInput.set(1, { insert: ' ', deleteLeft: 0 });
+      tokenizedInput.set(2, { insert: 'bc', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.deepEqual(preservedTransform, {
+        insert: 'a ',
+        deleteLeft: 1
+      });
+    });
+
+    it('handles multi-token delete with small insert correctly', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(-2, { insert: 'a', deleteLeft: 1 });
+      tokenizedInput.set(-1, { insert: '', deleteLeft: 1 });
+      tokenizedInput.set(0, { insert: '', deleteLeft: 1 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.isNotOk(preservedTransform);
+    });
+
+    it('handles multi-token insertion/deletion input correctly (1)', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(-2, { insert: 'a', deleteLeft: 1 });
+      tokenizedInput.set(-1, { insert: ' ', deleteLeft: 1 });
+      tokenizedInput.set(0, { insert: 'b', deleteLeft: 1 });
+      tokenizedInput.set(1, { insert: ' ', deleteLeft: 0 });
+      tokenizedInput.set(2, { insert: 'c', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.deepEqual(preservedTransform, {
+        insert: 'a b ',
+        deleteLeft: 3
+      });
+    });
+
+    it('handles multi-token insertion/deletion input correctly (2)', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(-4, { insert: 'a', deleteLeft: 1 });
+      tokenizedInput.set(-3, { insert: ' ', deleteLeft: 1 });
+      tokenizedInput.set(-2, { insert: 'b', deleteLeft: 1 });
+      tokenizedInput.set(-1, { insert: ' ', deleteLeft: 1 });
+      tokenizedInput.set(0, { insert: '', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.deepEqual(preservedTransform, {
+        insert: 'a b ',
+        deleteLeft: 4
+      });
+    });
+
+    it('handles multi-token insertion/deletion input correctly (3)', () => {
+      const tokenizedInput: Map<number, Transform> = new Map();
+      tokenizedInput.set(-4, { insert: 'a', deleteLeft: 1 });
+      tokenizedInput.set(-3, { insert: ' ', deleteLeft: 1 });
+      tokenizedInput.set(-2, { insert: 'b', deleteLeft: 1 });
+      tokenizedInput.set(-1, { insert: '', deleteLeft: 1 });
+      tokenizedInput.set(0, { insert: '', deleteLeft: 0 });
+
+      const preservedTransform = determineTaillessTrueKeystroke(tokenizedInput);
+      assert.deepEqual(preservedTransform, {
+        insert: 'a ',
+        deleteLeft: 2
+      });
+    });
+  });
 });

web/src/test/auto/headless/engine/predictive-text/worker-thread/prediction-helpers/determine-suggestion-alignment.tests.ts

@@ -51,7 +51,8 @@ describe('determineSuggestionAlignment', () => {
     const results = determineSuggestionAlignment(transition, transition.final.tokenization, plainCasedModel);
 
     assert.deepEqual(results.predictionContext, context);
-    assert.equal(results.deleteLeft, "techn".length);
+    assert.equal(results.correctionDeleteLeft, "techn".length /* does not include the deleted whitespace */);
+    assert.equal(results.committedDeleteLeft, 0);
   });
 
   it('handles extension of prior token after backspace', () => {
@@ -67,8 +68,14 @@ describe('determineSuggestionAlignment', () => {
     // transition, model
     const results = determineSuggestionAlignment(transition, transition.final.tokenization, plainCasedModel);
 
-    assert.deepEqual(results.predictionContext, context);
-    assert.equal(results.deleteLeft, "tech".length + 1 /* for the deleted whitespace */);
+    assert.deepEqual(results.predictionContext, {
+      ...context,
+      left: context.left.substring(0, context.left.length - 1),
+      right: '',
+      casingForm: undefined
+    });
+    assert.equal(results.correctionDeleteLeft, "tech".length /* does not include the deleted whitespace */);
+    assert.equal(results.committedDeleteLeft, 1 /* for the deleted whitespace */);
   });
 
   it('handles extension of prior token after complex input with delete-left', () => {
@@ -84,7 +91,13 @@ describe('determineSuggestionAlignment', () => {
     // transition, model
     const results = determineSuggestionAlignment(transition, transition.final.tokenization, plainCasedModel);
 
-    assert.deepEqual(results.predictionContext, context);
-    assert.equal(results.deleteLeft, "techn".length + 1 /* for the deleted whitespace */);
+    assert.deepEqual(results.predictionContext, {
+      ...context,
+      left: context.left.substring(0, context.left.length - 1),
+      right: '',
+      casingForm: undefined
+    });
+    assert.equal(results.correctionDeleteLeft, "tech".length /* does not include the deleted whitespace */);
+    assert.equal(results.committedDeleteLeft, 1 /* for the deleted whitespace */);
   });
 });