Use nominal type resolution to translate members of extensions

Use the new nominal type resolution functionality within jextract-swift
to resolve extensions to their corresponding nominal types, so that
members of those extensions show up as member of the resulting Java
class.

Author: DougGregor

Sources/JExtractSwift/NominalTypeResolution.swift

@@ -17,10 +17,6 @@ import SwiftSyntax
 /// their extended nominal types and mapping type names to their full names.
 @_spi(Testing)
 public class NominalTypeResolution {
-  /// A syntax node for a nominal type declaration.
-  @_spi(Testing)
-  public typealias NominalTypeDeclSyntaxNode = any DeclGroupSyntax & NamedDeclSyntax
-
   /// Mapping from the syntax identifier for a given type declaration node,
   /// such as StructDeclSyntax, to the set of extensions of this particular
   /// type.
@@ -39,6 +35,10 @@ public class NominalTypeResolution {
   @_spi(Testing) public init() { }
 }
 
+/// A syntax node for a nominal type declaration.
+@_spi(Testing)
+public typealias NominalTypeDeclSyntaxNode = any DeclGroupSyntax & NamedDeclSyntax
+
 // MARK: Nominal type name resolution.
 extension NominalTypeResolution {
   /// Compute the fully-qualified name of the given nominal type node.
@@ -170,6 +170,7 @@ extension NominalTypeResolution {
   ///
   /// Returns the list of extensions that could not be resolved.
   @_spi(Testing)
+  @discardableResult
   public func bindExtensions() -> [ExtensionDeclSyntax] {
     while !unresolvedExtensions.isEmpty {
       // Try to resolve all of the unresolved extensions.

Sources/JExtractSwift/Swift2JavaTranslator.swift

@@ -41,6 +41,8 @@ public final class Swift2JavaTranslator {
   /// type representation.
   public var importedTypes: [String: ImportedNominalType] = [:]
 
+  let nominalResolution: NominalTypeResolution = NominalTypeResolution()
+
   public init(
     javaPackage: String,
     swiftModuleName: String
@@ -82,6 +84,10 @@ extension Swift2JavaTranslator {
 
     let sourceFileSyntax = Parser.parse(source: text)
 
+    // Find all of the types and extensions, then bind the extensions.
+    nominalResolution.addSourceFile(sourceFileSyntax)
+    nominalResolution.bindExtensions()
+
     let visitor = Swift2JavaVisitor(
       moduleName: self.swiftModuleName,
       targetJavaPackage: self.javaPackage,

Sources/JExtractSwift/Swift2JavaVisitor.swift

@@ -25,6 +25,7 @@ final class Swift2JavaVisitor: SyntaxVisitor {
   /// store this along with type names as we import them.
   let targetJavaPackage: String
 
+  /// The current type name as a nested name like A.B.C.
   var currentTypeName: String? = nil
 
   var log: Logger { translator.log }
@@ -37,27 +38,46 @@ final class Swift2JavaVisitor: SyntaxVisitor {
     super.init(viewMode: .all)
   }
 
-  override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
-    guard node.shouldImport(log: log) else {
-      return .skipChildren
+  /// Try to resolve the given nominal type node into its imported
+  /// representation.
+  func resolveNominalType(
+    _ nominal: some DeclGroupSyntax & NamedDeclSyntax,
+    kind: NominalTypeKind
+  ) -> ImportedNominalType? {
+    if !nominal.shouldImport(log: log) {
+      return nil
+    }
+
+    guard let fullName = translator.nominalResolution.fullyQualifiedName(of: nominal) else {
+      return nil
+    }
+
+    if let alreadyImported = translator.importedTypes[fullName] {
+      return alreadyImported
     }
 
-    log.info("Import: \(node.kind) \(node.name)")
-    let typeName = node.name.text
-    currentTypeName = typeName
-    translator.importedTypes[typeName] = ImportedNominalType(
-      // TODO: support nested classes (parent name here)
+    let importedNominal = ImportedNominalType(
       name: ImportedTypeName(
-        swiftTypeName: typeName,
+        swiftTypeName: fullName,
         javaType: .class(
           package: targetJavaPackage,
-          name: typeName
+          name: fullName
         ),
-        swiftMangledName: node.mangledNameFromComment
+        swiftMangledName: nominal.mangledNameFromComment
       ),
-      kind: .class
+      kind: kind
     )
 
+    translator.importedTypes[fullName] = importedNominal
+    return importedNominal
+  }
+
+  override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
+    guard let importedNominalType = resolveNominalType(node, kind: .class) else {
+      return .skipChildren
+    }
+
+    currentTypeName = importedNominalType.name.swiftTypeName
     return .visitChildren
   }
 
@@ -68,6 +88,24 @@ final class Swift2JavaVisitor: SyntaxVisitor {
     }
   }
 
+  override func visit(_ node: ExtensionDeclSyntax) -> SyntaxVisitorContinueKind {
+    // Resolve the extended type of the extension as an imported nominal, and
+    // recurse if we found it.
+    guard let nominal = translator.nominalResolution.extendedType(of: node),
+          let importedNominalType = resolveNominalType(nominal, kind: .class) else {
+      return .skipChildren
+    }
+
+    currentTypeName = importedNominalType.name.swiftTypeName
+    return .visitChildren
+  }
+
+  override func visitPost(_ node: ExtensionDeclSyntax) {
+    if currentTypeName != nil {
+      currentTypeName = nil
+    }
+  }
+
   override func visit(_ node: FunctionDeclSyntax) -> SyntaxVisitorContinueKind {
     guard node.shouldImport(log: log) else {
       return .skipChildren
@@ -182,7 +220,7 @@ final class Swift2JavaVisitor: SyntaxVisitor {
   }
 }
 
-extension ClassDeclSyntax {
+extension DeclGroupSyntax where Self: NamedDeclSyntax {
   func shouldImport(log: Logger) -> Bool {
     guard (accessControlModifiers.first { $0.isPublic }) != nil else {
       log.trace("Cannot import \(self.name) because: is not public")

Tests/JExtractSwiftTests/FuncImportTests.swift

@@ -37,6 +37,11 @@ final class MethodImportTests: XCTestCase {
     // MANGLED NAME: $s14MySwiftLibrary23globalTakeLongIntString1l3i321sys5Int64V_s5Int32VSStF
     public func globalTakeIntLongString(i32: Int32, l: Int64, s: String)
 
+    extension MySwiftClass {
+      // MANGLED NAME: $s14MySwiftLibrary0aB5ClassC22helloMemberFunctionInExtension
+      public func helloMemberInExtension()
+    }
+
     // MANGLED NAME: $s14MySwiftLibrary0aB5ClassCMa
     public class MySwiftClass {
       // MANGLED NAME: $s14MySwiftLibrary0aB5ClassC3len3capACSi_SitcfC
@@ -213,6 +218,47 @@ final class MethodImportTests: XCTestCase {
     )
   }
 
+  func test_method_class_helloMemberInExtension_self_memorySegment() async throws {
+    let st = Swift2JavaTranslator(
+      javaPackage: "com.example.swift",
+      swiftModuleName: "__FakeModule"
+    )
+    st.log.logLevel = .trace
+
+    try await st.analyze(swiftInterfacePath: "/fake/__FakeModule/SwiftFile.swiftinterface", text: class_interfaceFile)
+
+    let funcDecl: ImportedFunc = st.importedTypes["MySwiftClass"]!.methods.first {
+      $0.baseIdentifier == "helloMemberInExtension"
+    }!
+
+    let output = CodePrinter.toString { printer in
+      st.printFuncDowncallMethod(&printer, decl: funcDecl, selfVariant: .memorySegment)
+    }
+
+    assertOutput(
+      output,
+      expected:
+        """
+        /**
+         * {@snippet lang=swift :
+         * public func helloMemberInExtension()
+         * }
+         */
+        public static void helloMemberInExtension(java.lang.foreign.MemorySegment self$) {
+            var mh$ = helloMemberInExtension.HANDLE;
+            try {
+                if (TRACE_DOWNCALLS) {
+                    traceDowncall(self$);
+                }
+                mh$.invokeExact(self$);
+            } catch (Throwable ex$) {
+                throw new AssertionError("should not reach here", ex$);
+            }
+        }
+        """
+    )
+  }
+
   func test_method_class_helloMemberFunction_self_wrapper() async throws {
     let st = Swift2JavaTranslator(
       javaPackage: "com.example.swift",