Properly parse file when extension is written before @INSTANTIABLE type

dfed · Dec 27, 2024 · cd58644 · cd58644
1 parent ebb3310
commit cd58644
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 3 deletions.
diff --git a/Sources/SafeDICore/Visitors/FileVisitor.swift b/Sources/SafeDICore/Visitors/FileVisitor.swift
@@ -107,6 +107,10 @@ public final class FileVisitor: SyntaxVisitor {
         return .visitChildren
     }
 
+    public override func visitPost(_: ExtensionDeclSyntax) {
+        exitType()
+    }
+
     // MARK: Public
 
     public private(set) var imports = [ImportStatement]()
@@ -142,8 +146,6 @@ public final class FileVisitor: SyntaxVisitor {
     }
 
     private func exitType() {
-        if let parentType {
-            self.parentType = parentType.popNested
-        }
+        parentType = parentType?.popNested
     }
 }
diff --git a/Tests/SafeDIToolTests/SafeDIToolCodeGenerationTests.swift b/Tests/SafeDIToolTests/SafeDIToolCodeGenerationTests.swift
@@ -5314,6 +5314,59 @@ final class SafeDIToolCodeGenerationTests: XCTestCase {
         )
     }
 
+    func test_run_writesConvenienceExtensionOnRootOfTree_whenTypeIsFulfilledByNestedTypeAndExtensionIsDefinedFirst() async throws {
+        let output = try await executeSafeDIToolTest(
+            swiftFileContent: [
+                """
+                extension ErasedType {
+                    // Extension defined before an @Instantiable should not make a difference.
+                }
+
+                public protocol ErasedType {}
+
+                public final class EnclosingType {
+                    @Instantiable
+                    public final class SomeType: ErasedType {
+                        public init() {}
+                    }
+                }
+                """,
+                """
+                extension ErasedType {
+                    // Extension defined before an @Instantiable should not make a difference.
+                }
+
+                @Instantiable
+                public final class TypeWithDependency {
+                    public init(erasedType: ErasedType ) {
+                        fatalError("SafeDI doesn't inspect the initializer body")
+                    }
+
+                    @Instantiated(fulfilledByType: "EnclosingType.SomeType") let erasedType: ErasedType 
+                }
+                """,
+            ],
+            buildDependencyTreeOutput: true,
+            filesToDelete: &filesToDelete
+        )
+
+        XCTAssertEqual(
+            try XCTUnwrap(output.dependencyTree),
+            """
+            // This file was generated by the SafeDIGenerateDependencyTree build tool plugin.
+            // Any modifications made to this file will be overwritten on subsequent builds.
+            // Please refrain from editing this file directly.
+
+            extension TypeWithDependency {
+                public convenience init() {
+                    let erasedType: ErasedType = EnclosingType.SomeType()
+                    self.init(erasedType: erasedType)
+                }
+            }
+            """
+        )
+    }
+
     // MARK: Private
 
     private var filesToDelete = [URL]()