Reduce key path COW copies in Graph forEach and compactMapValues

Sources/Testing/Support/Graph.swift

@@ -395,12 +395,12 @@ extension Graph {
   ///     path and leaf value of each node are passed to the closure.
   ///
   /// - Throws: Whatever is thrown by `body`.
-  private func _forEach<E>(keyPath: [K], _ body: (Element) throws(E) -> Void) throws(E) {
+  private func _forEach<E>(keyPath: inout [K], _ body: (Element) throws(E) -> Void) throws(E) {
     try body((keyPath, value))
     for (key, child) in children {
-      var childKeyPath = keyPath
-      childKeyPath.append(key)
-      try child._forEach(keyPath: childKeyPath, body)
+      keyPath.append(key)
+      try child._forEach(keyPath: &keyPath, body)
+      keyPath.removeLast()
     }
   }
 
@@ -413,12 +413,12 @@ extension Graph {
   ///     key path and leaf value of each node are passed to the closure.
   ///
   /// - Throws: Whatever is thrown by `body`.
-  private func _forEach<E>(keyPath: [K], _ body: (Element) async throws(E) -> Void) async throws(E) {
+  private func _forEach<E>(keyPath: inout [K], _ body: (Element) async throws(E) -> Void) async throws(E) {
     try await body((keyPath, value))
     for (key, child) in children {
-      var childKeyPath = keyPath
-      childKeyPath.append(key)
-      try await child._forEach(keyPath: childKeyPath, body)
+      keyPath.append(key)
+      try await child._forEach(keyPath: &keyPath, body)
+      keyPath.removeLast()
     }
   }
 
@@ -432,7 +432,8 @@ extension Graph {
   ///
   /// This function iterates depth-first.
   func forEach<E>(_ body: (Element) throws(E) -> Void) throws(E) {
-    try _forEach(keyPath: []) { (element) throws(E) in
+    var keyPath: [K] = []
+    try _forEach(keyPath: &keyPath) { (element) throws(E) in
       try body(element)
     }
   }
@@ -447,7 +448,8 @@ extension Graph {
   ///
   /// This function iterates depth-first.
   func forEach<E>(_ body: (Element) async throws(E) -> Void) async throws(E) {
-    try await _forEach(keyPath: []) { (element) async throws(E) in
+    var keyPath: [K] = []
+    try await _forEach(keyPath: &keyPath) { (element) async throws(E) in
       try await body(element)
     }
   }
@@ -546,7 +548,8 @@ extension Graph {
   ///
   /// This function iterates depth-first.
   func compactMapValues<U, E>(_ transform: (Element) throws(E) -> (U, recursivelyApply: Bool)?) throws(E) -> Graph<K, U>? {
-    try _compactMapValues(keyPath: []) { (element) throws(E) in
+    var keyPath: [K] = []
+    return try _compactMapValues(keyPath: &keyPath) { (element) throws(E) in
       try transform(element)
     }
   }
@@ -565,21 +568,21 @@ extension Graph {
   ///     child nodes are omitted from the new graph.
   ///
   /// - Throws: Whatever is thrown by `transform`.
-  private func _compactMapValues<U, E>(keyPath: [K], _ transform: (Element) throws(E) -> (U, recursivelyApply: Bool)?) throws(E) -> Graph<K, U>? {
+  private func _compactMapValues<U, E>(keyPath: inout [K], _ transform: (Element) throws(E) -> (U, recursivelyApply: Bool)?) throws(E) -> Graph<K, U>? {
     guard let (newValue, recursivelyApply) = try transform((keyPath, value)) else {
       return nil
     }
 
     var newChildren = [K: Graph<K,U>]()
     newChildren.reserveCapacity(children.count)
     for (key, child) in children {
-      var childKeyPath = keyPath
-      childKeyPath.append(key)
+      keyPath.append(key)
+      defer { keyPath.removeLast() }
 
       if recursivelyApply {
-        newChildren[key] = child._compactMapValues(keyPath: childKeyPath) { _ in (newValue, true) }
+        newChildren[key] = child._compactMapValues(keyPath: &keyPath) { _ in (newValue, true) }
       } else {
-        newChildren[key] = try child._compactMapValues(keyPath: childKeyPath, transform)
+        newChildren[key] = try child._compactMapValues(keyPath: &keyPath, transform)
       }
     }
 
@@ -606,7 +609,8 @@ extension Graph {
   ///
   /// This function iterates depth-first.
   func compactMapValues<U, E>(_ transform: (Element) async throws(E) -> (U, recursivelyApply: Bool)?) async throws(E) -> Graph<K, U>? {
-    try await _compactMapValues(keyPath: []) { (element) async throws(E) in
+    var keyPath: [K] = []
+    return try await _compactMapValues(keyPath: &keyPath) { (element) async throws(E) in
       try await transform(element)
     }
   }
@@ -625,21 +629,21 @@ extension Graph {
   ///     child nodes are omitted from the new graph.
   ///
   /// - Throws: Whatever is thrown by `transform`.
-  private func _compactMapValues<U, E>(keyPath: [K], _ transform: (Element) async throws(E) -> (U, recursivelyApply: Bool)?) async throws(E) -> Graph<K, U>? {
+  private func _compactMapValues<U, E>(keyPath: inout [K], _ transform: (Element) async throws(E) -> (U, recursivelyApply: Bool)?) async throws(E) -> Graph<K, U>? {
     guard let (newValue, recursivelyApply) = try await transform((keyPath, value)) else {
       return nil
     }
 
     var newChildren = [K: Graph<K,U>]()
     newChildren.reserveCapacity(children.count)
     for (key, child) in children {
-      var childKeyPath = keyPath
-      childKeyPath.append(key)
+      keyPath.append(key)
+      defer { keyPath.removeLast() }
 
       if recursivelyApply {
-        newChildren[key] = child._compactMapValues(keyPath: childKeyPath) { _ in (newValue, true) }
+        newChildren[key] = child._compactMapValues(keyPath: &keyPath) { _ in (newValue, true) }
       } else {
-        newChildren[key] = try await child._compactMapValues(keyPath: childKeyPath, transform)
+        newChildren[key] = try await child._compactMapValues(keyPath: &keyPath, transform)
       }
     }
 

Tests/TestingTests/Support/GraphTests.swift

@@ -588,4 +588,53 @@ struct GraphTests {
     // Standard deviation should be less than 25% of mean
     #expect(standardDeviation < mean * 0.25)
   }
+
+  @Test(.enabled(if: performanceTestsEnabled))
+  func `forEach traversal scales linearly with node count`() {
+    // This test verifies that `forEach` traversal time scales ~linearly with
+    // the number of nodes, regardless of depth. Previously, each recursion
+    // copied the parent key path before appending the child key, which made
+    // per-node traversal cost grow with depth (O(n * d) overall). With the
+    // shared inout key path, per-node cost stays constant (O(n) overall).
+
+    var perItemTimes = [Int: Double]()
+    let clock = Test.Clock()
+
+    for itemCount in [500, 1000, 5000, 10000] {
+      var graph = Graph<String, Int?>()
+
+      // Build a tree resembling a realistic test plan: a shared module/suite
+      // prefix with many leaf entries underneath. This means most nodes are
+      // visited at the maximum depth, where the old copying behavior was
+      // most expensive.
+      let sharedPrefix = ["ProjectTests", "ProjectSuite", "something(_:)"]
+      graph[sharedPrefix] = -1
+
+      for i in 0..<itemCount {
+        var path = sharedPrefix
+        path.append("Argument \(i)")
+        graph[path] = i
+      }
+
+      let startTime = clock.now
+
+      // Traverse the whole tree once. The body deliberately consumes the key
+      // path so the compiler can't elide it.
+      var visitedCount = 0
+      graph.forEach { keyPath, _ in
+        visitedCount &+= keyPath.count
+      }
+      #expect(visitedCount > 0)
+
+      let elapsedTime = startTime.duration(to: clock.now) / .seconds(1)
+      perItemTimes[itemCount] = elapsedTime / Double(itemCount)
+    }
+
+    let mean = perItemTimes.values.reduce(0, +) / Double(perItemTimes.count)
+    let variance = perItemTimes.values.reduce(0) { $0 + ($1 - mean) * ($1 - mean) }
+    let standardDeviation = (variance / (Double(perItemTimes.count) - 1)).squareRoot()
+
+    // Standard deviation should be less than 25% of mean
+    #expect(standardDeviation < mean * 0.25)
+  }
 }