feat: refactor snapshots when going back on Fabric (software-mansion#2134)

PR adding snapshots when going back on Fabric on Android and changing the behavior a bit on iOS.

Author: WoLewicki

RNScreens.podspec

@@ -4,7 +4,7 @@ package = JSON.parse(File.read(File.join(__dir__, "package.json")))
 
 new_arch_enabled = ENV['RCT_NEW_ARCH_ENABLED'] == '1'
 platform = new_arch_enabled ? "11.0" : "9.0"
-source_files = new_arch_enabled ? 'ios/**/*.{h,m,mm,cpp}' : ["ios/**/*.{h,m,mm}", "cpp/**/*.{cpp,h}"]
+source_files = new_arch_enabled ? 'ios/**/*.{h,m,mm,cpp}' : ["ios/**/*.{h,m,mm}", "cpp/RNScreensTurboModule.cpp", "cpp/RNScreensTurboModule.h"]
 
 Pod::Spec.new do |s|
   s.name         = "RNScreens"

android/CMakeLists.txt

@@ -2,11 +2,22 @@ cmake_minimum_required(VERSION 3.9.0)
 
 project(rnscreens)
 
+if(${RNS_NEW_ARCH_ENABLED})
 add_library(rnscreens
     SHARED
     ../cpp/RNScreensTurboModule.cpp
+    ../cpp/RNSScreenRemovalListener.cpp
     ./src/main/cpp/jni-adapter.cpp
+    ./src/main/cpp/NativeProxy.cpp
+    ./src/main/cpp/OnLoad.cpp
 )
+else()
+add_library(rnscreens
+    SHARED
+    ../cpp/RNScreensTurboModule.cpp
+    ./src/main/cpp/jni-adapter.cpp
+)
+endif()
 
 include_directories(
     ../cpp
@@ -19,9 +30,42 @@ set_target_properties(rnscreens PROPERTIES
     POSITION_INDEPENDENT_CODE ON
 )
 
+target_compile_definitions(
+    rnscreens
+    PRIVATE
+    -DFOLLY_NO_CONFIG=1
+)
+
 find_package(ReactAndroid REQUIRED CONFIG)
 
-target_link_libraries(rnscreens
-    ReactAndroid::jsi
-    android
-)
+if(${RNS_NEW_ARCH_ENABLED})
+    find_package(fbjni REQUIRED CONFIG)
+    
+    target_link_libraries(
+        rnscreens
+            ReactAndroid::jsi
+            ReactAndroid::react_nativemodule_core
+            ReactAndroid::react_utils
+            ReactAndroid::reactnativejni
+            ReactAndroid::fabricjni
+            ReactAndroid::react_debug
+            ReactAndroid::react_render_core
+            ReactAndroid::runtimeexecutor
+            ReactAndroid::react_render_graphics
+            ReactAndroid::rrc_view
+            ReactAndroid::yoga
+            ReactAndroid::rrc_text
+            ReactAndroid::glog
+            ReactAndroid::react_render_componentregistry
+            ReactAndroid::react_render_consistency
+            ReactAndroid::react_performance_timeline
+            ReactAndroid::react_render_observers_events
+            fbjni::fbjni
+            android
+    )
+else()
+    target_link_libraries(rnscreens
+        ReactAndroid::jsi
+        android
+    )
+endif()

android/build.gradle

@@ -81,6 +81,7 @@ def reactProperties = new Properties()
 file("$reactNativeRootDir/ReactAndroid/gradle.properties").withInputStream { reactProperties.load(it) }
 def REACT_NATIVE_VERSION = reactProperties.getProperty("VERSION_NAME")
 def REACT_NATIVE_MINOR_VERSION = REACT_NATIVE_VERSION.startsWith("0.0.0-") ? 1000 : REACT_NATIVE_VERSION.split("\\.")[1].toInteger()
+def IS_NEW_ARCHITECTURE_ENABLED = isNewArchitectureEnabled()
 
 android {
     compileSdkVersion safeExtGet('compileSdkVersion', rnsDefaultCompileSdkVersion)
@@ -106,13 +107,14 @@ android {
         targetSdkVersion safeExtGet('targetSdkVersion', rnsDefaultTargetSdkVersion)
         versionCode 1
         versionName "1.0"
-        buildConfigField "boolean", "IS_NEW_ARCHITECTURE_ENABLED", isNewArchitectureEnabled().toString()
+        buildConfigField "boolean", "IS_NEW_ARCHITECTURE_ENABLED", IS_NEW_ARCHITECTURE_ENABLED.toString()
         ndk {
             abiFilters (*reactNativeArchitectures())
         }
         externalNativeBuild {
             cmake {
-                arguments "-DANDROID_STL=c++_shared"
+                arguments "-DANDROID_STL=c++_shared",
+                        "-DRNS_NEW_ARCH_ENABLED=${IS_NEW_ARCHITECTURE_ENABLED}"
             }
         }
     }
@@ -143,13 +145,15 @@ android {
                 "META-INF/**",
                 "**/libjsi.so",
                 "**/libc++_shared.so",
-                "**/libreact_render*.so"
+                "**/libreact_render*.so",
+                "**/libreactnativejni.so",
+                "**/libreact_performance_timeline.so"
         ]
     }
     sourceSets.main {
         ext.androidResDir = "src/main/res"
         java {
-            if (isNewArchitectureEnabled()) {
+            if (IS_NEW_ARCHITECTURE_ENABLED) {
                 srcDirs += [
                     "src/fabric/java",
                 ]

android/src/fabric/java/com/swmansion/rnscreens/NativeProxy.kt

@@ -0,0 +1,53 @@
+package com.swmansion.rnscreens
+
+import android.util.Log
+import com.facebook.jni.HybridData
+import com.facebook.proguard.annotations.DoNotStrip
+import com.facebook.react.fabric.FabricUIManager
+import java.lang.ref.WeakReference
+import java.util.concurrent.ConcurrentHashMap
+
+class NativeProxy {
+    @DoNotStrip
+    @Suppress("unused")
+    private val mHybridData: HybridData
+
+    init {
+        mHybridData = initHybrid()
+    }
+
+    private external fun initHybrid(): HybridData
+
+    external fun nativeAddMutationsListener(fabricUIManager: FabricUIManager)
+
+    companion object {
+        // we use ConcurrentHashMap here since it will be read on the JS thread,
+        // and written to on the UI thread.
+        private val viewsMap = ConcurrentHashMap<Int, WeakReference<Screen>>()
+
+        fun addScreenToMap(
+            tag: Int,
+            view: Screen,
+        ) {
+            viewsMap[tag] = WeakReference(view)
+        }
+
+        fun removeScreenFromMap(tag: Int) {
+            viewsMap.remove(tag)
+        }
+
+        fun clearMapOnInvalidate() {
+            viewsMap.clear()
+        }
+    }
+
+    @DoNotStrip
+    public fun notifyScreenRemoved(screenTag: Int) {
+        val screen = viewsMap[screenTag]?.get()
+        if (screen is Screen) {
+            screen.startRemovalTransition()
+        } else {
+            Log.w("[RNScreens]", "Did not find view with tag $screenTag.")
+        }
+    }
+}

android/src/main/cpp/NativeProxy.cpp

@@ -0,0 +1,51 @@
+#include <fbjni/fbjni.h>
+#include <react/fabric/Binding.h>
+#include <react/renderer/scheduler/Scheduler.h>
+
+#include <string>
+
+#include "NativeProxy.h"
+
+using namespace facebook;
+using namespace react;
+
+namespace rnscreens {
+
+NativeProxy::NativeProxy(jni::alias_ref<NativeProxy::javaobject> jThis)
+    : javaPart_(jni::make_global(jThis)) {}
+
+NativeProxy::~NativeProxy() {}
+
+void NativeProxy::registerNatives() {
+  registerHybrid(
+      {makeNativeMethod("initHybrid", NativeProxy::initHybrid),
+       makeNativeMethod(
+           "nativeAddMutationsListener",
+           NativeProxy::nativeAddMutationsListener)});
+}
+
+void NativeProxy::nativeAddMutationsListener(
+    jni::alias_ref<facebook::react::JFabricUIManager::javaobject>
+        fabricUIManager) {
+  auto uiManager =
+      fabricUIManager->getBinding()->getScheduler()->getUIManager();
+  screenRemovalListener_ =
+      std::make_shared<RNSScreenRemovalListener>([this](int tag) {
+        static const auto method =
+            javaPart_->getClass()->getMethod<void(jint)>("notifyScreenRemoved");
+        method(javaPart_, tag);
+      });
+
+  uiManager->getShadowTreeRegistry().enumerate(
+      [this](const facebook::react::ShadowTree &shadowTree, bool &stop) {
+        shadowTree.getMountingCoordinator()->setMountingOverrideDelegate(
+            screenRemovalListener_);
+      });
+}
+
+jni::local_ref<NativeProxy::jhybriddata> NativeProxy::initHybrid(
+    jni::alias_ref<jhybridobject> jThis) {
+  return makeCxxInstance(jThis);
+}
+
+} // namespace rnscreens

android/src/main/cpp/NativeProxy.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include <fbjni/fbjni.h>
+#include <react/fabric/JFabricUIManager.h>
+#include "RNSScreenRemovalListener.h"
+
+#include <string>
+
+namespace rnscreens {
+using namespace facebook;
+using namespace facebook::jni;
+
+class NativeProxy : public jni::HybridClass<NativeProxy> {
+ public:
+  std::shared_ptr<RNSScreenRemovalListener> screenRemovalListener_;
+  static auto constexpr kJavaDescriptor =
+      "Lcom/swmansion/rnscreens/NativeProxy;";
+  static jni::local_ref<jhybriddata> initHybrid(
+      jni::alias_ref<jhybridobject> jThis);
+  static void registerNatives();
+
+  ~NativeProxy();
+
+ private:
+  friend HybridBase;
+  jni::global_ref<NativeProxy::javaobject> javaPart_;
+
+  explicit NativeProxy(jni::alias_ref<NativeProxy::javaobject> jThis);
+
+  void nativeAddMutationsListener(
+      jni::alias_ref<facebook::react::JFabricUIManager::javaobject>
+          fabricUIManager);
+};
+
+} // namespace rnscreens

android/src/main/cpp/OnLoad.cpp

@@ -0,0 +1,8 @@
+#include <fbjni/fbjni.h>
+
+#include "NativeProxy.h"
+
+JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *) {
+  return facebook::jni::initialize(
+      vm, [] { rnscreens::NativeProxy::registerNatives(); });
+}

android/src/main/java/com/swmansion/rnscreens/Screen.kt

@@ -6,6 +6,7 @@ import android.graphics.Paint
 import android.os.Parcelable
 import android.util.SparseArray
 import android.util.TypedValue
+import android.view.View
 import android.view.ViewGroup
 import android.view.WindowManager
 import android.webkit.WebView
@@ -37,6 +38,7 @@ class Screen(
     var screenOrientation: Int? = null
         private set
     var isStatusBarAnimated: Boolean? = null
+    var isBeingRemoved = false
 
     init {
         // we set layout params as WindowManager.LayoutParams to workaround the issue with TextInputs
@@ -280,6 +282,40 @@ class Screen(
 
     var nativeBackButtonDismissalEnabled: Boolean = true
 
+    fun startRemovalTransition() {
+        if (!isBeingRemoved) {
+            isBeingRemoved = true
+            startTransitionRecursive(this)
+        }
+    }
+
+    private fun startTransitionRecursive(parent: ViewGroup?) {
+        parent?.let {
+            for (i in 0 until it.childCount) {
+                val child = it.getChildAt(i)
+                if (child.javaClass.simpleName.equals("CircleImageView")) {
+                    // SwipeRefreshLayout class which has CircleImageView as a child,
+                    // does not handle `startViewTransition` properly.
+                    // It has a custom `getChildDrawingOrder` method which returns
+                    // wrong index if we called `startViewTransition` on the views on new arch.
+                    // We add a simple View to bump the number of children to make it work.
+                    // TODO: find a better way to handle this scenario
+                    it.addView(View(context), i)
+                } else {
+                    child?.let { view -> it.startViewTransition(view) }
+                }
+                if (child is ScreenStackHeaderConfig) {
+                    // we want to start transition on children of the toolbar too,
+                    // which is not a child of ScreenStackHeaderConfig
+                    startTransitionRecursive(child.toolbar)
+                }
+                if (child is ViewGroup) {
+                    startTransitionRecursive(child)
+                }
+            }
+        }
+    }
+
     private fun calculateHeaderHeight(): Pair<Double, Double> {
         val actionBarTv = TypedValue()
         val resolvedActionBarSize =

android/src/main/java/com/swmansion/rnscreens/ScreenStackHeaderConfig.kt

@@ -304,7 +304,7 @@ class ScreenStackHeaderConfig(
     }
 
     private fun maybeUpdate() {
-        if (parent != null && !isDestroyed) {
+        if (parent != null && !isDestroyed && screen?.isBeingRemoved == false) {
             onUpdate()
         }
     }

android/src/main/java/com/swmansion/rnscreens/ScreenStackHeaderConfigViewManager.kt

@@ -203,7 +203,7 @@ class ScreenStackHeaderConfigViewManager :
 
     // TODO: Find better way to handle platform specific props
     private fun logNotAvailable(propName: String) {
-        Log.w("RN SCREENS", "$propName prop is not available on Android")
+        Log.w("[RNScreens]", "$propName prop is not available on Android")
     }
 
     override fun setBackTitle(