From 334619086becbc4bb32b8c764017aea38f890984 Mon Sep 17 00:00:00 2001
From: Aleksandras Kostarevas <alex@futo.org>
Date: Fri, 13 Oct 2023 18:32:57 +0300
Subject: [PATCH] Trim context and fix case for misautocorections

---
 .../inputmethod/latin/xlm/LanguageModel.java  | 62 ++++++++++++++++---
 1 file changed, 54 insertions(+), 8 deletions(-)

diff --git a/java/src/org/futo/inputmethod/latin/xlm/LanguageModel.java b/java/src/org/futo/inputmethod/latin/xlm/LanguageModel.java
index 691ae3371..7868f4df4 100644
--- a/java/src/org/futo/inputmethod/latin/xlm/LanguageModel.java
+++ b/java/src/org/futo/inputmethod/latin/xlm/LanguageModel.java
@@ -19,6 +19,7 @@ import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Locale;
+import java.util.function.IntPredicate;
 
 
 public class LanguageModel extends Dictionary {
@@ -109,8 +110,16 @@ public class LanguageModel extends Dictionary {
             float weightForLocale,
             float[] inOutWeightOfLangModelVsSpatialModel
     ) {
-        if (mNativeState == 0) return null;
-        if (initThread != null && initThread.isAlive()) return null;
+        Log.d("LanguageModel", "getSuggestions called");
+
+        if (mNativeState == 0) {
+            Log.d("LanguageModel", "Exiting becuase mNativeState == 0");
+            return null;
+        }
+        if (initThread != null && initThread.isAlive()){
+            Log.d("LanguageModel", "Exiting because initThread");
+            return null;
+        }
 
         final InputPointers inputPointers = composedData.mInputPointers;
         final boolean isGesture = composedData.mIsBatchMode;
@@ -122,6 +131,26 @@ public class LanguageModel extends Dictionary {
             context = ngramContext.fullContext.trim();
         }
 
+        // Trim the context
+        while(context.length() > 128) {
+            if(context.contains("\n")) {
+                context = context.substring(context.indexOf("\n") + 1).trim();
+            }else if(context.contains(".") || context.contains("?") || context.contains("!")) {
+                int v = Arrays.stream(
+                        new int[]{
+                                context.indexOf("."),
+                                context.indexOf("?"),
+                                context.indexOf("!")
+                        }).filter(i -> i != -1).min().orElse(-1);
+
+                if(v == -1) break; // should be unreachable
+
+                context = context.substring(v + 1).trim();
+            } else {
+                break;
+            }
+        }
+
         String partialWord = composedData.mTypedWord;
 
         if(!partialWord.isEmpty() && context.endsWith(partialWord)) {
@@ -154,17 +183,32 @@ public class LanguageModel extends Dictionary {
         final ArrayList<SuggestedWords.SuggestedWordInfo> suggestions = new ArrayList<>();
 
         int kind = SuggestedWords.SuggestedWordInfo.KIND_PREDICTION;
+
+        boolean mustNotAutocorrect = false;
+        for(int i=0; i<maxResults; i++) {
+            if (outStrings[i] == null) continue;
+            if(!partialWord.isEmpty() && partialWord.trim().equalsIgnoreCase(outStrings[i].trim())) {
+                // If this prediction matches the partial word ignoring case, and this is the top
+                // prediction, then we can break.
+                // Otherwise, we cannot autocorrect to the top prediction, as it does not match the
+                // partial word but one of the top ones does.
+                if(i == 0) {
+                    break;
+                } else {
+                    mustNotAutocorrect = true;
+                }
+            }
+        }
+
+        if(!partialWord.isEmpty() && !mustNotAutocorrect) {
+            kind = SuggestedWords.SuggestedWordInfo.KIND_WHITELIST | SuggestedWords.SuggestedWordInfo.KIND_FLAG_APPROPRIATE_FOR_AUTO_CORRECTION;
+        }
+
         for(int i=0; i<maxResults; i++) {
             if(outStrings[i] == null) continue;
 
             String word = outStrings[i].trim();
 
-            if(!partialWord.isEmpty()) {
-                kind = SuggestedWords.SuggestedWordInfo.KIND_WHITELIST | SuggestedWords.SuggestedWordInfo.KIND_FLAG_APPROPRIATE_FOR_AUTO_CORRECTION;
-            }
-
-            Log.d("LanguageModel", "probability for word [" + word + "] is 100 * " + String.valueOf(outProbabilities[i]));
-
             suggestions.add(new SuggestedWords.SuggestedWordInfo( word, context, (int)(outProbabilities[i] * 100.0f), kind, this, 0, 0 ));
         }
 
@@ -178,6 +222,8 @@ public class LanguageModel extends Dictionary {
             }
         }
 
+        Log.d("LanguageModel", "returning " + String.valueOf(suggestions.size()) + " suggestions");
+
         return suggestions;
     }