Fix bug where a parent LK node purges child LK node's subviews

LayoutKitTests/ViewRecyclerTests.swift

@@ -14,7 +14,7 @@ class ViewRecyclerTests: XCTestCase {
     func testNilIdNotRecycledAndNotRemoved() {
         let root = View()
         let zero = View()
-        zero.isLayoutKitView = false    // default
+        zero.layoutKitRootView = nil    // default
         root.addSubview(zero)
 
         let recycler = ViewRecycler(rootView: root)
@@ -25,13 +25,13 @@ class ViewRecyclerTests: XCTestCase {
         XCTAssertEqual(v, expectedView)
 
         recycler.purgeViews()
-        XCTAssertNotNil(zero.superview, "`zero` should not be removed because `isLayoutKitView` is false")
+        XCTAssertNotNil(zero.superview, "`zero` should not be removed because `layoutKitRootView` is nil")
     }
 
     func testNilIdNotRecycledAndRemoved() {
         let root = View()
         let zero = View()
-        zero.isLayoutKitView = true // requires this flag to be removed by `ViewRecycler`
+        zero.layoutKitRootView = root
         root.addSubview(zero)
 
         let recycler = ViewRecycler(rootView: root)
@@ -42,13 +42,14 @@ class ViewRecyclerTests: XCTestCase {
         XCTAssertEqual(v, expectedView)
 
         recycler.purgeViews()
-        XCTAssertNil(zero.superview, "`zero` should be removed because `isLayoutKitView` is true")
+        XCTAssertNil(zero.superview, "`zero` should be removed because `layoutKitRootView` is set")
     }
 
     func testNonNilIdRecycled() {
         let root = View()
         let one = View(viewReuseId: "1")
         root.addSubview(one)
+        one.layoutKitRootView = root
 
         let recycler = ViewRecycler(rootView: root)
         let v: View? = recycler.makeOrRecycleView(havingViewReuseId: "1", viewProvider: {
@@ -91,13 +92,64 @@ class ViewRecyclerTests: XCTestCase {
         XCTAssertNotNil(one.superview)
     }
 
+    /// Create a hierarchy (rootOne > middleView > rootTwo > someSubview) where rootOne and rootTwo
+    /// are LayoutKit root nodes. Then ensure that purging rootOne doesn't remove someSubview.
+    func testDoesntPurgeOtherNodesSubviews() {
+        let rootOne = View(viewReuseId: "1")
+        let middleView = View()
+        let rootTwo = View(viewReuseId: "2")
+        let someSubview = View()
+
+        rootOne.addSubview(middleView)
+        middleView.addSubview(rootTwo)
+        rootTwo.addSubview(someSubview)
+
+        let recycler = ViewRecycler(rootView: rootOne)
+        _ = recycler.makeOrRecycleView(havingViewReuseId: "3") {
+            return middleView
+        }
+
+        let recycler2 = ViewRecycler(rootView: rootTwo)
+        _ = recycler2.makeOrRecycleView(havingViewReuseId: "4") {
+            return someSubview
+        }
+
+        recycler.purgeViews()
+        XCTAssertEqual(someSubview.superview, rootTwo)
+    }
+
+    /// Create a hierarchy (rootOne > middleView > rootTwo > someSubview) where rootOne and rootTwo
+    /// are LayoutKit root nodes. Then ensure that rootOne can't recycle someSubview.
+    func testDoesntRecycleOtherNodesViews() {
+        let rootOne = View(viewReuseId: "1")
+        let middleView = View(viewReuseId: "2")
+        let rootTwo = View(viewReuseId: "3")
+        let someSubview = View(viewReuseId: "4")
+
+        rootOne.addSubview(middleView)
+        middleView.layoutKitRootView = rootOne
+
+        middleView.addSubview(rootTwo)
+
+        rootTwo.addSubview(someSubview)
+        someSubview.layoutKitRootView = rootTwo
+
+        let recycler = ViewRecycler(rootView: rootOne)
+        let aDifferentView = View()
+        let v: View? = recycler.makeOrRecycleView(havingViewReuseId: "4") {
+            return aDifferentView
+        }
+        XCTAssertEqual(v, aDifferentView)
+    }
+
     /// Test for safe subview-purge in composite view e.g. UIButton.
     /// - SeeAlso: https://github.com/linkedin/LayoutKit/pull/85
     #if os(iOS) || os(tvOS)
     func testRecycledCompositeView() {
         let root = View()
         let button = UIButton(viewReuseId: "1")
         root.addSubview(button)
+        button.layoutKitRootView = root
 
         button.setTitle("dummy", for: .normal)
         button.layoutIfNeeded()

Sources/ViewRecycler.swift

@@ -20,11 +20,13 @@ class ViewRecycler {
 
     private var viewsById = [String: View]()
     private var unidentifiedViews = Set<View>()
+    private let rootView: View?
 
     /// Retains all subviews of rootView for recycling.
     init(rootView: View?) {
+        self.rootView = rootView
         rootView?.walkSubviews { (view) in
-            if let viewReuseId = view.viewReuseId {
+            if let viewReuseId = view.viewReuseId, view.layoutKitRootView == rootView {
                 self.viewsById[viewReuseId] = view
             } else {
                 self.unidentifiedViews.insert(view)
@@ -44,7 +46,7 @@ class ViewRecycler {
         }
 
         let providedView = viewProvider()
-        providedView.isLayoutKitView = true
+        providedView.layoutKitRootView = rootView
 
         // Remove the provided view from the list of cached views.
         if let viewReuseId = providedView.viewReuseId, let oldView = viewsById[viewReuseId], oldView == providedView {
@@ -63,15 +65,15 @@ class ViewRecycler {
         }
         viewsById.removeAll()
 
-        for view in unidentifiedViews where view.isLayoutKitView {
+        for view in unidentifiedViews where view.layoutKitRootView == rootView {
             view.removeFromSuperview()
         }
         unidentifiedViews.removeAll()
     }
 }
 
 private var viewReuseIdKey: UInt8 = 0
-private var isLayoutKitViewKey: UInt8 = 0
+private var layoutKitRootViewKey: UInt8 = 0
 
 extension View {
 
@@ -93,13 +95,17 @@ extension View {
         }
     }
 
-    /// Indicates the view is managed by LayoutKit that can be safely removed.
-    var isLayoutKitView: Bool {
+    /// All views managed by LayoutKit keep a weak reference to their LayoutKit root node view. This is to ensure
+    /// that when we purge views from a node, we only purge views that belong to that node - rather than purging
+    /// potentially any view in the subview hierarchy which could belong to a different LayoutKit root node.
+    var layoutKitRootView: View? {
         get {
-            return (objc_getAssociatedObject(self, &isLayoutKitViewKey) as? NSNumber)?.boolValue ?? false
+            return objc_getAssociatedObject(self, &layoutKitRootViewKey) as? View
         }
         set {
-            objc_setAssociatedObject(self, &isLayoutKitViewKey, NSNumber(value: newValue), .OBJC_ASSOCIATION_RETAIN)
+            // OBJC_ASSOCIATION_ASSIGN creates a weak reference which avoids a retain cycle since we'll
+            // have a view being referenced by another view which is somewhere in its subview hierarchy.
+            objc_setAssociatedObject(self, &layoutKitRootViewKey, newValue, .OBJC_ASSOCIATION_ASSIGN)
         }
     }
 }