fix(android): preserve inline styles during IME composition

android/src/main/java/com/swmansion/enriched/textinput/EnrichedTextInputConnectionWrapper.kt

@@ -1,11 +1,15 @@
 package com.swmansion.enriched.textinput
 
+import android.text.Editable
 import android.view.KeyEvent
+import android.view.inputmethod.BaseInputConnection
 import android.view.inputmethod.InputConnection
 import android.view.inputmethod.InputConnectionWrapper
 import com.facebook.react.bridge.ReactContext
 import com.facebook.react.uimanager.UIManagerHelper
+import com.swmansion.enriched.common.spans.interfaces.EnrichedInlineSpan
 import com.swmansion.enriched.textinput.events.OnInputKeyPressEvent
+import com.swmansion.enriched.textinput.spans.EnrichedSpans
 
 // This class is based on the implementation from Facebook React Native to provide 'onKeyPress' API on android.
 // Original source:
@@ -16,9 +20,33 @@ class EnrichedTextInputConnectionWrapper(
   private val editText: EnrichedTextInputView,
   private val experimentalSynchronousEvents: Boolean,
 ) : InputConnectionWrapper(target, false) {
+  private data class InlineSpanSnapshot(
+    val style: String,
+    val start: Int,
+    val end: Int,
+  )
+
+  private data class ComposingTextSnapshot(
+    val start: Int,
+    val end: Int,
+    val editableLength: Int,
+    val text: String,
+    val inlineSpans: List<InlineSpanSnapshot>,
+  )
+
   private var isBatchEdit = false
+  private var shouldPreserveInlineStyles = false
   private var key: String? = null
 
+  fun preserveInlineStylesForCurrentComposition() {
+    val editable = editText.text
+    val composingStart = editable?.let(BaseInputConnection::getComposingSpanStart) ?: -1
+    val composingEnd = editable?.let(BaseInputConnection::getComposingSpanEnd) ?: -1
+    shouldPreserveInlineStyles =
+      composingStart >= 0 &&
+      composingEnd > composingStart
+  }
+
   override fun beginBatchEdit(): Boolean {
     isBatchEdit = true
     return super.beginBatchEdit()
@@ -39,9 +67,23 @@ class EnrichedTextInputConnectionWrapper(
   ): Boolean {
     val previousSelectionStart = editText.selectionStart
     val previousSelectionEnd = editText.selectionEnd
+    val composingTextSnapshot =
+      if (shouldPreserveInlineStyles) {
+        captureComposingText()
+      } else {
+        null
+      }
+    if (shouldPreserveInlineStyles && composingTextSnapshot == null) {
+      shouldPreserveInlineStyles = false
+    }
 
     val consumed = super.setComposingText(text, newCursorPosition)
 
+    if (consumed && composingTextSnapshot != null) {
+      restoreInlineSpans(composingTextSnapshot)
+    }
+    clearPreservationIfCompositionFinished()
+
     val currentSelectionStart = editText.selectionStart
     val noPreviousSelection = previousSelectionStart == previousSelectionEnd
     val cursorDidNotMove = currentSelectionStart == previousSelectionStart
@@ -61,6 +103,149 @@ class EnrichedTextInputConnectionWrapper(
     return consumed
   }
 
+  private fun captureComposingText(): ComposingTextSnapshot? {
+    val editable = editText.text ?: return null
+    val composingStart = BaseInputConnection.getComposingSpanStart(editable)
+    val composingEnd = BaseInputConnection.getComposingSpanEnd(editable)
+    if (composingStart < 0 || composingEnd <= composingStart) return null
+
+    return ComposingTextSnapshot(
+      start = composingStart,
+      end = composingEnd,
+      editableLength = editable.length,
+      text = editable.substring(composingStart, composingEnd),
+      inlineSpans = captureInlineSpans(editable, composingStart, composingEnd),
+    )
+  }
+
+  private fun captureInlineSpans(
+    editable: Editable,
+    composingStart: Int,
+    composingEnd: Int,
+  ): List<InlineSpanSnapshot> =
+    editable
+      .getSpans(composingStart, composingEnd, EnrichedInlineSpan::class.java)
+      .mapNotNull { span ->
+        val style =
+          EnrichedSpans.inlineSpans.entries
+            .firstOrNull { (_, config) -> config.clazz.isInstance(span) }
+            ?.key
+            ?: return@mapNotNull null
+        val start = editable.getSpanStart(span).coerceAtLeast(composingStart)
+        val end = editable.getSpanEnd(span).coerceAtMost(composingEnd)
+
+        if (start < end) {
+          InlineSpanSnapshot(style, start - composingStart, end - composingStart)
+        } else {
+          null
+        }
+      }
+
+  private fun clearPreservationIfCompositionFinished() {
+    val editable = editText.text
+    if (
+      editable == null ||
+      BaseInputConnection.getComposingSpanStart(editable) < 0 ||
+      BaseInputConnection.getComposingSpanEnd(editable) < 0
+    ) {
+      shouldPreserveInlineStyles = false
+    }
+  }
+
+  private fun restoreInlineSpans(snapshot: ComposingTextSnapshot) {
+    val editable = editText.text ?: return
+    val previousComposingLength = snapshot.end - snapshot.start
+    val currentComposingLength =
+      editable.length - (snapshot.editableLength - previousComposingLength)
+    val currentComposingStart = snapshot.start.coerceAtMost(editable.length)
+    val currentComposingEnd =
+      (currentComposingStart + currentComposingLength)
+        .coerceAtLeast(currentComposingStart)
+        .coerceAtMost(editable.length)
+    if (currentComposingStart >= currentComposingEnd) return
+
+    val currentComposingText = editable.substring(currentComposingStart, currentComposingEnd)
+    val commonPrefixLength = commonPrefixLength(snapshot.text, currentComposingText)
+    val commonSuffixLength =
+      commonSuffixLength(
+        snapshot.text,
+        currentComposingText,
+        commonPrefixLength,
+      )
+
+    for (inlineSpan in snapshot.inlineSpans) {
+      restoreSnapshotIntersection(
+        inlineSpan,
+        oldRangeStart = 0,
+        oldRangeEnd = commonPrefixLength,
+        newRangeStart = 0,
+        composingStart = currentComposingStart,
+      )
+
+      val previousSuffixStart = snapshot.text.length - commonSuffixLength
+      val currentSuffixStart = currentComposingText.length - commonSuffixLength
+      restoreSnapshotIntersection(
+        inlineSpan,
+        oldRangeStart = previousSuffixStart,
+        oldRangeEnd = snapshot.text.length,
+        newRangeStart = currentSuffixStart,
+        composingStart = currentComposingStart,
+      )
+    }
+  }
+
+  private fun restoreSnapshotIntersection(
+    snapshot: InlineSpanSnapshot,
+    oldRangeStart: Int,
+    oldRangeEnd: Int,
+    newRangeStart: Int,
+    composingStart: Int,
+  ) {
+    val intersectionStart = snapshot.start.coerceAtLeast(oldRangeStart)
+    val intersectionEnd = snapshot.end.coerceAtMost(oldRangeEnd)
+    if (intersectionStart >= intersectionEnd) return
+
+    val mappedStart = newRangeStart + intersectionStart - oldRangeStart
+    val mappedEnd = newRangeStart + intersectionEnd - oldRangeStart
+    editText.inlineStyles?.restoreStyleOnRange(
+      snapshot.style,
+      composingStart + mappedStart,
+      composingStart + mappedEnd,
+    )
+  }
+
+  private fun commonPrefixLength(
+    previousText: String,
+    currentText: String,
+  ): Int {
+    val maximum = minOf(previousText.length, currentText.length)
+    var length = 0
+    while (length < maximum && previousText[length] == currentText[length]) {
+      length++
+    }
+    return length
+  }
+
+  private fun commonSuffixLength(
+    previousText: String,
+    currentText: String,
+    commonPrefixLength: Int,
+  ): Int {
+    val maximum =
+      minOf(
+        previousText.length - commonPrefixLength,
+        currentText.length - commonPrefixLength,
+      )
+    var length = 0
+    while (
+      length < maximum &&
+      previousText[previousText.lastIndex - length] == currentText[currentText.lastIndex - length]
+    ) {
+      length++
+    }
+    return length
+  }
+
   override fun commitText(
     text: CharSequence,
     newCursorPosition: Int,
@@ -73,7 +258,25 @@ class EnrichedTextInputConnectionWrapper(
       }
       dispatchKeyEventOrEnqueue(inputKey)
     }
-    return super.commitText(text, newCursorPosition)
+
+    val composingTextSnapshot =
+      if (shouldPreserveInlineStyles) {
+        captureComposingText()
+      } else {
+        null
+      }
+    val consumed = super.commitText(text, newCursorPosition)
+    if (consumed && composingTextSnapshot != null) {
+      restoreInlineSpans(composingTextSnapshot)
+    }
+    shouldPreserveInlineStyles = false
+    return consumed
+  }
+
+  override fun finishComposingText(): Boolean {
+    val consumed = super.finishComposingText()
+    shouldPreserveInlineStyles = false
+    return consumed
   }
 
   override fun deleteSurroundingText(

android/src/main/java/com/swmansion/enriched/textinput/EnrichedTextInputView.kt

@@ -146,6 +146,7 @@ class EnrichedTextInputView :
   private var defaultValueDirty: Boolean = false
 
   private var inputMethodManager: InputMethodManager? = null
+  private var enrichedInputConnection: EnrichedTextInputConnectionWrapper? = null
   private val spannableFactory = EnrichedTextInputSpannableFactory()
   private var contextMenuItems: List<Pair<Int, String>> = emptyList()
 
@@ -199,6 +200,7 @@ class EnrichedTextInputView :
         )
     }
 
+    enrichedInputConnection = inputConnection as? EnrichedTextInputConnectionWrapper
     return inputConnection
   }
 
