refactor(LayoutAnimations): make shared container ownership explicit

packages/react-native-reanimated/CHANGELOG.md

@@ -38,6 +38,7 @@
 - Fix animated styles on sticky headers throwing an immutable-object mutation error in development. ([#10389](https://github.com/software-mansion/react-native-reanimated/pull/10389) by [@ngocdevv](https://github.com/ngocdevv))
 - Fix shared element transitions never running on Android for npm installs by publishing `react-native.config.js`, which registers the Shared Transition Boundary component descriptor for autolinking. ([#10375](https://github.com/software-mansion/react-native-reanimated/pull/10375) by [@dennytosp](https://github.com/dennytosp))
 - Keep exiting views at their position in the host tree, so they no longer draw above later siblings. ([#10392](https://github.com/software-mansion/react-native-reanimated/pull/10392) by [@pawicao](https://github.com/pawicao))
+- Fix Shared Transitions cleanup, restoration, and positioning in native-stack screens and modals. ([#10370](https://github.com/software-mansion/react-native-reanimated/pull/10370) by [@bartlomiejbloniarz](https://github.com/bartlomiejbloniarz))
 - Fix Layout Animations replacement and completion ordering by centralizing native lifecycle operations. ([#10369](https://github.com/software-mansion/react-native-reanimated/pull/10369) by [@bartlomiejbloniarz](https://github.com/bartlomiejbloniarz))
 - Fix Layout Animations state leaking between Fabric surfaces by using one proxy per surface. ([#10368](https://github.com/software-mansion/react-native-reanimated/pull/10368) by [@bartlomiejbloniarz](https://github.com/bartlomiejbloniarz))
 

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsManager.cpp

@@ -27,7 +27,7 @@ void LayoutAnimationsManager::configureAnimationBatch(const std::vector<LayoutAn
     if (type == LayoutAnimationType::SHARED_ELEMENT_TRANSITION) {
       if (config == nullptr) {
         // TODO (future): if the view was transitioned (e.g. so we are on the second screen)
-        // and we remove the config, we should also bring back the view (probably using TransactionMeta::tagsToRestore)
+        // and we remove the config, we should also bring back the view (probably using TransactionMeta::nodesToRestore)
         sharedTransitions_.erase(tag);
         sharedTransitionManager_->tagToName_.erase(tag);
       } else {

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsManager.h

@@ -27,9 +27,6 @@ struct Transition {
   std::optional<Transform> transform[2];
 };
 
-using TransitionMap = std::unordered_map<SharedTag, Transition>;
-using Transitions = std::vector<std::pair<SharedTag, Transition>>;
-
 struct SharedTransitionManager {
   std::mutex mutex_;
   std::unordered_map<Tag, std::string> tagToName_;

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsProxy_Experimental.cpp

@@ -107,9 +107,9 @@ std::optional<MountingTransaction> LayoutAnimationsProxy_Experimental::pullTrans
       // To keep things simple, we put mutations related to source views before all muatations
       // and mutations to hide target views after all mutations.
       std::vector<ShadowViewMutation> mergedMutations;
-      hideTransitioningViews(BEFORE, transaction.transitions, mergedMutations, propsParserContext);
+      hideTransitioningViews(BEFORE, mergedMutations, transaction, propsParserContext);
       mergedMutations.insert(mergedMutations.end(), filteredMutations.begin(), filteredMutations.end());
-      hideTransitioningViews(AFTER, transaction.transitions, mergedMutations, propsParserContext);
+      hideTransitioningViews(AFTER, mergedMutations, transaction, propsParserContext);
       std::swap(filteredMutations, mergedMutations);
     }
 
@@ -716,13 +716,8 @@ void LayoutAnimationsProxy_Experimental::maybeCancelAnimation(const int tag) con
 
 void LayoutAnimationsProxy_Experimental::surfaceDidUnmount() {
   LayoutAnimationsProxyCommon::surfaceDidUnmount();
-  {
-    auto lock = std::unique_lock<std::recursive_mutex>(mutex);
-    auto sharedTransitionLock = std::unique_lock<std::mutex>(sharedTransitionManager_->mutex_);
-    for (const auto &[_, containerTag] : containerTags_) {
-      sharedTransitionManager_->tagToName_.erase(containerTag);
-    }
-  }
+  auto lock = std::unique_lock<std::recursive_mutex>(mutex);
+  sharedContainers_.clear();
 }
 
 // When entering animations start, we temporarily set opacity to 0
@@ -745,15 +740,17 @@ void LayoutAnimationsProxy_Experimental::cleanupAnimations(
     TransactionMeta &transaction,
     const PropsParserContext &propsParserContext) const {
   ReanimatedSystraceSection s("cleanupAnimations");
-  for (const auto &[tag, completedAnimation] : completedAnimations_) {
+  for (const auto &[tag, _] : completedAnimations_) {
     if (hasPendingLayoutAnimation(tag)) {
       continue;
     }
-    const auto restoreIt = restoreMap_.find(tag);
-    if (restoreIt == restoreMap_.end()) {
+    const auto containerIt = sharedContainers_.find(tag);
+    if (containerIt == sharedContainers_.end()) {
       continue;
     }
-    transaction.tagsToRestore.push_back(restoreIt->second[AFTER]);
+    if (containerIt->second.restoreAfterNode) {
+      transaction.nodesToRestore.push_back(containerIt->second.restoreAfterNode);
+    }
     removeSharedContainer(tag, transaction);
   }
 

packages/react-native-reanimated/Common/cpp/reanimated/LayoutAnimations/LayoutAnimationsProxy_Experimental.h

@@ -14,6 +14,7 @@
 #include <react/renderer/scheduler/Scheduler.h>
 #include <react/renderer/uimanager/UIManagerBinding.h>
 
+#include <array>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -39,28 +40,41 @@ struct PendingNodeAnimation {
   std::shared_ptr<Serializable> config;
 };
 
+struct SharedContainer {
+  SharedTag sharedTag;
+  std::shared_ptr<LightNode> node;
+  std::shared_ptr<LightNode> restoreBeforeNode;
+  std::shared_ptr<LightNode> restoreAfterNode;
+};
+
+struct CollectedTransition {
+  Transition transition;
+  std::array<std::shared_ptr<LightNode>, 2> nodes;
+};
+
+using CollectedTransitionMap = std::unordered_map<SharedTag, CollectedTransition>;
+using CollectedTransitions = std::vector<std::pair<SharedTag, CollectedTransition>>;
+
 struct TransactionMeta {
   ShadowViewMutationList filteredMutations;
   ShadowViewMutationList teardownMutations;
-  TransitionMap transitionMap;
-  Transitions transitions;
+  CollectedTransitionMap transitionMap;
+  CollectedTransitions transitions;
   std::vector<PendingNodeAnimation> layout;
   std::vector<PendingNodeAnimation> entering;
   std::vector<PendingNodeAnimation> exiting;
   std::vector<std::shared_ptr<LightNode>> containersToInsert;
-  std::vector<Tag> tagsToRestore;
-  std::vector<Tag> sharedContainersToRemove;
+  std::vector<std::shared_ptr<LightNode>> nodesToRestore;
+  std::vector<std::shared_ptr<LightNode>> containersToRemove;
 };
 
 struct LayoutAnimationsProxy_Experimental : public LayoutAnimationsProxyCommon {
-  mutable std::unordered_set<Tag> activeTransitions_;
   mutable Tag transitionTag_;
   mutable double transitionProgress_;
   mutable bool transitionUpdated_;
   mutable TransitionState transitionState_ = TransitionState::NONE;
   mutable std::shared_ptr<LightNode> topScreen_;
-  mutable std::unordered_map<Tag, Tag[2]> restoreMap_;
-  mutable std::unordered_map<std::string, Tag> containerTags_;
+  mutable std::unordered_map<Tag, SharedContainer> sharedContainers_;
   mutable bool synchronized_ = true;
   mutable Tag closingScreenTag_ = -1;
   std::shared_ptr<SharedTransitionManager> sharedTransitionManager_;
@@ -112,8 +126,8 @@ struct LayoutAnimationsProxy_Experimental : public LayoutAnimationsProxyCommon {
 
   void hideTransitioningViews(
       BeforeOrAfter index,
-      const Transitions &transitions,
       ShadowViewMutationList &mutations,
+      TransactionMeta &transaction,
       const PropsParserContext &propsParserContext) const;
 
   std::optional<SurfaceId> endLayoutAnimation(int tag, bool shouldRemove) override;
@@ -141,7 +155,11 @@ struct LayoutAnimationsProxy_Experimental : public LayoutAnimationsProxyCommon {
 
   std::vector<react::Point> getAbsolutePositionsForRootPathView(const std::shared_ptr<LightNode> &node) const;
 
-  Tag getOrCreateContainer(const ShadowView &before, const SharedTag &sharedTag, TransactionMeta &transaction) const;
+  Tag getOrCreateContainer(
+      const ShadowView &before,
+      const SharedTag &sharedTag,
+      const std::array<std::shared_ptr<LightNode>, 2> &nodes,
+      TransactionMeta &transaction) const;
 
   void overrideTransform(
       ShadowView &shadowView,