3 files changed, 102 insertions, 23 deletions

diff --git a/srcs/juloo.keyboard2/CurrentlyTypedWord.java b/srcs/juloo.keyboard2/CurrentlyTypedWord.java
index 2c7066d..b27c380 100644
--- a/srcs/juloo.keyboard2/CurrentlyTypedWord.java
+++ b/srcs/juloo.keyboard2/CurrentlyTypedWord.java
@@ -1,9 +1,11 @@
 package juloo.keyboard2;
 
+import android.os.Build.VERSION;
 import android.os.Handler;
 import android.view.KeyEvent;
 import android.view.inputmethod.EditorInfo;
 import android.view.inputmethod.InputConnection;
+import android.view.inputmethod.SurroundingText;
 import java.util.List;
 
 /** Keep track of the word being typed. This also tracks whether the selection
@@ -24,11 +26,14 @@ public final class CurrentlyTypedWord
   /** Used to avoid concurrent refreshes in [delayed_refresh()]. */
   boolean _refresh_pending = false;
 
-  /** The estimated cursor position. Used to avoid expensive IPC calls when the
-      typed word can be estimated locally with [typed]. When the cursor
-      position gets out of sync, the text before the cursor is queried again to
-      the editor. */
+  /** The estimated cursor position in code points. Used to avoid expensive IPC
+      calls when the typed word can be estimated locally with [typed]. When the
+      cursor position gets out of sync, the text before the cursor is queried
+      again to the editor. */
   int _cursor;
+  /** The cursor position within the current word relative to the end of the
+      word in chars. Equal to [0] when the cursor is at the end of the word. */
+  int _w_cursor;
 
   public CurrentlyTypedWord(Handler h, Callback cb)
   {
@@ -46,6 +51,12 @@ public final class CurrentlyTypedWord
     return _has_selection;
   }
 
+  /** The cursor position relative to the end of the word. */
+  public int cursor_relative()
+  {
+    return _w_cursor;
+  }
+
   public void started(Config conf, InputConnection ic)
   {
     _ic = ic;
@@ -53,8 +64,12 @@ public final class CurrentlyTypedWord
     EditorConfig e = conf.editor_config;
     _has_selection = e.initial_sel_start != e.initial_sel_end;
     _cursor = e.initial_sel_start;
+    _w_cursor = 0;
     if (!_has_selection)
+    {
       set_current_word(e.initial_text_before_cursor);
+      _w_cursor = -append_chars(e.initial_text_after_cursor); 
+    }
   }
 
   public void typed(String s)
@@ -70,12 +85,22 @@ public final class CurrentlyTypedWord
   {
     // Avoid the expensive [refresh_current_word] call when [typed] was called
     // before.
-    boolean new_has_sel = newSelStart != newSelEnd;
-    if (!_enabled || (newSelStart == _cursor && new_has_sel == _has_selection))
+    if (!_enabled)
       return;
-    _has_selection = new_has_sel;
-    _cursor = newSelStart;
-    refresh_current_word();
+    boolean new_has_sel = newSelStart != newSelEnd;
+    if (new_has_sel || _has_selection) // Selection was on or is now on.
+    {
+      _cursor = newSelStart;
+      _has_selection = new_has_sel;
+      refresh_current_word();
+    }
+    else if (newSelStart != _cursor)
+    {
+      _cursor = newSelStart;
+      _w_cursor += newSelStart - oldSelStart;
+      if (_w_cursor < -_w.length() || _w_cursor > 0)
+        refresh_current_word();
+    }
   }
 
   public void event_sent(int code, int meta)
@@ -87,31 +112,64 @@ public final class CurrentlyTypedWord
 
   void callback()
   {
-    _callback.currently_typed_word(_w.toString());
+    String w = _w.toString();
+    _callback.currently_typed_word(w);
   }
 
   /** Estimate the currently typed word after [chars] has been typed. */