@@ -933,6 +935,10 @@ class EnrichedTextInputView :
     val isValid = verifyStyle(name)
     if (!isValid) return
 
+    if (name in EnrichedSpans.inlineSpans) {
+      enrichedInputConnection?.preserveInlineStylesForCurrentComposition()
+    }
+
     val (rangeStart, rangeEnd) = getTargetRange(name)
 
     runAsATransaction {

android/src/main/java/com/swmansion/enriched/textinput/styles/InlineStyles.kt

@@ -161,6 +161,24 @@ class InlineStyles(
     setAndMergeSpans(spannable, type, start, end)
   }
 
+  fun restoreStyleOnRange(
+    name: String,
+    start: Int,
+    end: Int,
+  ) {
+    if (start >= end) return
+
+    val config = EnrichedSpans.inlineSpans[name] ?: return
+    val spannable = view.text as? Spannable ?: return
+    val spans = spannable.getSpans(start, end, config.clazz)
+
+    if (spans.any { spannable.getSpanStart(it) <= start && spannable.getSpanEnd(it) >= end }) {
+      return
+    }
+
+    setSpan(spannable, config.clazz, start, end)
+  }
+
   fun toggleStyle(name: String) {
     if (view.selection == null) return
     val (start, end) = view.selection.getInlineSelection()