Custom Transaction Execution (#212)

* doesnt work * works * doesn't work This reverts commit f2eb2d5. * wip * wip * Revert "wip" This reverts commit dff88f5. * Revert "wip" This reverts commit 361345c. * Revert "doesn't work" This reverts commit 64a9c38. * wip * wip * wip * wip * wip * clean up tests * clean up * fix test --------- Co-authored-by: Brandon Williams <[email protected]>
pointfreeco · Aug 22, 2024 · 4d04eb0 · 4d04eb0
1 parent 18d5db4
commit 4d04eb0
Show file tree

Hide file tree

Showing 11 changed files with 304 additions and 46 deletions.
diff --git a/Package.resolved b/Package.resolved
@@ -9,6 +9,15 @@
         "version" : "1.5.4"
       }
     },
+    {
+      "identity" : "swift-collections",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-collections",
+      "state" : {
+        "revision" : "3d2dc41a01f9e49d84f0a3925fb858bed64f702d",
+        "version" : "1.1.2"
+      }
+    },
     {
       "identity" : "swift-concurrency-extras",
       "kind" : "remoteSourceControl",

diff --git a/Package.swift b/Package.swift
@@ -25,6 +25,7 @@ let package = Package(
     ),
   ],
   dependencies: [
+    .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"),
     .package(url: "https://github.com/swiftlang/swift-docc-plugin", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-case-paths", from: "1.5.4"),
     .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.1.0"),
@@ -39,6 +40,7 @@ let package = Package(
         .product(name: "CasePaths", package: "swift-case-paths"),
         .product(name: "CustomDump", package: "swift-custom-dump"),
         .product(name: "ConcurrencyExtras", package: "swift-concurrency-extras"),
+        .product(name: "OrderedCollections", package: "swift-collections"),
         .product(name: "Perception", package: "swift-perception"),
       ]
     ),

diff --git a/[email protected] b/[email protected]
@@ -25,6 +25,7 @@ let package = Package(
     ),
   ],
   dependencies: [
+    .package(url: "https://github.com/apple/swift-collections", from: "1.0.0"),
     .package(url: "https://github.com/swiftlang/swift-docc-plugin", from: "1.0.0"),
     .package(url: "https://github.com/pointfreeco/swift-case-paths", from: "1.5.4"),
     .package(url: "https://github.com/pointfreeco/swift-concurrency-extras", from: "1.1.0"),
@@ -39,6 +40,7 @@ let package = Package(
         .product(name: "CasePaths", package: "swift-case-paths"),
         .product(name: "CustomDump", package: "swift-custom-dump"),
         .product(name: "ConcurrencyExtras", package: "swift-concurrency-extras"),
+        .product(name: "OrderedCollections", package: "swift-collections"),
         .product(name: "Perception", package: "swift-perception"),
       ]
     ),

