[Observation] Optimize the storage of registrar entries, provide KeyPath caching, and uniqueness notification (#78151)

Author: phausler

lib/Macros/Sources/ObservationMacros/ObservableMacro.swift

@@ -36,35 +36,70 @@ public struct ObservableMacro {
 
   static let registrarVariableName = "_$observationRegistrar"
 
-  static func registrarVariable(_ observableType: TokenSyntax) -> DeclSyntax {
+  static func registrarVariable(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
     return
       """
       @\(raw: ignoredMacroName) private let \(raw: registrarVariableName) = \(raw: qualifiedRegistrarTypeName)()
       """
   }
 
-  static func accessFunction(_ observableType: TokenSyntax) -> DeclSyntax {
-    return 
+  static func accessFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    return
       """
-      internal nonisolated func access<Member>(
-      keyPath: KeyPath<\(observableType), Member>
+      internal nonisolated func access<\(memberGeneric)>(
+        keyPath: KeyPath<\(observableType), \(memberGeneric)>
       ) {
-      \(raw: registrarVariableName).access(self, keyPath: keyPath)
+        \(raw: registrarVariableName).access(self, keyPath: keyPath)
       }
       """
   }
 
-  static func withMutationFunction(_ observableType: TokenSyntax) -> DeclSyntax {
-    return 
+  static func withMutationFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    let mutationGeneric = context.makeUniqueName("MutationResult")
+    return
       """
-      internal nonisolated func withMutation<Member, MutationResult>(
-      keyPath: KeyPath<\(observableType), Member>,
-      _ mutation: () throws -> MutationResult
-      ) rethrows -> MutationResult {
-      try \(raw: registrarVariableName).withMutation(of: self, keyPath: keyPath, mutation)
+      internal nonisolated func withMutation<\(memberGeneric), \(mutationGeneric)>(
+        keyPath: KeyPath<\(observableType), \(memberGeneric)>,
+        _ mutation: () throws -> \(mutationGeneric)
+      ) rethrows -> \(mutationGeneric) {
+        try \(raw: registrarVariableName).withMutation(of: self, keyPath: keyPath, mutation)
       }
       """
   }
+  
+  static func shouldNotifyObserversNonEquatableFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    return
+      """
+       private nonisolated func shouldNotifyObservers<\(memberGeneric)>(_ lhs: \(memberGeneric), _ rhs: \(memberGeneric)) -> Bool { true }
+      """
+  }
+  
+  static func shouldNotifyObserversEquatableFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    return
+      """
+      private nonisolated func shouldNotifyObservers<\(memberGeneric): Equatable>(_ lhs: \(memberGeneric), _ rhs: \(memberGeneric)) -> Bool { lhs != rhs }
+      """
+  }
+  
+  static func shouldNotifyObserversNonEquatableObjectFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    return
+      """
+       private nonisolated func shouldNotifyObservers<\(memberGeneric): AnyObject>(_ lhs: \(memberGeneric), _ rhs: \(memberGeneric)) -> Bool { lhs !== rhs }
+      """
+  }
+
+  static func shouldNotifyObserversEquatableObjectFunction(_ observableType: TokenSyntax, context: some MacroExpansionContext) -> DeclSyntax {
+    let memberGeneric = context.makeUniqueName("Member")
+    return
+      """
+      private nonisolated func shouldNotifyObservers<\(memberGeneric): Equatable & AnyObject>(_ lhs: \(memberGeneric), _ rhs: \(memberGeneric)) -> Bool { lhs != rhs }
+      """
+  }
 
   static var ignoredAttribute: AttributeSyntax {
     AttributeSyntax(
@@ -220,9 +255,13 @@ extension ObservableMacro: MemberMacro {
 
     var declarations = [DeclSyntax]()
 
-    declaration.addIfNeeded(ObservableMacro.registrarVariable(observableType), to: &declarations)
-    declaration.addIfNeeded(ObservableMacro.accessFunction(observableType), to: &declarations)
-    declaration.addIfNeeded(ObservableMacro.withMutationFunction(observableType), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.registrarVariable(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.accessFunction(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.withMutationFunction(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.shouldNotifyObserversNonEquatableFunction(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.shouldNotifyObserversEquatableFunction(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.shouldNotifyObserversNonEquatableObjectFunction(observableType, context: context), to: &declarations)
+    declaration.addIfNeeded(ObservableMacro.shouldNotifyObserversEquatableObjectFunction(observableType, context: context), to: &declarations)
 
     return declarations
   }
@@ -298,6 +337,10 @@ public struct ObservationTrackedMacro: AccessorMacro {
           let identifier = property.identifier?.trimmed else {
       return []
     }
+    
+    guard let container = context.lexicalContext[0].as(ClassDeclSyntax.self) else {
+      return []
+    }
 
     if property.hasMacroApplication(ObservableMacro.ignoredMacroName) {
       return []
@@ -307,34 +350,46 @@ public struct ObservationTrackedMacro: AccessorMacro {
       """
       @storageRestrictions(initializes: _\(identifier))
       init(initialValue) {
-      _\(identifier) = initialValue
+        _\(identifier) = initialValue
       }
       """
 
     let getAccessor: AccessorDeclSyntax =
       """
       get {
-      access(keyPath: \\.\(identifier))
-      return _\(identifier)
+        access(keyPath: \(container.trimmed.name)._cachedKeypath_\(identifier))
+        return _\(identifier)
       }
       """
 
     let setAccessor: AccessorDeclSyntax =
       """
       set {
-      withMutation(keyPath: \\.\(identifier)) {
-      _\(identifier) = newValue
-      }
+        guard shouldNotifyObservers(_\(identifier), newValue) else {
+          return
+        }
+        withMutation(keyPath: \(container.trimmed.name)._cachedKeypath_\(identifier)) {
+          _\(identifier) = newValue
+        }
       }
       """
 
+    // Note: this accessor cannot test the equality since it would incur
+    // additional CoW's on structural types. Most mutations in-place do
+    // not leave the value equal so this is "fine"-ish.
+    // Warning to future maintence: adding equality checks here can make
+    // container mutation O(N) instead of O(1).
+    // e.g. observable.array.append(element) should just emit a change
+    // to the new array, and NOT cause a copy of each element of the
+    // array to an entirely new array.
     let modifyAccessor: AccessorDeclSyntax =
       """
       _modify {
-      access(keyPath: \\.\(identifier))
-      \(raw: ObservableMacro.registrarVariableName).willSet(self, keyPath: \\.\(identifier))
-      defer { \(raw: ObservableMacro.registrarVariableName).didSet(self, keyPath: \\.\(identifier)) } 
-      yield &_\(identifier)
+        let keyPath = \(container.trimmed.name)._cachedKeypath_\(identifier)
+        access(keyPath: keyPath)
+        \(raw: ObservableMacro.registrarVariableName).willSet(self, keyPath: keyPath)
+        defer { \(raw: ObservableMacro.registrarVariableName).didSet(self, keyPath: keyPath) }
+        yield &_\(identifier)
       }
       """
 
@@ -352,7 +407,12 @@ extension ObservationTrackedMacro: PeerMacro {
     in context: Context
   ) throws -> [DeclSyntax] {
     guard let property = declaration.as(VariableDeclSyntax.self),
-          property.isValidForObservation else {
+          property.isValidForObservation,
+          let identifier = property.identifier?.trimmed else {
+      return []
+    }
+    
+    guard let container = context.lexicalContext[0].as(ClassDeclSyntax.self) else {
       return []
     }
 
@@ -362,7 +422,11 @@ extension ObservationTrackedMacro: PeerMacro {
     }
 
     let storage = DeclSyntax(property.privatePrefixed("_", addingAttribute: ObservableMacro.ignoredAttribute))
-    return [storage]
+    let cachedKeypath: DeclSyntax =
+      """
+      private static let _cachedKeypath_\(identifier) = \\\(container.name).\(identifier)
+      """
+    return [storage, cachedKeypath]
   }
 }
 

stdlib/public/Observation/Sources/Observation/Observable.swift

@@ -39,7 +39,7 @@ public protocol Observable { }
 ///        }
 ///     }
 @available(SwiftStdlib 5.9, *)
-@attached(member, names: named(_$observationRegistrar), named(access), named(withMutation))
+@attached(member, names: named(_$observationRegistrar), named(access), named(withMutation), named(shouldNotifyObservers))
 @attached(memberAttribute)
 @attached(extension, conformances: Observable)
 public macro Observable() =
@@ -51,7 +51,7 @@ public macro Observable() =
 /// framework isn't necessary.
 @available(SwiftStdlib 5.9, *)
 @attached(accessor, names: named(init), named(get), named(set), named(_modify))
-@attached(peer, names: prefixed(_))
+@attached(peer, names: prefixed(_), prefixed(_cachedKeypath_))
 public macro ObservationTracked() =
   #externalMacro(module: "ObservationMacros", type: "ObservationTrackedMacro")
 
@@ -61,7 +61,7 @@ public macro ObservationTracked() =
 /// is accessible to the observing object. To prevent observation of an
 /// accessible property, attach the `ObservationIgnored` macro to the property.
 @available(SwiftStdlib 5.9, *)
-@attached(accessor, names: named(willSet))
+@attached(accessor)
 public macro ObservationIgnored() =
   #externalMacro(module: "ObservationMacros", type: "ObservationIgnoredMacro")
 

stdlib/public/Observation/Sources/Observation/ObservationRegistrar.swift

@@ -40,8 +40,6 @@ public struct ObservationRegistrar: Sendable {
     private enum ObservationKind {
       case willSetTracking(@Sendable (AnyKeyPath) -> Void)
       case didSetTracking(@Sendable (AnyKeyPath) -> Void)
-      case computed(@Sendable (Any) -> Void)
-      case values(ValuesObserver)
     }
 
     private struct Observation {
@@ -70,42 +68,6 @@ public struct ObservationRegistrar: Sendable {
           return nil
         }
       }
-      
-      var observer: (@Sendable (Any) -> Void)? {
-        switch kind {
-        case .computed(let observer):
-          return observer
-        default:
-          return nil
-        }
-      }
-      
-      var isValueObserver: Bool {
-        switch kind {
-        case .values:
-          return true
-        default:
-          return false
-        }
-      }
-      
-      func emit<Element>(_ value: Element) -> Bool {
-        switch kind {
-        case .values(let observer):
-          return observer.emit(value)
-        default:
-          return false
-        }
-      }
-      
-      func cancel() {
-        switch kind {
-        case .values(let observer):
-          observer.cancel()
-        default:
-          break
-        }
-      }
     }
 
     private var id = 0
@@ -135,31 +97,6 @@ public struct ObservationRegistrar: Sendable {
       return id
     }
 
-    internal mutating func registerComputedValues(for properties: Set<AnyKeyPath>, observer: @Sendable @escaping (Any) -> Void) -> Int {
-      let id = generateId()
-      observations[id] = Observation(kind: .computed(observer), properties: properties)
-      for keyPath in properties {
-        lookups[keyPath, default: []].insert(id)
-      }
-      return id
-    }
-    
-    internal mutating func registerValues(for properties: Set<AnyKeyPath>, storage: ValueObservationStorage) -> Int {
-      let id = generateId()
-      observations[id] = Observation(kind: .values(ValuesObserver(storage: storage)), properties: properties)
-      for keyPath in properties {
-        lookups[keyPath, default: []].insert(id)
-      }
-      return id
-    }
-    
-    internal func valueObservers(for keyPath: AnyKeyPath) -> Set<Int> {
-      guard let ids = lookups[keyPath] else {
-        return []
-      }
-      return ids.filter { observations[$0]?.isValueObserver == true }
-    }
-    
     internal mutating func cancel(_ id: Int) {
       if let observation = observations.removeValue(forKey: id) {
         for keyPath in observation.properties {
@@ -170,14 +107,10 @@ public struct ObservationRegistrar: Sendable {
             }
           }
         }
-        observation.cancel()
       }
     }
 
     internal mutating func cancelAll() {
-      for observation in observations.values {
-        observation.cancel()
-      }
       observations.removeAll()
       lookups.removeAll()
     }
@@ -194,29 +127,16 @@ public struct ObservationRegistrar: Sendable {
       return trackers
     }
 
-    internal mutating func didSet<Subject: Observable, Member>(keyPath: KeyPath<Subject, Member>) -> ([@Sendable (Any) -> Void], [@Sendable (AnyKeyPath) -> Void]) {
-      var observers = [@Sendable (Any) -> Void]()
+    internal mutating func didSet<Subject: Observable, Member>(keyPath: KeyPath<Subject, Member>) -> [@Sendable (AnyKeyPath) -> Void] {
       var trackers = [@Sendable (AnyKeyPath) -> Void]()
       if let ids = lookups[keyPath] {
         for id in ids {
-          if let observer = observations[id]?.observer {
-            observers.append(observer)
-            cancel(id)
-          }
           if let tracker = observations[id]?.didSetTracker {
             trackers.append(tracker)
           }
         }
       }
-      return (observers, trackers)
-    }
-    
-    internal mutating func emit<Element>(_ value: Element, ids: Set<Int>) {
-      for id in ids {
-        if observations[id]?.emit(value) == true {
-          cancel(id)
-        }
-      }
+      return trackers
     }
   }
 
@@ -233,14 +153,6 @@ public struct ObservationRegistrar: Sendable {
       state.withCriticalRegion { $0.registerTracking(for: properties, didSet: observer) }
     }
 
-    internal func registerComputedValues(for properties: Set<AnyKeyPath>, observer: @Sendable @escaping (Any) -> Void) -> Int {
-      state.withCriticalRegion { $0.registerComputedValues(for: properties, observer: observer) }
-    }
-    
-    internal func registerValues(for properties: Set<AnyKeyPath>, storage: ValueObservationStorage) -> Int {
-      state.withCriticalRegion { $0.registerValues(for: properties, storage: storage) }
-    }
-    
     internal func cancel(_ id: Int) {
       state.withCriticalRegion { $0.cancel(id) }
     }
@@ -263,17 +175,10 @@ public struct ObservationRegistrar: Sendable {
       _ subject: Subject,
       keyPath: KeyPath<Subject, Member>
     ) {
-      let (ids, (actions, tracking)) = state.withCriticalRegion { ($0.valueObservers(for: keyPath), $0.didSet(keyPath: keyPath)) }
-      if !ids.isEmpty {
-        let value = subject[keyPath: keyPath]
-        state.withCriticalRegion { $0.emit(value, ids: ids) }
-      }
+      let tracking = state.withCriticalRegion { $0.didSet(keyPath: keyPath) }
       for action in tracking {
         action(keyPath)
       }
-      for action in actions {
-        action(subject)
-      }
     }
   }
 

stdlib/public/Observation/Sources/Observation/ObservationTracking.swift

@@ -115,7 +115,7 @@ public struct ObservationTracking: Sendable {
     })
   }
 
-  struct State {
+  struct State: @unchecked Sendable {
     var values = [ObjectIdentifier: ObservationTracking.Id]()
     var cancelled = false
     var changed: AnyKeyPath?