Preserve source text in prepared text layouts

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/fabric/FabricUIManager.java

@@ -706,13 +706,7 @@ public PreparedLayout prepareTextLayout(
   @UnstableReactNativeAPI
   public PreparedLayout reusePreparedLayoutWithNewReactTags(
       PreparedLayout preparedLayout, int[] reactTags) {
-    return new PreparedLayout(
-        preparedLayout.getLayout(),
-        preparedLayout.getMaximumNumberOfLines(),
-        preparedLayout.getVerticalOffset(),
-        reactTags,
-        preparedLayout.getTextBreakStrategy(),
-        preparedLayout.getJustificationMode());
+    return preparedLayout.copyWithReactTags(reactTags);
   }
 
   @AnyThread

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/text/MutableSpannableLayout.kt

@@ -12,14 +12,16 @@ import android.text.SpannableString
 import android.text.Spanned
 import com.facebook.react.common.annotations.UnstableReactNativeAPI
 import com.facebook.react.views.text.internal.span.StatefulSpan
+import java.util.IdentityHashMap
 
 /**
- * A delegating [Layout] subclass that clones the spannable text from [delegate] and replaces all
- * [StatefulSpan] instances with fresh clones. This gives each [PreparedLayoutTextView] independent
- * mutable span state (e.g. particle animation, dismiss state) even when the underlying [Layout] is
- * shared from a cache.
+ * A delegating [Layout] subclass that clones both the display-shaped text from [delegate] and its
+ * full source text. Corresponding [StatefulSpan] instances in both copies are replaced with the
+ * same fresh clone, while spans present in only one copy are cloned on demand. This gives each
+ * [PreparedLayoutTextView] independent mutable span state (e.g. particle animation, dismiss state)
+ * even when the underlying [Layout] is shared from a cache.
  *
- * The mutable [Spannable] can be useful for spans which affect display, but do not alter existing
+ * The mutable display text can be useful for spans which affect display, but do not alter existing
  * layout calculations.
  *
  * Line metrics are delegated to [delegate] so no expensive [StaticLayout] rebuild is needed.
@@ -30,6 +32,7 @@ internal class MutableSpannableLayout
 private constructor(
     private val delegate: Layout,
     clonedText: SpannableString,
+    val sourceText: Spanned,
 ) :
     Layout(
         clonedText,
@@ -43,22 +46,39 @@ private constructor(
   companion object {
     /** Returns a [MutableSpannableLayout] if [layout] contains stateful spans, else null. */
     @OptIn(UnstableReactNativeAPI::class)
-    fun createIfNeeded(layout: Layout): MutableSpannableLayout? {
-      val spanned = layout.text as? Spanned ?: return null
-      val statefulSpans = spanned.getSpans(0, spanned.length, StatefulSpan::class.java)
+    fun createIfNeeded(layout: Layout, sourceText: Spanned): MutableSpannableLayout? {
+      val layoutText = layout.text as? Spanned ?: return null
+      val statefulSpans = layoutText.getSpans(0, layoutText.length, StatefulSpan::class.java)
       if (statefulSpans.isEmpty()) {
         return null
       }
 
-      val cloned = SpannableString(spanned)
+      val spanClones = IdentityHashMap<StatefulSpan, StatefulSpan>()
+      val clonedSourceText = SpannableString(sourceText)
+      for (oldSpan in sourceText.getSpans(0, sourceText.length, StatefulSpan::class.java)) {
+        val newSpan = spanClones.getOrPut(oldSpan) { oldSpan.clone() }
+        replaceSpan(clonedSourceText, oldSpan, newSpan)
+      }
+
+      val clonedLayoutText = SpannableString(layoutText)
       for (oldSpan in statefulSpans) {
-        val start = cloned.getSpanStart(oldSpan)
-        val end = cloned.getSpanEnd(oldSpan)
-        val flags = cloned.getSpanFlags(oldSpan)
-        cloned.removeSpan(oldSpan)
-        cloned.setSpan(oldSpan.clone(), start, end, flags)
+        val newSpan = spanClones.getOrPut(oldSpan) { oldSpan.clone() }
+        replaceSpan(clonedLayoutText, oldSpan, newSpan)
       }
-      return MutableSpannableLayout(layout, cloned)
+      return MutableSpannableLayout(layout, clonedLayoutText, clonedSourceText)
+    }
+
+    @OptIn(UnstableReactNativeAPI::class)
+    private fun replaceSpan(
+        text: SpannableString,
+        oldSpan: StatefulSpan,
+        newSpan: StatefulSpan,
+    ) {
+      val start = text.getSpanStart(oldSpan)
+      val end = text.getSpanEnd(oldSpan)
+      val flags = text.getSpanFlags(oldSpan)
+      text.removeSpan(oldSpan)
+      text.setSpan(newSpan, start, end, flags)
     }
   }
 

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/text/PreparedLayout.kt

@@ -8,18 +8,31 @@
 package com.facebook.react.views.text
 
 import android.text.Layout
+import android.text.Spanned
 import com.facebook.proguard.annotations.DoNotStrip
 
 /**
- * Encapsulates an {android.text.Layout} along with any additional state needed to render or measure
- * it.
+ * Encapsulates an [android.text.Layout] along with the full source text and any additional state
+ * needed to render or measure it.
  */
 @DoNotStrip
 internal class PreparedLayout(
     val layout: Layout,
+    val sourceText: Spanned,
     val maximumNumberOfLines: Int,
     val verticalOffset: Float,
     val reactTags: IntArray,
     val textBreakStrategy: Int,
     val justificationMode: Int,
-)
+) {
+  fun copyWithReactTags(reactTags: IntArray): PreparedLayout =
+      PreparedLayout(
+          layout,
+          sourceText,
+          maximumNumberOfLines,
+          verticalOffset,
+          reactTags,
+          textBreakStrategy,
+          justificationMode,
+      )
+}

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/text/PreparedLayoutTextView.kt

