Optimize receiver lookup algorithm (#246)

Author: davdroman

Sources/Introspect.swift

@@ -20,17 +20,18 @@ extension View {
         customize: @escaping (PlatformSpecificEntity) -> Void
     ) -> some View {
         if let platform = platforms.first(where: \.isCurrent) {
-            let anchorID = IntrospectionAnchorID()
+            let introspectionViewID = IntrospectionViewID()
             self.background(
                 IntrospectionAnchorView(
-                    id: anchorID
+                    id: introspectionViewID
                 )
                 .frame(width: 0, height: 0)
             )
             .overlay(
                 IntrospectionView(
-                    selector: { entity in
-                        (platform.selector ?? .default)(entity, scope ?? viewType.scope, anchorID)
+                    id: introspectionViewID,
+                    selector: { controller in
+                        (platform.selector ?? .default)(controller, scope ?? viewType.scope)
                     },
                     customize: customize
                 )
@@ -53,9 +54,6 @@ public protocol PlatformEntity: AnyObject {
 
     @_spi(Internals)
     func isDescendant(of other: Base) -> Bool
-
-    @_spi(Internals)
-    func entityWithTag(_ tag: Int) -> Base?
 }
 
 extension PlatformEntity {
@@ -65,8 +63,8 @@ extension PlatformEntity {
     }
 
     @_spi(Internals)
-    public var allDescendants: [Base] {
-        self.descendants.reduce([self~]) { $0 + $1.allDescendants~ }
+    public var allDescendants: some Sequence<Base> {
+        recursiveSequence([self~], children: { $0.descendants~ }).dropFirst()
     }
 
     func nearestCommonAncestor(with other: Base) -> Base? {
@@ -79,37 +77,26 @@ extension PlatformEntity {
         return nearestAncestor
     }
 
-    func descendantsBetween(_ bottomEntity: Base, and topEntity: Base) -> [Base] {
-        var result: [Base] = []
-        var entered = false
-
-        for descendant in self.allDescendants {
-            if descendant === bottomEntity {
-                entered = true
-            } else if descendant === topEntity {
-                break
-            } else if entered {
-                result.append(descendant)
-            }
-        }
-
-        return result
+    func allDescendants(between bottomEntity: Base, and topEntity: Base) -> some Sequence<Base> {
+        self.allDescendants
+            .lazy
+            .drop(while: { $0 !== bottomEntity })
+            .prefix(while: { $0 !== topEntity })
     }
 
     func receiver<PlatformSpecificEntity: PlatformEntity>(
-        ofType type: PlatformSpecificEntity.Type,
-        anchorID: IntrospectionAnchorID
+        ofType type: PlatformSpecificEntity.Type
     ) -> PlatformSpecificEntity? {
         let frontEntity = self
         guard
-            let backEntity = Array(frontEntity.ancestors).last?.entityWithTag(anchorID.hashValue),
+            let backEntity = frontEntity.introspectionAnchorEntity,
             let commonAncestor = backEntity.nearestCommonAncestor(with: frontEntity~)
         else {
             return nil
         }
 
         return commonAncestor
-            .descendantsBetween(backEntity~, and: frontEntity~)
+            .allDescendants(between: backEntity~, and: frontEntity~)
             .compactMap { $0 as? PlatformSpecificEntity }
             .first
     }
@@ -134,11 +121,6 @@ extension PlatformView: PlatformEntity {
     public var descendants: [PlatformView] {
         subviews
     }
-
-    @_spi(Internals)
-    public func entityWithTag(_ tag: Int) -> PlatformView? {
-        viewWithTag(tag)
-    }
 }
 
 extension PlatformViewController: PlatformEntity {
@@ -156,17 +138,4 @@ extension PlatformViewController: PlatformEntity {
     public func isDescendant(of other: PlatformViewController) -> Bool {
         self.ancestors.contains(other)
     }
-
-    @_spi(Internals)
-    public func entityWithTag(_ tag: Int) -> PlatformViewController? {
-        if self.view.tag == tag {
-            return self
-        }
-        for child in children {
-            if let childWithTag = child.entityWithTag(tag) {
-                return childWithTag
-            }
-        }
-        return nil
-    }
 }

Sources/IntrospectionSelector.swift

@@ -6,28 +6,28 @@ public struct IntrospectionSelector<Target: PlatformEntity> {
     @_spi(Internals)
     public static func from<Entry: PlatformEntity>(_ entryType: Entry.Type, selector: @escaping (Entry) -> Target?) -> Self {
         .init(
-            receiverSelector: { controller, anchorID in
-                controller.as(Entry.self)?.receiver(ofType: Entry.self, anchorID: anchorID).flatMap(selector)
+            receiverSelector: { controller in
+                controller.as(Entry.Base.self)?.receiver(ofType: Entry.self).flatMap(selector)
             },
             ancestorSelector: { controller in
-                controller.as(Entry.self)?.ancestor(ofType: Entry.self).flatMap(selector)
+                controller.as(Entry.Base.self)?.ancestor(ofType: Entry.self).flatMap(selector)
             }
         )
     }
 
-    private var receiverSelector: (IntrospectionPlatformViewController, IntrospectionAnchorID) -> Target?
+    private var receiverSelector: (IntrospectionPlatformViewController) -> Target?
     private var ancestorSelector: (IntrospectionPlatformViewController) -> Target?
 
     private init(
-        receiverSelector: @escaping (IntrospectionPlatformViewController, IntrospectionAnchorID) -> Target?,
+        receiverSelector: @escaping (IntrospectionPlatformViewController) -> Target?,
         ancestorSelector: @escaping (IntrospectionPlatformViewController) -> Target?
     ) {
         self.receiverSelector = receiverSelector
         self.ancestorSelector = ancestorSelector
     }
 
     @_spi(Internals)
-    public func withReceiverSelector(_ selector: @escaping (PlatformViewController, IntrospectionAnchorID) -> Target?) -> Self {
+    public func withReceiverSelector(_ selector: @escaping (PlatformViewController) -> Target?) -> Self {
         var copy = self
         copy.receiverSelector = selector
         return copy
@@ -40,14 +40,10 @@ public struct IntrospectionSelector<Target: PlatformEntity> {
         return copy
     }
 
-    func callAsFunction(
-        _ controller: IntrospectionPlatformViewController,
-        _ scope: IntrospectionScope,
-        _ anchorID: IntrospectionAnchorID
-    ) -> Target? {
+    func callAsFunction(_ controller: IntrospectionPlatformViewController, _ scope: IntrospectionScope) -> Target? {
         if
             scope.contains(.receiver),
-            let target = receiverSelector(controller, anchorID)
+            let target = receiverSelector(controller)
         {
             return target
         }
@@ -62,14 +58,14 @@ public struct IntrospectionSelector<Target: PlatformEntity> {
 }
 
 extension PlatformViewController {
-    func `as`<Entity: PlatformEntity>(_ entityType: Entity.Type) -> (any PlatformEntity)? {
-        if Entity.Base.self == PlatformView.self {
+    func `as`<Base: PlatformEntity>(_ baseType: Base.Type) -> (any PlatformEntity)? {
+        if Base.self == PlatformView.self {
             #if canImport(UIKit)
             return viewIfLoaded
             #elseif canImport(AppKit)
             return isViewLoaded ? view : nil
             #endif
-        } else if Entity.Base.self == PlatformViewController.self {
+        } else if Base.self == PlatformViewController.self {
             return self
         }
         return nil

Sources/IntrospectionView.swift

@@ -1,7 +1,30 @@
 import SwiftUI
 
-@_spi(Internals)
-public typealias IntrospectionAnchorID = UUID
+typealias IntrospectionViewID = UUID
+
+fileprivate enum IntrospectionStore {
+    static var shared: [IntrospectionViewID: Pair] = [:]
+
+    struct Pair {
+        weak var controller: IntrospectionPlatformViewController?
+        weak var anchor: IntrospectionAnchorPlatformViewController?
+    }
+}
+
+extension PlatformEntity {
+    var introspectionAnchorEntity: Base? {
+        if let introspectionController = self as? IntrospectionPlatformViewController {
+            return IntrospectionStore.shared[introspectionController.id]?.anchor~
+        }
+        if
+            let view = self as? PlatformView,
+            let introspectionController = view.introspectionController
+        {
+            return IntrospectionStore.shared[introspectionController.id]?.anchor?.view~
+        }
+        return nil
+    }
+}
 
 /// ⚓️
 struct IntrospectionAnchorView: PlatformViewControllerRepresentable {
@@ -14,9 +37,9 @@ struct IntrospectionAnchorView: PlatformViewControllerRepresentable {
     @Binding
     private var observed: Void // workaround for state changes not triggering view updates
 
-    let id: IntrospectionAnchorID
+    let id: IntrospectionViewID
 
-    init(id: IntrospectionAnchorID) {
+    init(id: IntrospectionViewID) {
         self._observed = .constant(())
         self.id = id
     }
@@ -31,36 +54,19 @@ struct IntrospectionAnchorView: PlatformViewControllerRepresentable {
 }
 
 final class IntrospectionAnchorPlatformViewController: PlatformViewController {
-    let id: IntrospectionAnchorID
-
-    init(id: IntrospectionAnchorID) {
-        self.id = id
+    init(id: IntrospectionViewID) {
         super.init(nibName: nil, bundle: nil)
+        IntrospectionStore.shared[id, default: .init()].anchor = self
     }
 
     @available(*, unavailable)
     required init?(coder: NSCoder) {
         fatalError("init(coder:) has not been implemented")
     }
 
-    #if canImport(UIKit)
-    override func viewDidLoad() {
-        super.viewDidLoad()
-        view.tag = id.hashValue
-    }
-    #elseif canImport(AppKit)
-    final class TaggableView: NSView {
-        private var _tag: Int?
-        override var tag: Int {
-            get { _tag ?? super.tag }
-            set { _tag = newValue }
-        }
-    }
-
+    #if canImport(AppKit) && !targetEnvironment(macCatalyst)
     override func loadView() {
-        let view = TaggableView()
-        view.tag = id.hashValue
-        self.view = view
+        view = NSView()
     }
     #endif
 }
@@ -78,14 +84,17 @@ struct IntrospectionView<Target: PlatformEntity>: PlatformViewControllerRepresen
 
     @Binding
     private var observed: Void // workaround for state changes not triggering view updates
+    private let id: IntrospectionViewID
     private let selector: (IntrospectionPlatformViewController) -> Target?
     private let customize: (Target) -> Void
 
     init(
+        id: IntrospectionViewID,
         selector: @escaping (IntrospectionPlatformViewController) -> Target?,
         customize: @escaping (Target) -> Void
     ) {
         self._observed = .constant(())
+        self.id = id
         self.selector = selector
         self.customize = customize
     }
@@ -95,7 +104,7 @@ struct IntrospectionView<Target: PlatformEntity>: PlatformViewControllerRepresen
     }
 
     func makePlatformViewController(context: Context) -> IntrospectionPlatformViewController {
-        let controller = IntrospectionPlatformViewController { controller in
+        let controller = IntrospectionPlatformViewController(id: id) { controller in
             guard let target = selector(controller) else {
                 return
             }
@@ -128,16 +137,22 @@ struct IntrospectionView<Target: PlatformEntity>: PlatformViewControllerRepresen
 }
 
 final class IntrospectionPlatformViewController: PlatformViewController {
+    let id: IntrospectionViewID
     var handler: (() -> Void)? = nil
 
-    fileprivate init(handler: ((IntrospectionPlatformViewController) -> Void)?) {
+    fileprivate init(
+        id: IntrospectionViewID,
+        handler: ((IntrospectionPlatformViewController) -> Void)?
+    ) {
+        self.id = id
         super.init(nibName: nil, bundle: nil)
         self.handler = { [weak self] in
             guard let self = self else {
                 return
             }
             handler?(self)
         }
+        IntrospectionStore.shared[id, default: .init()].controller = self
     }
 
     @available(*, unavailable)
@@ -146,13 +161,14 @@ final class IntrospectionPlatformViewController: PlatformViewController {
     }
 
     #if canImport(UIKit)
-    override func didMove(toParent parent: UIViewController?) {
-        super.didMove(toParent: parent)
+    override func viewDidLoad() {
+        super.viewDidLoad()
+        view.introspectionController = self
         handler?()
     }
 
-    override func viewDidLoad() {
-        super.viewDidLoad()
+    override func didMove(toParent parent: UIViewController?) {
+        super.didMove(toParent: parent)
         handler?()
     }
 
@@ -168,6 +184,7 @@ final class IntrospectionPlatformViewController: PlatformViewController {
     #elseif canImport(AppKit)
     override func loadView() {
         view = NSView()
+        view.introspectionController = self
     }
 
     override func viewDidLoad() {
@@ -181,3 +198,18 @@ final class IntrospectionPlatformViewController: PlatformViewController {
     }
     #endif
 }
+
+import ObjectiveC
+
+extension PlatformView {
+    fileprivate var introspectionController: IntrospectionPlatformViewController? {
+        get {
+            let key = unsafeBitCast(Selector(#function), to: UnsafeRawPointer.self)
+            return objc_getAssociatedObject(self, key) as? IntrospectionPlatformViewController
+        }
+        set {
+            let key = unsafeBitCast(Selector(#function), to: UnsafeRawPointer.self)
+            objc_setAssociatedObject(self, key, newValue, .OBJC_ASSOCIATION_ASSIGN)
+        }
+    }
+}