From c98ed5ee16e1ad4e28aaffda7d8c65e4d85271ae Mon Sep 17 00:00:00 2001
From: Adlai Holler <adlai@google.com>
Date: Wed, 7 Nov 2018 15:08:01 -0800
Subject: [PATCH] Push root node ref down the tree

---
 Source/ASDisplayNode.mm                | 33 +++++++++++++++++++++-----
 Source/Private/ASDisplayNodeInternal.h |  6 +++--
 Tests/ASDisplayNodeTests.mm            | 18 ++++++++++++++
 3 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/Source/ASDisplayNode.mm b/Source/ASDisplayNode.mm
index 89237626..fe9dfb69 100644
--- a/Source/ASDisplayNode.mm
+++ b/Source/ASDisplayNode.mm
@@ -2108,16 +2108,23 @@ ASDISPLAYNODE_INLINE BOOL subtreeIsRasterized(ASDisplayNode *node) {
 // NOTE: This method must be dealloc-safe (should not retain self).
 - (ASDisplayNode *)supernode
 {
-#if CHECK_LOCKING_SAFETY
-  if (__instanceLock__.locked()) {
-    NSLog(@"WARNING: Accessing supernode while holding recursive instance lock of this node is worrisome. It's likely that you will soon try to acquire the supernode's lock, and this can easily cause deadlocks.");
-  }
-#endif
-  
   ASDN::MutexLocker l(__instanceLock__);
   return _supernode;
 }
 
+- (ASDN::UniqueLock)acquireRootLock __unused {
+  for (;; std::this_thread::yield()) {
+    if (!__instanceLock__.try_lock()) {
+      continue;
+    }
+    ASDN::UniqueLock rootLock(_rootNode->__instanceLock__, std::try_to_lock);
+    if (!rootLock.owns_lock()) {
+      continue;
+    }
+    return rootLock;
+  }
+}
+
 - (void)_setSupernode:(ASDisplayNode *)newSupernode
 {
   BOOL supernodeDidChange = NO;
@@ -2129,6 +2136,20 @@ ASDISPLAYNODE_INLINE BOOL subtreeIsRasterized(ASDisplayNode *node) {
                                   // in case supernode implementation must access one of our properties.
       _supernode = newSupernode;
       supernodeDidChange = YES;
+      
+      unowned ASDisplayNode *newRoot = newSupernode ? newSupernode->_rootNode : self;
+      if (newRoot != self) {
+        _rootNode = newRoot;
+      } else {
+        _rootNode = nil;
+      }
+      // Push new root node ref down while holding self.
+      // In the future, tree modifications could require the root node to be locked and thus
+      // we could guarantee a stable view of a tree on-demand.
+      ASDisplayNodePerformBlockOnEverySubnode(self, NO, ^(ASDisplayNode *node) {
+        ASDN::MutexLocker l(node->__instanceLock__);
+        node->_rootNode = newRoot;
+      });
     }
   }
   
diff --git a/Source/Private/ASDisplayNodeInternal.h b/Source/Private/ASDisplayNodeInternal.h
index 1e907d3e..d0b1d9b3 100644
--- a/Source/Private/ASDisplayNodeInternal.h
+++ b/Source/Private/ASDisplayNodeInternal.h
@@ -78,7 +78,9 @@ AS_EXTERN NSString * const ASRenderingEngineDidDisplayNodesScheduledBeforeTimest
 
 @interface ASDisplayNode () <_ASTransitionContextCompletionDelegate>
 {
-@package
+@public
+  // Note: These ivars are declared public only for testing. Don't access them from outside please!
+  
   ASDN::RecursiveMutex __instanceLock__;
 
   _ASPendingState *_pendingViewState;
@@ -126,7 +128,7 @@ AS_EXTERN NSString * const ASRenderingEngineDidDisplayNodesScheduledBeforeTimest
     unsigned isDeallocating:1;
   } _flags;
   
-@protected
+  unowned ASDisplayNode *_rootNode;
   ASDisplayNode * __weak _supernode;
   NSMutableArray<ASDisplayNode *> *_subnodes;
 
diff --git a/Tests/ASDisplayNodeTests.mm b/Tests/ASDisplayNodeTests.mm
index f1618e49..7d8e5db7 100644
--- a/Tests/ASDisplayNodeTests.mm
+++ b/Tests/ASDisplayNodeTests.mm
@@ -2698,4 +2698,22 @@ static bool stringContainsPointer(NSString *description, id p) {
   }
 }
 
+- (void)testRootNodePropagation
+{
+  // [A [B [C, D]]
+  ASDisplayNode *a = [ASDisplayNode new];
+  XCTAssertNil(a->_rootNode);
+  ASDisplayNode *b = [ASDisplayNode new];
+  [a addSubnode:b];
+  XCTAssertEqual(b->_rootNode, a);
+  ASDisplayNode *c = [ASDisplayNode new];
+  [b addSubnode:c];
+  ASDisplayNode *d = [ASDisplayNode new];
+  [b addSubnode:d];
+  XCTAssertEqual(d->_rootNode, a);
+  [b removeFromSupernode];
+  XCTAssertNil(b->_rootNode);
+  XCTAssertEqual(c->_rootNode, b);
+}
+
 @end