-  void type_chars(String s)
+  void type_chars(CharSequence s, int start, int end)
   {
-    int len = s.length();
-    for (int i = 0; i < len;)
+    int insert_start = 0;
+    // Iterate over code points as that's the unit of [_cursor].
+    for (int i = start; i < end;)
     {
-      int c = s.codePointAt(i);
-      if (Character.isLetter(c))
-        _w.appendCodePoint(c);
-      else
-        _w.setLength(0);
+      int c = Character.codePointAt(s, i);
+      i += Character.charCount(c);
       _cursor++;
+      if (!Character.isLetter(c))
+        insert_start = i;
+    }
+    if (insert_start > 0)
+      _w.setLength(0);
+    _w.insert(_w.length() + _w_cursor, s, insert_start, end);
+  }
+
+  void type_chars(CharSequence s)
+  {
+    type_chars(s, 0, s.length());
+  }
+
+  /** Append chars to the current word without moving the cursor. Return the
+      number of characters that were added in the current word. */
+  int append_chars(CharSequence s, int start, int end)
+  {
+    int i = start;
+    while (i < end)
+    {
+      int c = Character.codePointAt(s, i);
+      if (!Character.isLetter(c))
+        break;
+      _w.appendCodePoint(c);
       i += Character.charCount(c);
     }
+    return i - start;
+  }
+
+  int append_chars(CharSequence s)
+  {
+    return append_chars(s, 0, s.length());
   }
 
   /** Refresh the current word by immediately querying the editor. */
   void refresh_current_word()
   {
+    Logs.debug("Refresh current word");
     _refresh_pending = false;
+    _w_cursor = 0;
     if (_has_selection)
       set_current_word("");
+    else if (VERSION.SDK_INT >= 31)
+      set_current_word(_ic.getSurroundingText(20, 20, 0));
     else
       set_current_word(_ic.getTextBeforeCursor(20, 0));
   }
@@ -128,6 +186,21 @@ public final class CurrentlyTypedWord
     callback();
   }
 
+  /** Like above but take the text after the cursor into account. */
+  void set_current_word(SurroundingText st)
+  {
+    _w.setLength(0);
+    if (st == null)
+      return;
+    int saved_cursor = _cursor;
+    int st_sel = st.getSelectionStart();
+    CharSequence st_text = st.getText();
+    type_chars(st_text, 0, st_sel);
+    _w_cursor = -append_chars(st_text, st_sel, st_text.length());
+    _cursor = saved_cursor;
+    callback();
+  }
+
   /** Wait some time to let the editor finishes reacting to changes and call
       [refresh_current_word]. */
   void delayed_refresh()
diff --git a/srcs/juloo.keyboard2/EditorConfig.java b/srcs/juloo.keyboard2/EditorConfig.java
index a198775..fa6a9dd 100644
--- a/srcs/juloo.keyboard2/EditorConfig.java
+++ b/srcs/juloo.keyboard2/EditorConfig.java
@@ -32,6 +32,7 @@ public final class EditorConfig
 
   /** CurrentlyTypedWord. */
   public CharSequence initial_text_before_cursor = null; // Might be [null].
+  public CharSequence initial_text_after_cursor = null; // Might be [null].
   public int initial_sel_start;
   public int initial_sel_end;
 
@@ -92,7 +93,10 @@ public final class EditorConfig
     caps_initially_updated = caps_should_update_state(info);
     /* CurrentlyTypedWord */
     if (VERSION.SDK_INT >= 30)
-      initial_text_before_cursor = info.getInitialTextBeforeCursor(10, 0);
+    {
+      initial_text_before_cursor = info.getInitialTextBeforeCursor(20, 0);
+      initial_text_after_cursor = info.getInitialTextAfterCursor(20, 0);
+    }
     initial_sel_start = info.initialSelStart;
     initial_sel_end = info.initialSelEnd;
     /* Suggestions */
diff --git a/srcs/juloo.keyboard2/KeyEventHandler.java b/srcs/juloo.keyboard2/KeyEventHandler.java
index fd283df..8691cda 100644
--- a/srcs/juloo.keyboard2/KeyEventHandler.java
+++ b/srcs/juloo.keyboard2/KeyEventHandler.java
@@ -129,7 +129,8 @@ public final class KeyEventHandler
   public void suggestion_entered(String text)
   {
     String old = _typedword.get();
-    replace_text_before_cursor(old.length(), text + " ");
+    int cur_rel = _typedword.cursor_relative();
+    replace_surrounding_text(old.length() + cur_rel, -cur_rel, text + " ");
     last_replaced_word = old;
     last_replacement_word_len = text.length() + 1;
   }
@@ -256,13 +257,14 @@ public final class KeyEventHandler
     clear_space_bar_state();
   }
 
-  void replace_text_before_cursor(int remove_length, String new_text)
+  void replace_surrounding_text(int remove_before, int remove_after,
+      String new_text)
   {
     InputConnection conn = _recv.getCurrentInputConnection();
     if (conn == null)
       return;
     conn.beginBatchEdit();
-    conn.deleteSurroundingText(remove_length, 0);
+    conn.deleteSurroundingText(remove_before, remove_after);
     conn.commitText(new_text, 1);
     conn.endBatchEdit();
   }
@@ -540,7 +542,7 @@ public final class KeyEventHandler
   {
     if (last_replaced_word != null)
     {
-      replace_text_before_cursor(last_replacement_word_len,
+      replace_surrounding_text(last_replacement_word_len, 0,
           last_replaced_word + " ");
       last_replaced_word = null;
     }