From 49210f45a62fe981b805cb7c86e6fd37def1dd3b Mon Sep 17 00:00:00 2001
From: Ian Leitch <ian@leitch.io>
Date: Mon, 18 Dec 2023 11:51:46 +0000
Subject: [PATCH] Retain @dynamicMemberLookup subscript functions. Closes #684

---
 CHANGELOG.md                                   |  2 +-
 .../DynamicMemberLookupReferenceBuilder.swift  | 18 ++++++++++++++++++
 .../SourceGraph/SourceGraphMutatorRunner.swift |  1 +
 ...stRetainsDynamicMemberLookupSubscript.swift | 10 ++++++++++
 Tests/PeripheryTests/RetentionTest.swift       |  9 +++++++++
 Tests/Shared/DeclarationDescription.swift      |  4 ++++
 6 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 Sources/PeripheryKit/SourceGraph/Mutators/DynamicMemberLookupReferenceBuilder.swift
 create mode 100644 Tests/Fixtures/RetentionFixtures/testRetainsDynamicMemberLookupSubscript.swift

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bb9206e2a..8377bcc2a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,7 +10,7 @@
 
 ##### Bug Fixes
 
-- None.
+- Subscript functions required by `@dynamicMemberLookup` are now retained.
 
 ## 2.17.1 (2023-12-04)
 
diff --git a/Sources/PeripheryKit/SourceGraph/Mutators/DynamicMemberLookupReferenceBuilder.swift b/Sources/PeripheryKit/SourceGraph/Mutators/DynamicMemberLookupReferenceBuilder.swift
new file mode 100644
index 000000000..56807d264
--- /dev/null
+++ b/Sources/PeripheryKit/SourceGraph/Mutators/DynamicMemberLookupReferenceBuilder.swift
@@ -0,0 +1,18 @@
+import Foundation
+import Shared
+
+final class DynamicMemberLookupReferenceBuilder: SourceGraphMutator {
+    private let graph: SourceGraph
+
+    required init(graph: SourceGraph, configuration: Configuration) {
+        self.graph = graph
+    }
+
+    func mutate() throws {
+        for decl in graph.declarations(ofKind: .functionSubscript) {
+            if decl.name == "subscript(dynamicMember:)", decl.parent?.attributes.contains("dynamicMemberLookup") ?? false {
+                graph.markRetained(decl)
+            }
+        }
+    }
+}
diff --git a/Sources/PeripheryKit/SourceGraph/SourceGraphMutatorRunner.swift b/Sources/PeripheryKit/SourceGraph/SourceGraphMutatorRunner.swift
index 92dce5730..91752a223 100644
--- a/Sources/PeripheryKit/SourceGraph/SourceGraphMutatorRunner.swift
+++ b/Sources/PeripheryKit/SourceGraph/SourceGraphMutatorRunner.swift
@@ -30,6 +30,7 @@ public final class SourceGraphMutatorRunner {
         DefaultConstructorReferenceBuilder.self,
         ComplexPropertyAccessorReferenceBuilder.self,
         EnumCaseReferenceBuilder.self,
+        DynamicMemberLookupReferenceBuilder.self,
 
         UnusedParameterRetainer.self,
         AssetReferenceRetainer.self,
diff --git a/Tests/Fixtures/RetentionFixtures/testRetainsDynamicMemberLookupSubscript.swift b/Tests/Fixtures/RetentionFixtures/testRetainsDynamicMemberLookupSubscript.swift
new file mode 100644
index 000000000..7705fc277
--- /dev/null
+++ b/Tests/Fixtures/RetentionFixtures/testRetainsDynamicMemberLookupSubscript.swift
@@ -0,0 +1,10 @@
+@dynamicMemberLookup
+public struct FixtureStruct7 {
+    subscript(dynamicMember member: String) -> String {
+        ""
+    }
+
+    subscript(other: String) -> String {
+        ""
+    }
+}
diff --git a/Tests/PeripheryTests/RetentionTest.swift b/Tests/PeripheryTests/RetentionTest.swift
index 79340bdfa..59447c3bd 100644
--- a/Tests/PeripheryTests/RetentionTest.swift
+++ b/Tests/PeripheryTests/RetentionTest.swift
@@ -1032,6 +1032,15 @@ final class RetentionTest: FixtureSourceGraphTestCase {
         }
     }
 
+    func testRetainsDynamicMemberLookupSubscript() {
+        analyze(retainPublic: true) {
+            assertReferenced(.struct("FixtureStruct7")) {
+                self.assertReferenced(.functionSubscript("subscript(dynamicMember:)"))
+                self.assertNotReferenced(.functionSubscript("subscript(_:)"))
+            }
+        }
+    }
+
     // MARK: - Assign-only properties
 
     func testSimplePropertyAssignedButNeverRead() {
diff --git a/Tests/Shared/DeclarationDescription.swift b/Tests/Shared/DeclarationDescription.swift
index 6424c8d7f..dfed294b1 100644
--- a/Tests/Shared/DeclarationDescription.swift
+++ b/Tests/Shared/DeclarationDescription.swift
@@ -70,6 +70,10 @@ struct DeclarationDescription: CustomStringConvertible {
         self.init(kind: .functionDestructor, name: name, line: line)
     }
 
+    static func functionSubscript(_ name: String, line: Int? = nil) -> Self {
+        self.init(kind: .functionSubscript, name: name, line: line)
+    }
+
     static func varStatic(_ name: String, line: Int? = nil) -> Self {
         self.init(kind: .varStatic, name: name, line: line)
     }