fix: keep the parser registration alive across effect remounts

android/src/main/cpp/MarkdownParser.cpp

@@ -7,14 +7,46 @@ using namespace facebook;
 
 namespace expensify {
 namespace livemarkdown {
+  jni::local_ref<MarkdownParser::jhybriddata> MarkdownParser::initHybrid(jni::alias_ref<jclass>) {
+    return makeCxxInstance();
+  }
+
+  void MarkdownParser::nativeSetParserId(const int parserId) {
+    std::unique_lock<std::mutex> lock(mutex_);
+    if (parserId_ == parserId) {
+      return;
+    }
+    const auto markdownWorklet = findMarkdownWorklet(parserId);
+    if (markdownWorklet == nullptr) {
+      return;
+    }
+    parserId_ = parserId;
+    markdownWorklet_ = markdownWorklet;
+  }
+
+  // A parse for the current id uses the worklet kept alive by `nativeSetParserId`.
+  // Any other id is looked up in the registry the way it always was.
+  std::shared_ptr<SerializableWorklet> MarkdownParser::workletForParserId(const int parserId) {
+    {
+      std::unique_lock<std::mutex> lock(mutex_);
+      if (parserId_ == parserId) {
+        return markdownWorklet_;
+      }
+    }
+
+    return findMarkdownWorklet(parserId);
+  }
+
   jni::local_ref<jni::JString> MarkdownParser::nativeParse(
-      jni::alias_ref<jhybridobject> jThis,
       jni::alias_ref<jni::JString> text,
       const int parserId) {
-    const auto markdownRuntime = expensify::livemarkdown::getMarkdownRuntime();
-    jsi::Runtime &rt = markdownRuntime->getJSIRuntime();
+    const auto markdownWorklet = workletForParserId(parserId);
+    if (markdownWorklet == nullptr) {
+      return jni::make_jstring("[]");
+    }
 
-    const auto markdownWorklet = expensify::livemarkdown::getMarkdownWorklet(parserId);
+    const auto markdownRuntime = getMarkdownRuntime();
+    jsi::Runtime &rt = markdownRuntime->getJSIRuntime();
 
     const auto input = jsi::String::createFromUtf8(rt, text->toStdString());
     const auto output = markdownRuntime->runGuarded(markdownWorklet, input);
@@ -25,6 +57,8 @@ namespace livemarkdown {
 
   void MarkdownParser::registerNatives() {
     registerHybrid({
+        makeNativeMethod("initHybrid", MarkdownParser::initHybrid),
+        makeNativeMethod("nativeSetParserId", MarkdownParser::nativeSetParserId),
         makeNativeMethod("nativeParse", MarkdownParser::nativeParse)});
   }
 

android/src/main/cpp/MarkdownParser.h

@@ -10,26 +10,48 @@
 #include <fbjni/fbjni.h>
 #include <jsi/jsi.h>
 
+#include <worklets/WorkletRuntime/WorkletRuntime.h>
+
+#include <memory>
+#include <mutex>
+
 using namespace facebook;
+using namespace worklets;
 
 namespace expensify {
 namespace livemarkdown {
 
-  class MarkdownParser : public jni::HybridClass<MarkdownParser>,
-                        public jsi::HostObject {
+  class MarkdownParser : public jni::HybridClass<MarkdownParser> {
   public:
     static constexpr auto kJavaDescriptor =
         "Lcom/expensify/livemarkdown/MarkdownParser;";
 
-    static jni::local_ref<jni::JString> nativeParse(
-        jni::alias_ref<jhybridobject> jThis,
+    static jni::local_ref<jhybriddata> initHybrid(jni::alias_ref<jclass>);
+
+    // Looks up the worklet registered under `parserId` and keeps it alive until
+    // another registered id is set or this parser is released. JS unregisters
+    // the id when React cleans up effects, which also happens for an input that
+    // is hidden but still mounted, so the registry can't be asked again at
+    // parse time. An id the registry doesn't know leaves the previous worklet
+    // in place.
+    void nativeSetParserId(const int parserId);
+
+    jni::local_ref<jni::JString> nativeParse(
         jni::alias_ref<jni::JString> text,
         const int parserId);
 
     static void registerNatives();
 
   private:
     friend HybridBase;
+
+    MarkdownParser() = default;
+
+    std::shared_ptr<SerializableWorklet> workletForParserId(const int parserId);
+
+    std::mutex mutex_;
+    int parserId_ = 0;
+    std::shared_ptr<SerializableWorklet> markdownWorklet_;
   };
 
 } // namespace livemarkdown

android/src/main/java/com/expensify/livemarkdown/MarkdownParser.java

@@ -2,6 +2,8 @@
 
 import androidx.annotation.NonNull;
 
+import com.facebook.jni.HybridData;
+import com.facebook.jni.annotations.DoNotStrip;
 import com.facebook.react.bridge.ReactContext;
 import com.facebook.react.util.RNLog;
 import com.facebook.soloader.SoLoader;
@@ -20,17 +22,34 @@ public class MarkdownParser {
     SoLoader.loadLibrary("livemarkdown");
   }
 
+  @DoNotStrip
+  @SuppressWarnings("unused")
+  private final HybridData mHybridData;
+
   private final @NonNull ReactContext mReactContext;
   private String mPrevText;
   private int mPrevParserId;
   private List<MarkdownRange> mPrevMarkdownRanges;
 
   public MarkdownParser(@NonNull ReactContext reactContext) {
     mReactContext = reactContext;
+    mHybridData = initHybrid();
   }
 
+  private static native HybridData initHybrid();
+
+  private native void nativeSetParserId(int parserId);
+
   private native String nativeParse(@NonNull String text, int parserId);
 
+  /**
+   * Keeps the worklet registered under {@code parserId} alive in native code for as long as this parser lives, so a
+   * later parse still works after JS has unregistered the id. See {@code MarkdownParser.h} for why that happens.
+   */
+  public synchronized void setParserId(int parserId) {
+    nativeSetParserId(parserId);
+  }
+
   public synchronized List<MarkdownRange> parse(@NonNull String text, int parserId) {
     try {
       Systrace.beginSection(0, "parse");

android/src/main/java/com/expensify/livemarkdown/MarkdownTextInputDecoratorView.java

@@ -15,12 +15,18 @@ public class MarkdownTextInputDecoratorView extends ReactViewGroup {
 
   public MarkdownTextInputDecoratorView(Context context) {
     super(context);
+    mMarkdownParser = new MarkdownParser((ReactContext) context);
   }
 
   private MarkdownStyle mMarkdownStyle;
 
   private int mParserId;
 
+  // Owned by the view rather than by `mMarkdownUtils`, which is recreated every
+  // time the view is attached, so the parser worklet stays alive for as long as
+  // the view is mounted.
+  private final MarkdownParser mMarkdownParser;
+
   private MarkdownUtils mMarkdownUtils;
 
   private ReactEditText mReactEditText;
@@ -33,7 +39,7 @@ protected void onAttachedToWindow() {
 
     View child = getChildAt(0);
     if (child instanceof ReactEditText) {
-      mMarkdownUtils = new MarkdownUtils((ReactContext) getContext());
+      mMarkdownUtils = new MarkdownUtils((ReactContext) getContext(), mMarkdownParser);
       mMarkdownUtils.setMarkdownStyle(mMarkdownStyle);
       mMarkdownUtils.setParserId(mParserId);
       mReactEditText = (ReactEditText) child;
@@ -64,6 +70,7 @@ protected void setMarkdownStyle(MarkdownStyle markdownStyle) {
 
   protected void setParserId(int parserId) {
     mParserId = parserId;
+    mMarkdownParser.setParserId(parserId);
     if (mMarkdownUtils != null) {
       mMarkdownUtils.setParserId(mParserId);
     }

android/src/main/java/com/expensify/livemarkdown/MarkdownUtils.java

@@ -11,7 +11,11 @@
 
 public class MarkdownUtils {
   public MarkdownUtils(@NonNull ReactContext reactContext) {
-    mMarkdownParser = new MarkdownParser(reactContext);
+    this(reactContext, new MarkdownParser(reactContext));
+  }
+
+  public MarkdownUtils(@NonNull ReactContext reactContext, @NonNull MarkdownParser markdownParser) {
+    mMarkdownParser = markdownParser;
     mMarkdownFormatter = new MarkdownFormatter(reactContext.getAssets());
   }
 
@@ -27,6 +31,7 @@ public void setMarkdownStyle(@NonNull MarkdownStyle markdownStyle) {
 
   public void setParserId(int parserId) {
     mParserId = parserId;
+    mMarkdownParser.setParserId(parserId);
   }
 
   public void applyMarkdownFormatting(SpannableStringBuilder ssb) {

apple/MarkdownParser.h

@@ -5,6 +5,12 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface MarkdownParser : NSObject
 
+// Looks up the worklet registered under `parserId` and keeps it alive until
+// another registered id is set or this parser is released. JS unregisters the
+// id when React cleans up effects, which also happens for an input that is
+// hidden but still mounted, so the registry can't be asked again at parse time.
+- (void)setParserId:(nonnull NSNumber *)parserId;
+
 - (NSArray<MarkdownRange *> *)parse:(nonnull NSString *)text
                        withParserId:(nonnull NSNumber *)parserId;
 

apple/MarkdownParser.mm

@@ -56,6 +56,10 @@ @implementation MarkdownParser {
   NSNumber *_pendingParserId;
   void (^_pendingCompletion)(void);
   BOOL _warmupScheduled;
+
+  // The worklet registered under `_parserId`, kept alive here (see the header).
+  NSNumber *_parserId;
+  std::shared_ptr<SerializableWorklet> _markdownWorklet;
 }
 
 - (instancetype)init
@@ -82,6 +86,39 @@ + (dispatch_queue_t)cacheWarmupQueue
   return queue;
 }
 
+// An id the registry doesn't know leaves the previous worklet in place. The
+// measure path shares one parser between shadow node clones, so a clone that
+// still carries an older, already unregistered id must not drop the worklet
+// the current id resolved to.
+- (void)setParserId:(nonnull NSNumber *)parserId
+{
+  @synchronized (self) {
+    if ([_parserId isEqualToNumber:parserId]) {
+      return;
+    }
+    const auto markdownWorklet = expensify::livemarkdown::findMarkdownWorklet([parserId intValue]);
+    if (markdownWorklet == nullptr) {
+      return;
+    }
+    _parserId = parserId;
+    _markdownWorklet = markdownWorklet;
+  }
+}
+
+// A parse for the current id uses the worklet kept alive by `setParserId:`.
+// Any other id comes from a shadow node clone that still carries an older id,
+// so it is looked up in the registry the way it always was.
+- (std::shared_ptr<SerializableWorklet>)workletForParserId:(nonnull NSNumber *)parserId
+{
+  @synchronized (self) {
+    if ([_parserId isEqualToNumber:parserId]) {
+      return _markdownWorklet;
+    }
+  }
+
+  return expensify::livemarkdown::findMarkdownWorklet([parserId intValue]);
+}
+
 - (nullable NSArray<MarkdownRange *> *)cachedRangesForText:(nonnull NSString *)text
                                               withParserId:(nonnull NSNumber *)parserId
 {
@@ -209,16 +246,16 @@ - (void)drainPendingWarmups
 - (NSArray<MarkdownRange *> *)parseUncached:(nonnull NSString *)text
                                withParserId:(nonnull NSNumber *)parserId
 {
-  const auto &markdownRuntime = expensify::livemarkdown::getMarkdownRuntime();
-  jsi::Runtime &rt = markdownRuntime->getJSIRuntime();
-
-  std::shared_ptr<SerializableWorklet> markdownWorklet;
-  try {
-    markdownWorklet = expensify::livemarkdown::getMarkdownWorklet([parserId intValue]);
-  } catch (const std::out_of_range &error) {
+  // The first commit carries parserId 0, before JS has created the worklet runtime.
+  // Resolve the worklet before accessing that runtime.
+  const auto markdownWorklet = [self workletForParserId:parserId];
+  if (markdownWorklet == nullptr) {
     return @[];
   }
 
+  const auto &markdownRuntime = expensify::livemarkdown::getMarkdownRuntime();
+  jsi::Runtime &rt = markdownRuntime->getJSIRuntime();
+
   const auto &input = jsi::String::createFromUtf8(rt, [text UTF8String]);
 
   jsi::Value output;

apple/RCTMarkdownUtils.mm

@@ -17,6 +17,12 @@ - (instancetype)init
   return self;
 }
 
+- (void)setParserId:(NSNumber *)parserId
+{
+  _parserId = parserId;
+  [_markdownParser setParserId:parserId];
+}
+
 - (void)applyMarkdownFormatting:(nonnull NSMutableAttributedString *)attributedString
       withDefaultTextAttributes:(nonnull NSDictionary<NSAttributedStringKey, id> *)defaultTextAttributes
 {
@@ -49,6 +55,7 @@ - (void)applyMarkdownFormatting:(nonnull NSMutableAttributedString *)attributedS
     _markdownStyle = markdownStyle;
     _parserId = parserId;
   }
+  [_markdownParser setParserId:parserId];
 
   NSString *text = attributedString.string;
   NSArray<MarkdownRange *> *markdownRanges = [_markdownParser cachedRangesForText:text withParserId:parserId];

cpp/MarkdownGlobal.cpp

@@ -34,9 +34,10 @@ void unregisterMarkdownWorklet(const int parserId) {
   globalMarkdownShareableWorklets.erase(parserId);
 }
 
-std::shared_ptr<SerializableWorklet> getMarkdownWorklet(const int parserId) {
+std::shared_ptr<SerializableWorklet> findMarkdownWorklet(const int parserId) {
   std::unique_lock<std::mutex> lock(globalMarkdownShareableWorkletsMutex);
-  return globalMarkdownShareableWorklets.at(parserId);
+  const auto it = globalMarkdownShareableWorklets.find(parserId);
+  return it == globalMarkdownShareableWorklets.end() ? nullptr : it->second;
 }
 
 } // namespace livemarkdown

cpp/MarkdownGlobal.h

@@ -18,7 +18,10 @@ const int registerMarkdownWorklet(const std::shared_ptr<SerializableWorklet> &ma
 
 void unregisterMarkdownWorklet(const int parserId);
 
-std::shared_ptr<SerializableWorklet> getMarkdownWorklet(const int parserId);
+// Returns nullptr when nothing is registered under `parserId`. Callers keep the
+// result: JS drops the entry when React cleans up effects, which also happens
+// for an input that is hidden but still mounted.
+std::shared_ptr<SerializableWorklet> findMarkdownWorklet(const int parserId);
 
 } // namespace livemarkdown
 } // namespace expensify

package.json

@@ -74,6 +74,7 @@
     "@release-it/conventional-changelog": "^5.0.0",
     "@types/jest": "^29.5.14",
     "@types/react": "^19.2.0",
+    "@types/react-dom": "^19.2.7",
     "@typescript-eslint/eslint-plugin": "^8.53.1",
     "@typescript-eslint/parser": "^8.53.1",
     "del-cli": "^5.0.0",