Improve performance of module lookups

Sources/Knit/Module/DependencyBuilder.swift

@@ -40,35 +40,42 @@ final class DependencyBuilder {
 
         let toAssemble = assemblyCache.toAssemble
 
+        var inputLookup: [AssemblyReference: any ModuleAssembly] = [:]
+        for module in modules {
+            let ref = AssemblyReference(type(of: module))
+            inputLookup[ref] = module
+            for replaced in ref.type.replaces {
+                inputLookup[AssemblyReference(replaced)] = module
+            }
+        }
+
         // Instantiate all types
         for ref in toAssemble {
-            guard !self.isRegisteredInParent(ref) else {
-                continue
-            }
-            assemblies.append(try instantiate(moduleType: ref.type))
+            assemblies.append(try instantiate(moduleRef: ref, inputLookup: inputLookup))
         }
     }
 
-    private func instantiate(moduleType: any ModuleAssembly.Type) throws -> any ModuleAssembly {
-        let inputModule = inputModules.first(where: { type(of: $0) == moduleType})
-        let existingType = inputModules.first { assembly in
-            return type(of: assembly).doesReplace(type: moduleType)
-        }
-        if let existingType {
-            return existingType
+    private func instantiate(
+        moduleRef: AssemblyReference,
+        inputLookup: [AssemblyReference: any ModuleAssembly]
+    ) throws -> any ModuleAssembly {
+        if let inputModule = inputLookup[moduleRef] {
+            return inputModule
         }
-        if let overrideType = try defaultOverride(moduleType, fromInput: inputModule != nil),
+
+        if let overrideType = try defaultOverride(moduleRef.type, fromInput: false),
            let autoInit = overrideType as? any AutoInitModuleAssembly.Type {
             return autoInit.init()
         }
-        if let inputModule {
-            return inputModule
-        }
-        if let autoInit = moduleType as? any AutoInitModuleAssembly.Type {
+
+        if let autoInit = moduleRef.type as? any AutoInitModuleAssembly.Type {
             return autoInit.init()
         }
 
-        throw DependencyBuilderError.moduleNotProvided(moduleType, dependencyTree.sourcePathString(moduleType: moduleType))
+        throw DependencyBuilderError.moduleNotProvided(
+            moduleRef.type,
+            dependencyTree.sourcePathString(moduleType: moduleRef.type)
+        )
     }
 
     private func gatherDependencies(

Sources/Knit/Module/DependencyTree.swift

@@ -8,7 +8,7 @@ import Foundation
 public struct DependencyTree: CustomDebugStringConvertible {
 
     // Assemblies that were originally provided to create the tree
-    private let inputModules: [AssemblyReference]
+    private let inputModules: Set<AssemblyReference>
 
     // List of all registered modules
     private var allModules: Set<AssemblyReference>
@@ -17,8 +17,8 @@ public struct DependencyTree: CustomDebugStringConvertible {
     private var moduleSources: [AssemblyReference: AssemblyReference] = [:]
 
     init(inputModules: [any ModuleAssembly]) {
-        self.inputModules = inputModules.map { AssemblyReference(type(of: $0)) }
-        allModules = Set(self.inputModules)
+        self.inputModules = Set(inputModules.map { AssemblyReference(type(of: $0)) })
+        allModules = self.inputModules
     }
 
     mutating func add(
@@ -48,7 +48,10 @@ public struct DependencyTree: CustomDebugStringConvertible {
     }
 
     public var debugDescription: String {
-        return inputModules.flatMap { debugDescription(assemblyRef: $0, indent: "") }.joined(separator: "\n")
+        let sortedInputs = inputModules
+            .map { ($0, $0.debugDescription) }
+            .sorted(by: {$0.1 < $1.1})
+        return sortedInputs.flatMap { debugDescription(assemblyRef: $0.0, indent: "") }.joined(separator: "\n")
     }
 
     public func dumpGraph() {

Tests/KnitTests/DependencyBuilderTests.swift

@@ -18,9 +18,9 @@ final class DependencyBuilderTests: XCTestCase {
         XCTAssertEqual(
             builder.dependencyTree.debugDescription,
             """
-            Assembly3
             Assembly1
               - Assembly2
+            Assembly3
             """
         )
 
@@ -71,9 +71,9 @@ final class DependencyBuilderTests: XCTestCase {
         XCTAssertEqual(
             builder.dependencyTree.debugDescription,
             """
+            Assembly1
             Assembly4
               - Assembly2
-            Assembly1
             """
         )
     }