diff --git a/...tNavigation/Internal/AssumeIsolated.swift → ...tNavigation/Internal/AssumeIsolated.swift b/...tNavigation/Internal/AssumeIsolated.swift → ...tNavigation/Internal/AssumeIsolated.swift
@@ -2,7 +2,7 @@ import Foundation
 
 extension MainActor {
   // NB: This functionality was not back-deployed in Swift 5.9
-  static func _assumeIsolated<T: Sendable>(
+  package static func _assumeIsolated<T: Sendable>(
     _ operation: @MainActor () throws -> T,
     file: StaticString = #fileID,
     line: UInt = #line

diff --git a/Sources/UIKitNavigation/Observe.swift → ...es/SwiftNavigation/NSObject+Observe.swift b/Sources/UIKitNavigation/Observe.swift → ...es/SwiftNavigation/NSObject+Observe.swift
@@ -1,6 +1,6 @@
-#if canImport(UIKit)
-  @_spi(Internals) import SwiftNavigation
-  import UIKit
+#if canImport(ObjectiveC)
+  import Dispatch
+  import ObjectiveC
 
   @MainActor
   extension NSObject {
@@ -123,34 +123,7 @@
     ) -> ObserveToken {
       let token = SwiftNavigation.observe { transaction in
         MainActor._assumeIsolated {
-          withUITransaction(transaction) {
-            #if os(watchOS)
-              apply(transaction)
-            #else
-              if transaction.uiKit.disablesAnimations {
-                UIView.performWithoutAnimation { apply(transaction) }
-                for completion in transaction.uiKit.animationCompletions {
-                  completion(true)
-                }
-              } else if let animation = transaction.uiKit.animation {
-                return animation.perform(
-                  { apply(transaction) },
-                  completion: transaction.uiKit.animationCompletions.isEmpty
-                    ? nil
-                    : {
-                      for completion in transaction.uiKit.animationCompletions {
-                        completion($0)
-                      }
-                    }
-                )
-              } else {
-                apply(transaction)
-                for completion in transaction.uiKit.animationCompletions {
-                  completion(true)
-                }
-              }
-            #endif
-          }
+          apply(transaction)
         }
       } task: { transaction, work in
         DispatchQueue.main.async {

diff --git a/Sources/SwiftNavigation/Observe.swift b/Sources/SwiftNavigation/Observe.swift
@@ -102,8 +102,7 @@ private actor ActorProxy {
   }
 }
 
-@_spi(Internals)
-public func observe(
+func observe(
   _ apply: @escaping @Sendable (_ transaction: UITransaction) -> Void,
   task: @escaping @Sendable (
     _ transaction: UITransaction, _ operation: @escaping @Sendable () -> Void
@@ -118,7 +117,21 @@ public func observe(
         let token,
         !token.isCancelled
       else { return }
-      apply(transaction)
+
+      var perform: @Sendable () -> Void = { apply(transaction) }
+      for key in transaction.storage.keys {
+        guard let keyType = key.keyType as? any _UICustomTransactionKey.Type
+        else { continue }
+        func open<K: _UICustomTransactionKey>(_: K.Type) {
+          perform = { [perform] in
+            K.perform(value: transaction[K.self]) {
+              perform()
+            }
+          }
+        }
+        open(keyType)
+      }
+      perform()
     },
     task: task
   )

diff --git a/Sources/SwiftNavigation/UITransaction.swift b/Sources/SwiftNavigation/UITransaction.swift
@@ -1,3 +1,5 @@
+import OrderedCollections
+
 /// Executes a closure with the specified transaction and returns the result.
 ///
 /// - Parameters:
@@ -8,7 +10,10 @@ public func withUITransaction<Result>(
   _ transaction: UITransaction,
   _ body: () throws -> Result
 ) rethrows -> Result {
-  try UITransaction.$current.withValue(transaction, operation: body)
+  try UITransaction.$current.withValue(
+    UITransaction.current.merging(transaction),
+    operation: body
+  )
 }
 
 /// Executes a closure with the specified transaction key path and value and returns the result.
@@ -24,7 +29,7 @@ public func withUITransaction<R, V>(
   _ value: V,
   _ body: () throws -> R
 ) rethrows -> R {
-  var transaction = UITransaction.current
+  var transaction = UITransaction()
   transaction[keyPath: keyPath] = value
   return try withUITransaction(transaction, body)
 }
@@ -36,7 +41,7 @@ public func withUITransaction<R, V>(
 public struct UITransaction: Sendable {
   @TaskLocal package static var current = Self()
 
-  private var storage: [Key: any Sendable] = [:]
+  var storage: OrderedDictionary<Key, any Sendable> = [:]
 
   /// Creates a transaction.
   public init() {}
@@ -68,7 +73,15 @@ public struct UITransaction: Sendable {
     storage.isEmpty
   }
 
-  private struct Key: Hashable {
+  fileprivate func merging(_ other: Self) -> Self {
+    Self(storage: storage.merging(other.storage, uniquingKeysWith: { $1 }))
+  }
+
+  private init(storage: OrderedDictionary<Key, any Sendable>) {
+    self.storage = storage
+  }
+
+  struct Key: Hashable {
     let keyType: Any.Type
     init<K: UITransactionKey>(_ keyType: K.Type) {
       self.keyType = keyType
@@ -92,3 +105,10 @@ public protocol UITransactionKey {
   /// The default value for the transaction key.
   static var defaultValue: Value { get }
 }
+
+public protocol _UICustomTransactionKey: UITransactionKey, Sendable {
+  static func perform(
+    value: Value,
+    operation: @Sendable () -> Void
+  )
+}
diff --git a/Sources/UIKitNavigation/UITransaction.swift b/Sources/UIKitNavigation/UITransaction.swift
@@ -1,4 +1,7 @@
 #if canImport(UIKit) && !os(watchOS)
+  import SwiftNavigation
+  import UIKit
+
   extension UITransaction {
     /// Creates a transaction and assigns its animation property.
     ///
@@ -14,8 +17,42 @@
       set { self[UIKitKey.self] = newValue }
     }
 
-    private enum UIKitKey: UITransactionKey {
+    private enum UIKitKey: _UICustomTransactionKey {
       static let defaultValue = UIKit()
+
+      static func perform(
+        value: UIKit,
+        operation: @Sendable () -> Void
+      ) {
+        MainActor._assumeIsolated {
+          #if os(watchOS)
+            operation()
+          #else
+            if value.disablesAnimations {
+              UIView.performWithoutAnimation { operation() }
+              for completion in value.animationCompletions {
+                completion(true)
+              }
+            } else if let animation = value.animation {
+              return animation.perform(
+                { operation() },
+                completion: value.animationCompletions.isEmpty
+                  ? nil
+                  : {
+                    for completion in value.animationCompletions {
+                      completion($0)
+                    }
+                  }
+              )
+            } else {
+              operation()
+              for completion in value.animationCompletions {
+                completion(true)
+              }
+            }
+          #endif
+        }
+      }
     }
 
     /// UIKit-specific data associated with a ``UITransaction``.

diff --git a/Tests/SwiftNavigationTests/ObserveTests.swift b/Tests/SwiftNavigationTests/ObserveTests.swift
@@ -1,16 +1,28 @@
-#if swift(>=6)
-  import SwiftNavigation
-  import XCTest
+import SwiftNavigation
+import XCTest
 
-  class ObserveTests: XCTestCase {
+class ObserveTests: XCTestCase {
+  #if swift(>=6)
     @MainActor
-    func testCompiles() {
+    func testIsolation() {
       var count = 0
       let token = SwiftNavigation.observe {
         count = 1
       }
       XCTAssertEqual(count, 1)
       _ = token
     }
+  #endif
+
+  @MainActor
+  func testTokenStorage() {
+    var count = 0
+    observe {
+      count += 1
+    }
+    observe {
+      count += 1
+    }
+    XCTAssertEqual(count, 2)
   }
-#endif
+}