From b678c7d8a8b0f9f98f01975cf20ea4ac368a91a1 Mon Sep 17 00:00:00 2001
From: Owen Voorhees <ovoorhees@apple.com>
Date: Mon, 3 Jun 2024 13:35:11 -0700
Subject: [PATCH] Improve performance of precompiled module invalidation in
 incremental builds

---
 .../FirstWaveComputer.swift                   | 58 ++++++-------------
 1 file changed, 18 insertions(+), 40 deletions(-)

diff --git a/Sources/SwiftDriver/IncrementalCompilation/FirstWaveComputer.swift b/Sources/SwiftDriver/IncrementalCompilation/FirstWaveComputer.swift
index 15cc6887f..aa94cfdac 100644
--- a/Sources/SwiftDriver/IncrementalCompilation/FirstWaveComputer.swift
+++ b/Sources/SwiftDriver/IncrementalCompilation/FirstWaveComputer.swift
@@ -149,12 +149,11 @@ extension IncrementalCompilationState.FirstWaveComputer {
   throws -> Set<ModuleDependencyId> {
     let mainModuleInfo = moduleDependencyGraph.mainModule
     var modulesRequiringRebuild: Set<ModuleDependencyId> = []
-    var validatedModules: Set<ModuleDependencyId> = []
+    var visitedModules: Set<ModuleDependencyId> = []
     // Scan from the main module's dependencies to avoid reporting
     // the main module itself in the results.
     for dependencyId in mainModuleInfo.directDependencies ?? [] {
-      try outOfDateModuleScan(on: moduleDependencyGraph, from: dependencyId,
-                              pathSoFar: [], visitedValidated: &validatedModules,
+      try outOfDateModuleScan(on: moduleDependencyGraph, from: dependencyId, visited: &visitedModules,
                               modulesRequiringRebuild: &modulesRequiringRebuild)
     }
 
@@ -167,54 +166,33 @@ extension IncrementalCompilationState.FirstWaveComputer {
   /// modules, along with the path from the root to this module.
   private func outOfDateModuleScan(on moduleDependencyGraph: InterModuleDependencyGraph,
                                    from moduleId: ModuleDependencyId,
-                                   pathSoFar: [ModuleDependencyId],
-                                   visitedValidated: inout Set<ModuleDependencyId>,
+                                   visited: inout Set<ModuleDependencyId>,
                                    modulesRequiringRebuild: inout Set<ModuleDependencyId>) throws {
     let moduleInfo = try moduleDependencyGraph.moduleInfo(of: moduleId)
-    let isMainModule = moduleId == .swift(moduleDependencyGraph.mainModuleName)
-
-    // Routine to invalidate the path from root to this node
-    let invalidatePath = { (modulesRequiringRebuild: inout Set<ModuleDependencyId>) in
-      if let reporter {
-        for pathModuleId in pathSoFar {
-          if !modulesRequiringRebuild.contains(pathModuleId) &&
-             !isMainModule {
-            reporter.reportExplicitDependencyWillBeReBuilt(pathModuleId.moduleNameForDiagnostic,
-                                                           reason: "Invalidated by downstream dependency")
-          }
-        }
-      }
-      modulesRequiringRebuild.formUnion(pathSoFar)
-    }
-
-    // Routine to invalidate this node and the path that led to it
-    let invalidateOutOfDate = { (modulesRequiringRebuild: inout Set<ModuleDependencyId>) in
-      reporter?.reportExplicitDependencyWillBeReBuilt(moduleId.moduleNameForDiagnostic, reason: "Out-of-date")
-      modulesRequiringRebuild.insert(moduleId)
-      invalidatePath(&modulesRequiringRebuild)
-    }
-
     // Visit the module's dependencies
+    var hasOutOfDateModuleDependency = false
     for dependencyId in moduleInfo.directDependencies ?? [] {
-      if !visitedValidated.contains(dependencyId) {
-        let newPath = pathSoFar + [moduleId]
-        try outOfDateModuleScan(on: moduleDependencyGraph, from: dependencyId, pathSoFar: newPath,
-                                visitedValidated: &visitedValidated,
+      // If we have not already visited this module, recurse.
+      if !visited.contains(dependencyId) {
+        try outOfDateModuleScan(on: moduleDependencyGraph, from: dependencyId,
+                                visited: &visited,
                                 modulesRequiringRebuild: &modulesRequiringRebuild)
       }
+      // Even if we're not revisiting a dependency, we must check if it's already known to be out of date.
+      hasOutOfDateModuleDependency = hasOutOfDateModuleDependency || modulesRequiringRebuild.contains(dependencyId)
     }
 
-    if modulesRequiringRebuild.contains(moduleId) {
-      invalidatePath(&modulesRequiringRebuild)
+    if hasOutOfDateModuleDependency {
+        reporter?.reportExplicitDependencyWillBeReBuilt(moduleId.moduleNameForDiagnostic, reason: "Invalidated by downstream dependency")
+        modulesRequiringRebuild.insert(moduleId)
     } else if try !IncrementalCompilationState.IncrementalDependencyAndInputSetup.verifyModuleDependencyUpToDate(moduleID: moduleId, moduleInfo: moduleInfo,
                                                                                                                  fileSystem: fileSystem, reporter: reporter) {
-      invalidateOutOfDate(&modulesRequiringRebuild)
-    } else {
-      // Only if this module is known to be up-to-date with respect to its inputs
-      // and dependencies do we mark it as visited. We may need to re-visit
-      // out-of-date modules in order to collect all possible paths to them.
-      visitedValidated.insert(moduleId)
+        reporter?.reportExplicitDependencyWillBeReBuilt(moduleId.moduleNameForDiagnostic, reason: "Out-of-date")
+        modulesRequiringRebuild.insert(moduleId)
     }
+
+    // Now that we've determined if this module must be rebuilt, mark it as visited.
+    visited.insert(moduleId)
   }
 
   /// In an explicit module build, filter out dependency module pre-compilation tasks