SwiftExtract: separate protocol requirements from default implementations and extension members

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+JavaBindingsPrinting.swift

@@ -173,7 +173,11 @@ extension JNISwift2JavaGenerator {
     let filename = "\(decl.effectiveJavaSimpleName).java"
     logger.debug("Printing contents: \(filename)")
 
-    resetForNewOutputFile(decl.initializers + decl.variables + decl.methods)
+    let methods =
+      decl.swiftNominal.kind == .protocol
+      ? self.supportedProtocolRequirements(of: decl)
+      : decl.initializers + decl.variables + decl.methods
+    resetForNewOutputFile(methods)
 
     printHeader(&printer)
     printPackage(&printer)

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator+SwiftThunkPrinting.swift

@@ -310,9 +310,11 @@ extension JNISwift2JavaGenerator {
     printHeader(&printer)
     printer.println()
 
-    self.currentJavaIdentifiers = JavaIdentifierFactory(
-      type.initializers + type.variables + type.methods
-    )
+    let methods =
+      type.swiftNominal.kind == .protocol
+      ? self.allProtocolRequirementMethods(of: type)
+      : type.initializers + type.variables + type.methods
+    self.currentJavaIdentifiers = JavaIdentifierFactory(methods)
 
     switch type.swiftNominal.kind {
     case .actor, .class, .enum, .struct:

Sources/JExtractSwiftLib/JNI/JNISwift2JavaGenerator.swift

@@ -204,7 +204,8 @@ extension JNISwift2JavaGenerator {
   /// doesn't translate (e.g. referencing `Self`/associated types).
   func supportedProtocolRequirements(of type: ExtractedNominalType) -> [ExtractedFunc] {
     (type.methods + type.variables).filter { requirement in
-      !requirement.isStatic && !requirement.isClass && self.translatedDecl(for: requirement) != nil
+      !requirement.isStatic && !requirement.isClass
+        && (try? self.javaTranslator.translate(requirement)) != nil
     }
   }
 
@@ -218,9 +219,16 @@ extension JNISwift2JavaGenerator {
     var visited: Set<ObjectIdentifier> = []
     var queue: [ExtractedNominalType] = [type]
     var methods: [ExtractedFunc] = []
-    while let current = queue.popLast() {
+    while !queue.isEmpty {
+      let current = queue.removeFirst()
       guard visited.insert(ObjectIdentifier(current)).inserted else { continue }
-      methods.append(contentsOf: self.supportedProtocolRequirements(of: current))
+      for requirement in self.supportedProtocolRequirements(of: current) {
+        if !methods.contains(where: {
+          ExtractedNominalType.satisfiesProtocolRequirement(candidate: requirement, requirement: $0)
+        }) {
+          methods.append(requirement)
+        }
+      }
       queue.append(contentsOf: inheritedProtocols(of: current))
     }
     return methods

Sources/SwiftExtract/AnalysisResult.swift

@@ -45,6 +45,8 @@ public struct AnalysisResult {
     for type in self.extractedTypes.values {
       type.methods = type.methods.flatMap { $0.expandingVariadicOverloads(maxOverloads: maxOverloads) }
       type.initializers = type.initializers.flatMap { $0.expandingVariadicOverloads(maxOverloads: maxOverloads) }
+      type.protocolDefaultImplementations = type.protocolDefaultImplementations.flatMap { $0.expandingVariadicOverloads(maxOverloads: maxOverloads) }
+      type.protocolExtensionMethods = type.protocolExtensionMethods.flatMap { $0.expandingVariadicOverloads(maxOverloads: maxOverloads) }
     }
   }
 }

Sources/SwiftExtract/ExtractedDecls.swift

@@ -54,6 +54,16 @@ public final class ExtractedNominalType: ExtractedSwiftDecl {
   public var methods: [ExtractedFunc] = []
   public var variables: [ExtractedFunc] = []
   public var cases: [ExtractedEnumCase] = []
+
+  /// If this nominal is a protocol, stores default implementations provided in protocol extensions.
+  public var protocolDefaultImplementations: [ExtractedFunc] = []
+
+  /// If this nominal is a protocol, stores methods declared in protocol extensions that are not protocol requirements.
+  public var protocolExtensionMethods: [ExtractedFunc] = []
+
+  /// If this nominal is a protocol, stores variables declared in protocol extensions that are not protocol requirements.
+  public var protocolExtensionVariables: [ExtractedFunc] = []
+
   public var inheritedTypes: [SwiftType]
   public var parent: SwiftNominalTypeDeclaration?
 
@@ -119,6 +129,9 @@ public final class ExtractedNominalType: ExtractedSwiftDecl {
     self.methods = base.methods.map { $0.clone(for: selfType) }
     self.variables = base.variables.map { $0.clone(for: selfType) }
     self.cases = base.cases.map { $0.clone(for: selfType) }
+    self.protocolDefaultImplementations = base.protocolDefaultImplementations.map { $0.clone(for: selfType) }
+    self.protocolExtensionMethods = base.protocolExtensionMethods.map { $0.clone(for: selfType) }
+    self.protocolExtensionVariables = base.protocolExtensionVariables.map { $0.clone(for: selfType) }
     self.inheritedTypes = base.inheritedTypes
     self.parent = base.parent
 
@@ -238,6 +251,90 @@ public final class ExtractedNominalType: ExtractedSwiftDecl {
     }
     return false
   }
+
+  /// Returns whether `candidate` satisfies the protocol requirement `requirement`.
+  public static func satisfiesProtocolRequirement(
+    candidate: ExtractedFunc,
+    requirement: ExtractedFunc
+  ) -> Bool {
+    guard candidate.apiKind == requirement.apiKind,
+      candidate.name == requirement.name,
+      candidate.isStatic == requirement.isStatic,
+      candidate.isClass == requirement.isClass
+    else {
+      return false
+    }
+
+    let cSig = candidate.functionSignature
+    let rSig = requirement.functionSignature
+
+    guard cSig.isAsync == rSig.isAsync else {
+      return false
+    }
+
+    // A non-throwing candidate can satisfy a throwing requirement,
+    // but a throwing candidate cannot satisfy a non-throwing requirement.
+    if cSig.isThrowing && !rSig.isThrowing {
+      return false
+    }
+
+    guard cSig.parameters.count == rSig.parameters.count else {
+      return false
+    }
+
+    for (cParam, rParam) in zip(cSig.parameters, rSig.parameters) {
+      guard cParam.argumentLabel == rParam.argumentLabel,
+        cParam.type == rParam.type,
+        cParam.convention == rParam.convention,
+        cParam.isVariadic == rParam.isVariadic
+      else {
+        return false
+      }
+    }
+
+    guard cSig.result == rSig.result else {
+      return false
+    }
+
+    return true
+  }
+
+  /// Finds an existing protocol requirement that `candidate` satisfies in this protocol
+  /// or in any transitively inherited protocol.
+  public func findProtocolRequirement(
+    satisfying candidate: ExtractedFunc,
+    in extractedTypes: [SwiftTypeName: ExtractedNominalType] = [:]
+  ) -> ExtractedFunc? {
+    var visited: Set<ObjectIdentifier> = []
+    var queue: [ExtractedNominalType] = [self]
+
+    while !queue.isEmpty {
+      let current = queue.removeFirst()
+      guard visited.insert(ObjectIdentifier(current)).inserted else { continue }
+
+      let pool =
+        (candidate.apiKind == .getter || candidate.apiKind == .setter)
+        ? current.variables
+        : current.methods
+
+      if let match = pool.first(where: {
+        Self.satisfiesProtocolRequirement(candidate: candidate, requirement: $0)
+      }) {
+        return match
+      }
+
+      for inherited in current.inheritedTypes {
+        if let nominalDecl = inherited.asNominalTypeDeclaration,
+          nominalDecl.kind == .protocol,
+          let parentType = extractedTypes[nominalDecl.qualifiedName]
+        {
+          queue.append(parentType)
+        }
+      }
+    }
+
+    return nil
+  }
 }
 
 public struct SpecializationError: Error {

Sources/SwiftExtract/SwiftAnalysisVisitor.swift

@@ -37,6 +37,16 @@ final class SwiftAnalysisVisitor {
   }
   private var deferredConstrainedExtensions: [DeferredConstrainedExtension] = []
 
+  /// Unconstrained protocol extensions deferred until all nominal declarations have been visited
+  private struct DeferredProtocolExtension {
+    var node: ExtensionDeclSyntax
+    var target: ExtractedNominalType
+    var sourceFilePath: String
+  }
+  private var deferredProtocolExtensions: [DeferredProtocolExtension] = []
+
+  private var isVisitingExtension: Bool = false
+
   private static func operatorKind(_ node: FunctionDeclSyntax) -> SwiftAPIKind? {
     switch node.name.tokenKind {
     case .binaryOperator:
@@ -170,14 +180,24 @@ final class SwiftAnalysisVisitor {
     }
 
     guard !constraints.isEmpty else {
+      if extractedNominalType.swiftNominal.kind == .protocol {
+        self.deferredProtocolExtensions.append(
+          .init(node: node, target: extractedNominalType, sourceFilePath: sourceFilePath)
+        )
+        return
+      }
+
       // The extension is unconstrained: add to the base type (visible through all specializations)
       extractedNominalType.inheritedTypes +=
         node.inheritanceClause?.inheritedTypes.compactMap {
           try? SwiftType($0.type, lookupContext: analyzer.lookupContext)
         } ?? []
+      let wasVisitingExtension = self.isVisitingExtension
+      self.isVisitingExtension = true
       for memberItem in node.memberBlock.members {
         self.visit(decl: memberItem.decl, in: extractedNominalType, sourceFilePath: sourceFilePath)
       }
+      self.isVisitingExtension = wasVisitingExtension
       return
     }
 
@@ -205,11 +225,14 @@ final class SwiftAnalysisVisitor {
     }
 
     // Visit members in each matching specialization, not the base type
+    let wasVisitingExtension = self.isVisitingExtension
+    self.isVisitingExtension = true
     for specialized in matchingSpecializations {
       for memberItem in node.memberBlock.members {
         self.visit(decl: memberItem.decl, in: specialized, sourceFilePath: sourceFilePath)
       }
     }
+    self.isVisitingExtension = wasVisitingExtension
   }
 
   // ==== -----------------------------------------------------------------------
@@ -382,7 +405,17 @@ final class SwiftAnalysisVisitor {
 
     log.debug("Record extracted method \(node.qualifiedNameForDebug)")
     if let typeContext {
-      typeContext.methods.append(extracted)
+      if typeContext.swiftNominal.kind == .protocol && isVisitingExtension {
+        if let requirement = typeContext.findProtocolRequirement(satisfying: extracted, in: analyzer.extractedTypes) {
+          log.debug("Record protocol default implementation for \(requirement.displayName): \(extracted.displayName)")
+          typeContext.protocolDefaultImplementations.append(extracted)
+        } else {
+          log.debug("Record protocol extension method \(extracted.displayName)")
+          typeContext.protocolExtensionMethods.append(extracted)
+        }
+      } else {
+        typeContext.methods.append(extracted)
+      }
     } else {
       analyzer.extractedGlobalFuncs.append(extracted)
     }
@@ -532,7 +565,11 @@ final class SwiftAnalysisVisitor {
       sourceFilePath: sourceFilePath,
     )
 
-    typeContext.initializers.append(extracted)
+    if typeContext.swiftNominal.kind == .protocol && isVisitingExtension {
+      typeContext.protocolExtensionMethods.append(extracted)
+    } else {
+      typeContext.initializers.append(extracted)
+    }
   }
 
   private func visit(
@@ -646,7 +683,17 @@ final class SwiftAnalysisVisitor {
       "Record extracted variable accessor \(kind == .getter || kind == .subscriptGetter ? "getter" : "setter"):\(node.qualifiedNameForDebug)"
     )
     if let typeContext {
-      typeContext.variables.append(extracted)
+      if typeContext.swiftNominal.kind == .protocol && isVisitingExtension {
+        if let requirement = typeContext.findProtocolRequirement(satisfying: extracted, in: analyzer.extractedTypes) {
+          log.debug("Record protocol default implementation for variable \(requirement.displayName): \(extracted.displayName)")
+          typeContext.protocolDefaultImplementations.append(extracted)
+        } else {
+          log.debug("Record protocol extension variable \(extracted.displayName)")
+          typeContext.protocolExtensionVariables.append(extracted)
+        }
+      } else {
+        typeContext.variables.append(extracted)
+      }
     } else {
       analyzer.extractedGlobalVariables.append(extracted)
     }
@@ -805,6 +852,30 @@ final class SwiftAnalysisVisitor {
     deferredConstrainedExtensions.removeAll()
   }
 
+  /// Process unconstrained protocol extensions that were deferred until all nominal declarations were visited
+  func processDeferredProtocolExtensions() {
+    let deferred = self.deferredProtocolExtensions
+    self.deferredProtocolExtensions = []
+
+    for item in deferred {
+      let extractedNominalType = item.target
+      let node = item.node
+      let sourceFilePath = item.sourceFilePath
+
+      extractedNominalType.inheritedTypes +=
+        node.inheritanceClause?.inheritedTypes.compactMap {
+          try? SwiftType($0.type, lookupContext: analyzer.lookupContext)
+        } ?? []
+
+      let wasVisitingExtension = self.isVisitingExtension
+      self.isVisitingExtension = true
+      for memberItem in node.memberBlock.members {
+        self.visit(decl: memberItem.decl, in: extractedNominalType, sourceFilePath: sourceFilePath)
+      }
+      self.isVisitingExtension = wasVisitingExtension
+    }
+  }
+
   // ==== -----------------------------------------------------------------------
   // MARK: Constrained extension merging
 

Sources/SwiftExtract/SwiftAnalyzer.swift

@@ -174,6 +174,8 @@ extension SwiftAnalyzer {
       visitor.visit(inputFile: input)
     }
 
+    visitor.processDeferredProtocolExtensions()
+
     try hook(self)
 
     // Apply any specializations registered after their target types were visited

Tests/JExtractSwiftTests/Asserts/TextAssertions.swift

@@ -45,6 +45,7 @@ func assertOutput(
   moduleJavaPackages: [String: String] = [:],
   expectedChunks: [String],
   notExpectedChunks: [String] = [],
+  expectedOccurrences: [String: Int] = [:],
   fileID: String = #fileID,
   filePath: String = #filePath,
   line: Int = #line,
@@ -112,6 +113,7 @@ func assertOutput(
     dump: dump,
     expectedChunks: expectedChunks,
     notExpectedChunks: notExpectedChunks,
+    expectedOccurrences: expectedOccurrences,
     detectChunkByInitialLines: _detectChunkByInitialLines,
     fileID: fileID,
     filePath: filePath,
@@ -126,6 +128,7 @@ func assertOutput(
   dump: Bool = false,
   expectedChunks: [String],
   notExpectedChunks: [String] = [],
+  expectedOccurrences: [String: Int] = [:],
   detectChunkByInitialLines _detectChunkByInitialLines: Int = 4,
   fileID: String = #fileID,
   filePath: String = #filePath,
@@ -149,6 +152,19 @@ func assertOutput(
     )
   }
 
+  for (expectedChunk, expectedCount) in expectedOccurrences {
+    let actualCount = output.components(separatedBy: expectedChunk).count - 1
+    #expect(
+      actualCount == expectedCount,
+      """
+      Expected '\(expectedChunk)' to occur \(expectedCount) time(s), but found \(actualCount).
+      ==== Got output ====
+      \(output)
+      """,
+      sourceLocation: sourceLocation
+    )
+  }
+
   let gotLines = output.split(separator: "\n").filter { l in
     l.trimmingCharacters(in: CharacterSet.whitespacesAndNewlines).count > 0
   }