From 8b1af46bdbd05a09be9a71374ff76fd04ea0be45 Mon Sep 17 00:00:00 2001
From: Ian Leitch <ian@leitch.io>
Date: Sun, 22 Sep 2024 09:52:37 +0100
Subject: [PATCH] Workaround Swift 6 bug for missing related references to
 external overriden declarations. Closes #820

---
 CHANGELOG.md                                  |  1 +
 .../SourceGraph/Elements/Declaration.swift    |  4 ++
 .../Mutators/ExternalOverrideRetainer.swift   | 46 +++++++++++++++++++
 .../SourceGraphMutatorRunner.swift            |  1 +
 4 files changed, 52 insertions(+)
 create mode 100644 Sources/SourceGraph/Mutators/ExternalOverrideRetainer.swift

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8e8d92c7e..05e1a23c9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@
 
 - Enums with the `@main` attribute are now retained.
 - Unused public/exported imports are excluded from the results even if unused in the declaring file as the exported symbols may be referenced in other files, and thus removing the import would result in a build failure.
+- Fixed issue in Swift 6 where declarations that override external type members are incorrectly identified as unused.
 
 ## 2.21.0 (2024-06-15)
 
diff --git a/Sources/SourceGraph/Elements/Declaration.swift b/Sources/SourceGraph/Elements/Declaration.swift
index 7aa5c17b0..7b1a41bc3 100644
--- a/Sources/SourceGraph/Elements/Declaration.swift
+++ b/Sources/SourceGraph/Elements/Declaration.swift
@@ -164,6 +164,10 @@ public final class Declaration {
             functionKinds.union(variableKinds).union(globalKinds)
         }
 
+        static var overrideKinds: Set<Kind> {
+            [.functionMethodInstance, .varInstance]
+        }
+
         public var isAccessorKind: Bool {
             rawValue.hasPrefix("function.accessor")
         }
diff --git a/Sources/SourceGraph/Mutators/ExternalOverrideRetainer.swift b/Sources/SourceGraph/Mutators/ExternalOverrideRetainer.swift
new file mode 100644
index 000000000..a515244a6
--- /dev/null
+++ b/Sources/SourceGraph/Mutators/ExternalOverrideRetainer.swift
@@ -0,0 +1,46 @@
+import Configuration
+import Foundation
+import Shared
+
+/// Retains instance functions/vars that override external declarations.
+///
+/// It's not possible to determine if a declaration that overrides an external declaration is used, as the
+/// external implementation may call the overridden declaration.
+final class ExternalOverrideRetainer: SourceGraphMutator {
+    private let graph: SourceGraph
+    private let isSwift6FixEnabled: Bool
+
+    required init(graph: SourceGraph, configuration _: Configuration, swiftVersion: SwiftVersion) {
+        self.graph = graph
+        isSwift6FixEnabled = swiftVersion.version.isVersion(greaterThanOrEqualTo: "6.0")
+    }
+
+    func mutate() {
+        for decl in graph.declarations(ofKinds: Declaration.Kind.overrideKinds) {
+            guard decl.isOverride else { continue }
+
+            var didIdentifyRelatedRef = false
+
+            for relatedRef in decl.related {
+                if relatedRef.kind == decl.kind,
+                   relatedRef.name == decl.name,
+                   relatedRef.location == decl.location
+                {
+                    didIdentifyRelatedRef = true
+
+                    if graph.explicitDeclaration(withUsr: relatedRef.usr) == nil {
+                        // The related decl is external.
+                        graph.markRetained(decl)
+                    }
+
+                    break
+                }
+            }
+
+            // https://github.com/swiftlang/swift/issues/76628
+            if !didIdentifyRelatedRef, isSwift6FixEnabled {
+                graph.markRetained(decl)
+            }
+        }
+    }
+}
diff --git a/Sources/SourceGraph/SourceGraphMutatorRunner.swift b/Sources/SourceGraph/SourceGraphMutatorRunner.swift
index 71b9025ad..e9090c65f 100644
--- a/Sources/SourceGraph/SourceGraphMutatorRunner.swift
+++ b/Sources/SourceGraph/SourceGraphMutatorRunner.swift
@@ -42,6 +42,7 @@ public final class SourceGraphMutatorRunner {
         ResultBuilderRetainer.self,
         CapitalSelfFunctionCallRetainer.self,
         CodablePropertyRetainer.self,
+        ExternalOverrideRetainer.self,
 
         AncestralReferenceEliminator.self,
         AssignOnlyPropertyReferenceEliminator.self,