Merge pull request #101 from squareup/skorulis/source-loop

Prevent infinite loops when generating source paths
cashapp · Dec 12, 2023 · 62bfe7d · 62bfe7d
2 parents ace7453 + 118a6d8
commit 62bfe7d
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 8 deletions.
diff --git a/Sources/KnitLib/Module/DependencyBuilder.swift b/Sources/KnitLib/Module/DependencyBuilder.swift
@@ -81,7 +81,12 @@ final class DependencyBuilder {
         result: inout [any ModuleAssembly.Type],
         source: (any ModuleAssembly.Type)?
     ) throws {
-        moduleSources[String(describing: from)] = source
+        // Add a source for the original type
+        let originalFromInput = inputModules.contains(where: { type(of: $0).matches(moduleType: from) })
+        if !originalFromInput {
+            moduleSources[String(describing: from)] = source
+        }
+
         let resolved = try resolvedType(from)
 
         // Assembly validation should be performed "up front"
@@ -98,8 +103,12 @@ final class DependencyBuilder {
         guard !result.contains(where: {$0 == resolved}) else {
             return
         }
-        // Add a source for both the original and the resolved types
-        moduleSources[String(describing: resolved)] = source
+        // Add a source for the resolved type
+        let resolvedFromInput = inputModules.contains(where: { type(of: $0).matches(moduleType: resolved) })
+        if !resolvedFromInput {
+            moduleSources[String(describing: resolved)] = source
+        }
+
         result.insert(resolved, at: insertionPoint)
         for dep in getDependencies(resolved) {
             try gatherDependencies(
@@ -139,15 +148,34 @@ final class DependencyBuilder {
         return type
     }
 
-    func sourcePath(moduleType: any ModuleAssembly.Type) -> [String] {
-        return sourcePath(moduleName: String(describing: moduleType))
+    private func buildSourcePath(moduleType: any ModuleAssembly.Type, path: inout [String]) {
+        return buildSourcePath(moduleName: String(describing: moduleType), path: &path)
     }
 
-    func sourcePath(moduleName: String) -> [String] {
+    private func buildSourcePath(moduleName: String, path: inout [String]) {
+        path.insert(moduleName, at: 0)
         guard let source = moduleSources[moduleName] else {
-            return [moduleName]
+            return
+        }
+
+        // Prevent an infinite loop
+        let sourceName = String(describing: source)
+        if path.contains(sourceName) {
+            return
         }
-        return sourcePath(moduleType: source) + [moduleName]
+        return buildSourcePath(moduleType: source, path: &path)
+    }
+
+    func sourcePath(moduleType: any ModuleAssembly.Type) -> [String] {
+        var path: [String] = []
+        buildSourcePath(moduleType: moduleType, path: &path)
+        return path
+    }
+
+    func sourcePath(moduleName: String) -> [String] {
+        var path: [String] = []
+        buildSourcePath(moduleName: moduleName, path: &path)
+        return path
     }
 
     func sourcePathString(moduleName: String) -> String {

diff --git a/Tests/KnitLibTests/ModuleCycleTests.swift b/Tests/KnitLibTests/ModuleCycleTests.swift
@@ -24,6 +24,20 @@ final class ModuleCycleTests: XCTestCase {
             ["\(Assembly1.self)", "\(Assembly3.self)"]
         )
     }
+
+    func test_sourceCycle() {
+        let assembler = ModuleAssembler([Assembly5()])
+        XCTAssertEqual(
+            assembler.builder.sourcePath(moduleType: Assembly5.self),
+            ["\(Assembly5.self)"]
+        )
+
+        XCTAssertEqual(
+            assembler.builder.sourcePath(moduleType: Assembly7.self),
+            ["\(Assembly5.self)", "\(Assembly6.self)", "\(Assembly7.self)"]
+        )
+    }
+
 }
 
 // Assembly1 depends on Assembly2
@@ -52,3 +66,20 @@ private struct Assembly4: AutoInitModuleAssembly {
     static var dependencies: [any ModuleAssembly.Type] { [] }
     func assemble(container: Container) {}
 }
+
+// Assembly 5-6-7 form a dependency circle
+
+private struct Assembly5: AutoInitModuleAssembly {
+    static var dependencies: [any ModuleAssembly.Type] { [Assembly6.self] }
+    func assemble(container: Container) {}
+}
+
+private struct Assembly6: AutoInitModuleAssembly {
+    static var dependencies: [any ModuleAssembly.Type] { [Assembly7.self] }
+    func assemble(container: Container) {}
+}
+
+private struct Assembly7: AutoInitModuleAssembly {
+    static var dependencies: [any ModuleAssembly.Type] { [Assembly5.self] }
+    func assemble(container: Container) {}
+}