Fix protocol DefaultImpl with associated types + generic functions

Generator/Sources/Internal/Crawlers/Crawler.swift

@@ -76,7 +76,8 @@ final class Crawler: SyntaxVisitor {
             attributes: attributes(from: node.attributes),
             accessibility: accessibility(from: node.modifiers) ?? (container as? HasAccessibility)?.accessibility ?? .internal,
             name: node.name.filteredDescription,
-            genericParameters: (genericParameters(from: node.primaryAssociatedTypeClause?.primaryAssociatedTypes) + associatedTypes(from: node.memberBlock.members)).merged(),
+            associatedTypes: associatedTypes(from: node.memberBlock.members),
+            primaryAssociatedTypes: genericParameters(from: node.primaryAssociatedTypeClause?.primaryAssociatedTypes),
             genericRequirements: genericRequirements(from: node.genericWhereClause?.requirements),
             inheritedTypes: inheritedTypes,
             members: []

Generator/Sources/Internal/GeneratorHelper.swift

@@ -40,7 +40,7 @@
 
         let matchableWhereConstraints = method.signature.parameters.enumerated().map { index, parameter -> String in
             let type = parameter.type.isOptional ? "OptionalMatchedType" : "MatchedType"
-            return "M\(index + 1).\(type) == \(genericSafeType(from: parameter.type.withoutAttributes(except: ["@Sendable"]).unoptionaled.description))"
+            return "M\(index + 1).\(type) == \(genericSafeType(from: parameter.type.withoutAttributes(except: ["@MainActor", "@Sendable"]).unoptionaled.description))"
         }
         let methodWhereConstraints = method.signature.whereConstraints
         return " where \((matchableWhereConstraints + methodWhereConstraints).joined(separator: ", "))"
@@ -57,7 +57,7 @@
     private static func parameterMatchers(for parameters: [MethodParameter]) -> String {
         guard parameters.isEmpty == false else { return "let matchers: [Cuckoo.ParameterMatcher<Void>] = []" }
 
-        let tupleType = parameters.map { $0.type.withoutAttributes(except: ["@Sendable"]).description }.joined(separator: ", ")
+        let tupleType = parameters.map { $0.type.withoutAttributes(except: ["@MainActor", "@Sendable"]).description }.joined(separator: ", ")
         let matchers = parameters
             // Enumeration is done after filtering out parameters without usable names.
             .enumerated()
@@ -76,7 +76,7 @@
     private static func openNestedClosure(for method: Method) -> String {
         var fullString = ""
         for (index, parameter) in method.signature.parameters.enumerated() {
             if !parameter.type.containsAttribute(named: "@escaping"), let closure = parameter.type.findClosure() {
                 if fullString.isEmpty {
                     fullString = "\n"
                 }

Generator/Sources/Internal/Templates/MockTemplate.swift

@@ -15,11 +15,16 @@ extension Templates {
 {% if container.hasParent %}
 extension {{ container.parentFullyQualifiedName }} {
 {% endif %}
+{% if container.hasPrimaryAssociatedTypes %}
+@available(iOS 16.0.0, macOS 13.0.0, watchOS 9.0, tvOS 16, *) // runtime support for constrained protocols with primary associated types
+{% endif %}
 {{ container.accessibility|withSpace }}class {{ container.mockName }}{{ container.genericParameters }}:{% if container.isNSObjectProtocol %} NSObject,{% endif %} {{ container.name }}{% if container.isImplementation %}{{ container.genericArguments }}{% endif %},{% if container.isImplementation %} Cuckoo.ClassMock{% else %} Cuckoo.ProtocolMock{% endif %}, @unchecked Sendable {
-    {% if container.isGeneric and not container.isImplementation %}
+    {% if container.isGeneric and not container.isImplementation and not container.hasOnlyPrimaryAssociatedTypes %}
     {{ container.accessibility|withSpace }}typealias MocksType = \(typeErasureClassName)
-    {% else %}
+    {% elif container.isImplementation %}
     {{ container.accessibility|withSpace }}typealias MocksType = {{ container.name }}{{ container.genericArguments }}
+    {% else %}
+    {{ container.accessibility|withSpace }}typealias MocksType = any {{ container.name }}{{ container.genericArguments }}
     {% endif %}
     {{ container.accessibility|withSpace }}typealias Stubbing = __StubbingProxy_{{ container.name }}
     {{ container.accessibility|withSpace }}typealias Verification = __VerificationProxy_{{ container.name }}
@@ -31,7 +36,7 @@ extension {{ container.parentFullyQualifiedName }} {
 
     {{ container.accessibility|withSpace }}let cuckoo_manager = Cuckoo.MockManager.preconfiguredManager ?? Cuckoo.MockManager(hasParent: {{ container.isImplementation }})
 
-    {% if container.isGeneric and not container.isImplementation %}
+    {% if container.isGeneric and not container.isImplementation and not container.hasOnlyPrimaryAssociatedTypes %}
 \(Templates.typeErasure.indented())
 
     private var __defaultImplStub: \(typeErasureClassName)?
@@ -43,7 +48,7 @@ extension {{ container.parentFullyQualifiedName }} {
     }
 
     {{ container.accessibility|withSpace }}func enableDefaultImplementation<\(staticGenericParameter): {{ container.name }}>(mutating stub: UnsafeMutablePointer<\(staticGenericParameter)>) where {{ container.genericProtocolIdentity }} {
-        __defaultImplStub = \(typeErasureClassName)(from: stub, keeping: nil)
+        __defaultImplStub = \(typeErasureClassName)(from: stub, keeping: stub.pointee)
         cuckoo_manager.enableDefaultStubImplementation()
     }
     {% else %}

Generator/Sources/Internal/Templates/TypeErasureTemplate.swift

@@ -3,7 +3,7 @@ import Foundation
 extension Templates {
     static let typeErasure = """
     {{ container.accessibility|withSpace }}class \(typeErasureClassName): {{ container.name }}, @unchecked Sendable {
-        private let reference: Any
+        private let reference: () -> any {{ container.name }}{{ container.genericPrimaryAssociatedTypeArguments }}
 
         {% for property in container.properties %}
         private let _getter_storage$${{ property.name }}: () -> {{ property.type }}
@@ -18,9 +18,9 @@ extension Templates {
         }
 
         {% endfor %}
-        {# For developers: The `keeping reference: Any?` is necessary because when called from the `enableDefaultImplementation(stub:)` method
+        {# For developers: The `keeping reference: Any` is necessary because when called from the `enableDefaultImplementation(stub:)` method
         instead of `enableDefaultImplementation(mutating:)`, we need to prevent the struct getting deallocated. #}
-        init<\(staticGenericParameter): {{ container.name }}>(from defaultImpl: UnsafeMutablePointer<\(staticGenericParameter)>, keeping reference: @escaping @autoclosure () -> Any?) where {{ container.genericProtocolIdentity }} {
+        init<\(staticGenericParameter): {{ container.name }}>(from defaultImpl: UnsafeMutablePointer<\(staticGenericParameter)>, keeping reference: @escaping @autoclosure () -> \(staticGenericParameter)) where {{ container.genericProtocolIdentity }} {
             self.reference = reference
 
             {% for property in container.properties %}
@@ -30,7 +30,9 @@ extension Templates {
             {% endif %}
             {% endfor %}
             {% for method in container.methods %}
+            {% if not method.hasGenericParams %}
             _storage${{ forloop.counter }}${{ method.name }} = defaultImpl.pointee.{{ method.name }}
+            {% endif %}
             {% endfor %}
         }
         {% if container.initializers %}
@@ -43,9 +45,20 @@ extension Templates {
         {% endfor %}
 
         {% for method in container.methods +%}
+        {% if not method.hasGenericParams %}
         private let _storage${{ forloop.counter }}${{ method.name }}: ({{ method.inputTypes }}) {% if method.isAsync %} async{% endif %} {% if method.isThrowing %} throws{% endif %} -> {{ method.returnType }}
+        {% endif %}
         {{ container.accessibility|withSpace }}func {{ method.name|escapeReservedKeywords }}{{ method.signature }} {
+            {% if method.hasGenericParams and method.hasNonPrimaryAssociatedTypeParams %}
+            func openExistential<\(staticGenericParameter): {{ container.name }}{{ container.genericPrimaryAssociatedTypeArguments }}>(_ opened: \(staticGenericParameter)) {% if method.isAsync %} async{% endif %} {% if method.isThrowing %} throws{% endif %} -> {{ method.returnType }} {
+                return {% if method.isThrowing %} try{% endif %} {% if method.isAsync %} await{% endif %} opened.{{ method.name }}{{ method.staticGenericCall }}
+            }
+            return {% if method.isThrowing %} try{% endif %} {% if method.isAsync %} await{% endif %} openExistential(reference())
+            {% elif method.hasGenericParams %}
+            return {% if method.isThrowing %} try{% endif %} {% if method.isAsync %} await{% endif %} reference().{{ method.name }}({{ method.call }})
+            {% else %}
             return {% if method.isThrowing %} try{% endif %} {% if method.isAsync %} await{% endif %} _storage${{ forloop.counter }}${{ method.name }}({{ method.parameterNames }})
+            {% endif %}
         }
         {% endfor %}
     }

Generator/Sources/Internal/Tokens/Capabilities/HasGenerics.swift

@@ -19,15 +19,37 @@
     var isGeneric: Bool {
         !genericParameters.isEmpty
     }
+
+    var hasPrimaryAssociatedTypes: Bool {
+        guard let protocolDeclaration = asProtocol else { return false }
+        return !protocolDeclaration.primaryAssociatedTypes.isEmpty
+    }
+
+    var hasOnlyPrimaryAssociatedTypes: Bool {
+        guard let protocolDeclaration = asProtocol else { return false }
+        return protocolDeclaration.primaryAssociatedTypes.count == protocolDeclaration.genericParameters.count
+    }
 
     func genericsSerialize() -> GeneratorContext {
-        let genericProtocolIdentity = isProtocol ? genericParameters.map { "\(Templates.staticGenericParameter).\($0.name) == \($0.name)" }.joined(separator: ", ") : nil
-
+        var genericProtocolIdentity: String?
+        var genericPrimaryAssociatedTypeArguments: String?
+
+        if let protocolDeclaration = asProtocol {
+            genericProtocolIdentity = genericParameters.map { "\(Templates.staticGenericParameter).\($0.name) == \($0.name)" }.joined(separator: ", ")
+            if !protocolDeclaration.primaryAssociatedTypes.isEmpty {
+                let arguments = protocolDeclaration.primaryAssociatedTypes.map { $0.name }.joined(separator: ", ")
+                genericPrimaryAssociatedTypeArguments = "<\(arguments)>"
+            }
+        }
+
         return [
             "isGeneric": isGeneric,
             "genericParameters": genericParametersString,
             "genericArguments": genericArgumentsString,
+            "hasPrimaryAssociatedTypes": hasPrimaryAssociatedTypes,
+            "hasOnlyPrimaryAssociatedTypes": hasOnlyPrimaryAssociatedTypes,
             "genericProtocolIdentity": genericProtocolIdentity,
+            "genericPrimaryAssociatedTypeArguments": genericPrimaryAssociatedTypeArguments,
         ]
         .compactMapValues { $0 }
     }

Generator/Sources/Internal/Tokens/ComplexType.swift

@@ -20,7 +20,7 @@
             self = .attributed(
                 attributes: [
                     attributedType.attributes.map { $0.trimmedDescription },
                     attributedType.specifier.map { [$0.trimmedDescription] } ?? [],
                 ].flatMap { $0 },
                 baseType: ComplexType(syntax: attributedType.baseType)
             )
@@ -241,6 +241,83 @@
             nil
         }
     }
+
+    func replaceType(named typeName: String, with replacement: String) -> ComplexType? {
+        switch self {
+        case .attributed(attributes: let attributes, baseType: let baseType):
+            return baseType.replaceType(named: typeName, with: replacement)
+                .map { ComplexType.attributed(attributes: attributes, baseType: $0) }
+        case .optional(wrappedType: let wrappedType, isImplicit: let isImplicit):
+            return wrappedType.replaceType(named: typeName, with: replacement)
+                .map { ComplexType.optional(wrappedType: $0, isImplicit: isImplicit) }
+        case .array(elementType: let elementType):
+            return elementType.replaceType(named: typeName, with: replacement)
+                .map { ComplexType.array(elementType: $0) }
+        case .dictionary(keyType: let keyType, valueType: let valueType):
+            let newKey = keyType.replaceType(named: typeName, with: replacement)
+            let newValue = valueType.replaceType(named: typeName, with: replacement)
+            if newKey == nil && newValue == nil { return nil }
+            return .dictionary(
+                keyType: newKey ?? keyType,
+                valueType: newValue ?? valueType
+            )
+        case .closure(let closure):
+            var changed = false
+            let newParams: [ComplexType.Closure.Parameter] = closure.parameters.map { param in
+                if let newType = param.type.replaceType(named: typeName, with: replacement) {
+                    changed = true
+                    return ComplexType.Closure.Parameter(label: param.label, type: newType)
+                } else {
+                    return param
+                }
+            }
+            let newReturn = closure.returnType.replaceType(named: typeName, with: replacement)
+            if !changed && newReturn == nil { return nil }
+            return .closure(.init(parameters: newParams, effects: closure.effects, returnType: newReturn ?? closure.returnType))
+        case .type(let name):
+            return name == typeName ? ComplexType.type(replacement) : nil
+        }
+    }
+
+    func replaceTypes(named typeNames: [String], with replacement: (String) -> String) -> ComplexType? {
+        var changed = false
+        var type = self
+        for typeName in typeNames {
+            if let replaced = type.replaceType(named: typeName, with: replacement(typeName)) {
+                changed = true
+                type = replaced
+            }
+        }
+
+        return changed ? type : nil
+    }
+
+    func containsType(named typeName: String) -> Bool {
+        switch self {
+        case .attributed(attributes: _, baseType: let baseType):
+            baseType.containsType(named: typeName)
+        case .optional(wrappedType: let wrappedType, isImplicit: _):
+            wrappedType.containsType(named: typeName)
+        case .array(elementType: let elementType):
+            elementType.containsType(named: typeName)
+        case .dictionary(keyType: let keyType, valueType: let valueType):
+            keyType.containsType(named: typeName) || valueType.containsType(named: typeName)
+        case .closure(let closure):
+            (closure.parameters.map(\.type) + [closure.returnType]).contains(where: { $0.containsType(named: typeName)})
+        case .type(let name):
+            name == typeName
+        }
+    }
+
+    func containsTypes(named typeNames: [String]) -> Bool {
+        typeNames.contains(where: { containsType(named: $0) })
+    }
+}
+
+extension String {
+    func forceCast(as type: ComplexType) -> String {
+        "\(self) as! \(type.withoutAttributes(except: ["@MainActor", "@Sendable"]).description)"
+    }
 }
 
 extension ComplexType.Closure.Effects {

Generator/Sources/Internal/Tokens/Method.swift

@@ -55,24 +55,41 @@ extension Method {
     var hasOptionalParams: Bool {
         signature.parameters.contains { $0.type.isOptional }
     }
+
+    var hasGenericParams: Bool {
+        !signature.genericParameters.isEmpty
+    }
+
+    var hasNonPrimaryAssociatedTypeParams: Bool {
+        guard let parent = parent?.asProtocol else { return false }
+        return signature.containsTypes(named: parent.nonPrimaryAssociatedTypes.map(\.name))
+    }
 
     func serialize() -> [String : Any] {
-        let call = signature.parameters
-            .map { parameter in
-                let name = escapeReservedKeywords(for: parameter.usableName)
-                let value = "\(parameter.isInout ? "&" : "")\(name)\(parameter.type.containsAttribute(named: "@autoclosure") ? "()" : "")"
-                if parameter.name == "_" {
-                    return value
-                } else {
-                    return "\(parameter.name): \(value)"
-                }
-            }
-            .joined(separator: ", ")
-
         guard let parent else {
             fatalError("Failed to find parent of method \(fullSignature). Please file a bug.")
         }
-
+
+        let call = signature.parameters
+            .map(\.call)
+            .joined(separator: ", ")
+
+        var staticGenericCall = "(\(call))"
+
+        if let parent = parent.asProtocol, !parent.nonPrimaryAssociatedTypes.isEmpty {
+            let nonPrimary = parent.nonPrimaryAssociatedTypes.map(\.name)
+
+            let staticGenericCallableParameters = signature.parameters
+                .map { $0.callAndCastTypes(named: nonPrimary, as: { Templates.staticGenericParameter + ".\($0)" }) }
+                .joined(separator: ", ")
+
+            staticGenericCall = "(\(staticGenericCallableParameters))"
+
+            if let returnType, returnType.containsTypes(named: nonPrimary) {
+                staticGenericCall = staticGenericCall.forceCast(as: returnType)
+            }
+        }
+
         let stubFunctionPrefix = parent.isClass ? "Class" : "Protocol"
         let returnString = returnType?.isVoid == false ? "" : "NoReturn"
         let throwingString = isThrowing ? "Throwing" : ""
@@ -108,16 +125,19 @@ extension Method {
             "throwTypeError": signature.throwType?.type ?? "",
             "fullyQualifiedName": fullyQualifiedName,
             "call": call,
+            "staticGenericCall": staticGenericCall,
             "parameterSignature": signature.parameters.map { $0.description }.joined(separator: ", "),
             "parameterSignatureWithoutNames": signature.parameters.map { "\($0.name): \($0.type)" }.joined(separator: ", "),
             "argumentSignature": signature.parameters.map { $0.type.description }.joined(separator: ", "),
             "stubFunction": stubFunction,
-            "inputTypes": signature.parameters.map { $0.type.withoutAttributes(except: ["@escaping", "@Sendable"]).description }.joined(separator: ", "),
-            "genericInputTypes": signature.parameters.map { $0.type.withoutAttributes(except: ["@Sendable"]).description }.joined(separator: ", "),
+            "inputTypes": signature.parameters.map { $0.type.withoutAttributes(except: ["@escaping", "@MainActor", "@Sendable"]).description }.joined(separator: ", "),
+            "genericInputTypes": signature.parameters.map { $0.type.withoutAttributes(except: ["@MainActor", "@Sendable"]).description }.joined(separator: ", "),
             "isOptional": isOptional,
             "hasClosureParams": hasClosureParams,
             "hasOptionalParams": hasOptionalParams,
+            "hasGenericParams": hasGenericParams,
             "genericParameters": signature.genericParameters.sourceDescription,
+            "hasNonPrimaryAssociatedTypeParams": hasNonPrimaryAssociatedTypeParams,
             "hasUnavailablePlatforms": hasUnavailablePlatforms,
             "unavailablePlatformsCheck": unavailablePlatformsCheck,
         ]

Generator/Sources/Internal/Tokens/MethodParameter.swift

@@ -25,6 +25,27 @@ struct MethodParameter: Token {
     var isEscaping: Bool {
         type.isClosure && (type.containsAttribute(named: "@escaping") || type.isOptional)
     }
+
+    var call: String {
+        let escapedName = escapeReservedKeywords(for: usableName)
+        let value = "\(isInout ? "&" : "")\(escapedName)\(type.containsAttribute(named: "@autoclosure") ? "()" : "")"
+        if name == "_" {
+            return value
+        } else {
+            return "\(name): \(value)"
+        }
+    }
+
+    func callAndCastTypes(named typeNames: [String], as replacement: (String) -> String) -> String {
+        let replaced = type.replaceTypes(named: typeNames, with: replacement)
+
+        let callToCast = call
+        if let replaced {
+            return callToCast.forceCast(as: replaced)
+        } else {
+            return callToCast
+        }
+    }
 
     func serialize() -> [String: Any] {
         return [

Generator/Sources/Internal/Tokens/MethodSignature.swift

@@ -64,3 +64,14 @@ extension Method.Signature {
         && whereConstraints == other.whereConstraints
     }
 }
+
+extension Method.Signature {
+    func containsType(named typeName: String) -> Bool {
+        parameters.map(\.type)
+            .contains(where: { $0.containsType(named: typeName) })
+    }
+
+    func containsTypes(named typeNames: [String]) -> Bool {
+        typeNames.contains(where: { containsType(named: $0) })
+    }
+}