@@ -93,7 +93,7 @@ internal class PreparedLayoutTextView(context: Context) : ViewGroup(context), Re
   val text: CharSequence?
     // Avoid mangling the getter name, to allow black box E2E tests to read text content via
     // reflection
-    @DoNotStrip get() = preparedLayout?.layout?.text
+    @DoNotStrip get() = preparedLayout?.sourceText
 
   init {
     // ViewGroup by default says only its children will draw
@@ -136,7 +136,7 @@ internal class PreparedLayoutTextView(context: Context) : ViewGroup(context), Re
             selectionColor ?: DefaultStyleValuesUtil.getDefaultTextColorHighlight(context)
       }
 
-      val spanned = text as? Spanned
+      val spanned = layout.text as? Spanned
       val canvasEffectSpans =
           spanned?.getSpans(0, spanned.length, CanvasEffectSpan::class.java) ?: emptyArray()
 
@@ -242,7 +242,7 @@ internal class PreparedLayoutTextView(context: Context) : ViewGroup(context), Re
     val action = event.actionMasked
     if (action == MotionEvent.ACTION_CANCEL) {
       // Forward ACTION_CANCEL to all TouchableSpans so they can reset pressed/animation state
-      val spanned = text as? Spanned
+      val spanned = preparedLayout?.layout?.text as? Spanned
       for (span in spanned?.getSpans(0, spanned.length, TouchableSpan::class.java).orEmpty()) {
         span.onTouchEvent(action, 0f, 0f)
       }
@@ -314,7 +314,7 @@ internal class PreparedLayoutTextView(context: Context) : ViewGroup(context), Re
       return null
     }
 
-    val spanned = text as? Spanned ?: return null
+    val spanned = preparedLayout?.layout?.text as? Spanned ?: return null
 
     val spans = spanned.getSpans(offset, offset, clazz)
     if (spans.isEmpty()) {
@@ -490,13 +490,15 @@ internal class PreparedLayoutTextView(context: Context) : ViewGroup(context), Re
     }
 
     /**
-     * If the layout contains [StatefulSpan]s, returns a new [PreparedLayout] whose spannable has
-     * independent clones of those spans. Otherwise returns the receiver unchanged.
+     * If the layout contains [StatefulSpan]s, returns a new [PreparedLayout] whose display and
+     * source text share independent clones of those spans. Otherwise returns the receiver
+     * unchanged.
      */
     private fun PreparedLayout.maybeProxyStatefulSpans(): PreparedLayout {
-      val proxyLayout = MutableSpannableLayout.createIfNeeded(layout) ?: return this
+      val proxyLayout = MutableSpannableLayout.createIfNeeded(layout, sourceText) ?: return this
       return PreparedLayout(
           proxyLayout,
+          proxyLayout.sourceText,
           maximumNumberOfLines,
           verticalOffset,
           reactTags,

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/text/ReactTextViewManager.kt

@@ -17,6 +17,7 @@ import android.text.Spanned
 import android.text.TextUtils
 import android.text.util.Linkify
 import android.view.Gravity
+import androidx.annotation.VisibleForTesting
 import com.facebook.common.logging.FLog
 import com.facebook.react.R
 import com.facebook.react.bridge.Dynamic
@@ -207,13 +208,13 @@ public constructor(
   /**
    * Constructs a [ReactTextUpdate] from a [PreparedLayout] received via [ReferenceStateWrapper].
    */
-  private fun getReactTextUpdateFromPreparedLayout(
+  @VisibleForTesting
+  internal fun getReactTextUpdateFromPreparedLayout(
       view: ReactTextView,
       preparedLayout: PreparedLayout,
   ): ReactTextUpdate {
     val layout = preparedLayout.layout
-    val text = layout.text
-    val spanned = if (text is Spannable) text else SpannableString(text)
+    val spanned = SpannableString(preparedLayout.sourceText)
     view.setSpanned(spanned)
     view.setPreparedLayout(preparedLayout)
 

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/text/TextLayoutManager.kt

@@ -1328,6 +1328,7 @@ internal object TextLayoutManager {
 
     return PreparedLayout(
         result.layout,
+        text,
         maximumNumberOfLines,
         verticalOffset,
         reactTags,