FFMSwift2JavaGenerator+SwiftThunkPrinting.swift

//===----------------------------------------------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2024 Apple Inc. and the Swift.org project authors
// Licensed under Apache License v2.0
//
// See LICENSE.txt for license information
// See CONTRIBUTORS.txt for the list of Swift.org project authors
//
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//

import CodePrinting
import SwiftSyntax
import SwiftSyntaxBuilder

extension FFMSwift2JavaGenerator {
  package func writeSwiftThunkSources() throws {
    var printer = CodePrinter()
    try writeSwiftThunkSources(printer: &printer)
  }

  package func writeSwiftExpectedEmptySources() throws {
    let pendingFileCount = self.expectedOutputSwiftFileNames.count
    guard pendingFileCount > 0 else {
      return // no need to write any empty files, yay
    }

    log.info(
      "[swift-java] Write empty [\(self.expectedOutputSwiftFileNames.count)] 'expected' files in: \(swiftOutputDirectory)/"
    )

    for expectedFileName in self.expectedOutputSwiftFileNames {
      log.info("Write SwiftPM-'expected' empty file: \(expectedFileName.bold)")

      var printer = CodePrinter()
      printer.print("// Empty file generated on purpose")
      _ = try printer.writeContents(
        outputDirectory: self.swiftOutputDirectory,
        javaPackagePath: nil,
        filename: expectedFileName,
      )
    }
  }

  package func writeSwiftThunkSources(printer: inout CodePrinter) throws {
    // Skip global thunks when generating for a single type
    if config.singleType == nil {
      let moduleFilenameBase = "\(self.swiftModuleName)Module+SwiftJava"
      let moduleFilename = "\(moduleFilenameBase).swift"
      do {
        log.debug("Printing contents: \(moduleFilename)")

        try printGlobalSwiftThunkSources(&printer)

        if let outputFile = try printer.writeContents(
          outputDirectory: self.swiftOutputDirectory,
          javaPackagePath: nil,
          filename: moduleFilename,
        ) {
          log.info("Generated: \(moduleFilenameBase.bold).swift (at \(outputFile.absoluteString))")
          self.expectedOutputSwiftFileNames.remove(moduleFilename)
        }
      } catch {
        log.warning("Failed to write to Swift thunks: \(moduleFilename)")
      }
    }

    // === All types
    // We have to write all types to their corresponding output file that matches the file they were declared in,
    // because otherwise SwiftPM plugins will not pick up files apropriately -- we expect 1 output +SwiftJava.swift file for every input.

    let filteredTypes: [String: ImportedNominalType]
    if let singleType = config.singleType {
      filteredTypes = self.analysis.importedTypes.filter { $0.key == singleType }
    } else {
      filteredTypes = self.analysis.importedTypes
    }

    for group: (key: String, value: [Dictionary<String, ImportedNominalType>.Element]) in Dictionary(
      grouping: filteredTypes,
      by: { $0.value.sourceFilePath },
    ) {
      log.warning("Writing types in file group: \(group.key): \(group.value.map(\.key))")

      let importedTypesForThisFile = group.value
        .map(\.value)
        .sorted(by: { $0.qualifiedName < $1.qualifiedName })

      let inputFileName = "\(group.key)".split(separator: "/").last ?? "__Unknown.swift"
      let filename = "\(inputFileName)".replacing(/\.swift(interface)?/, with: "+SwiftJava.swift")

      // Print file header before all type thunks
      printer.print(
        """
        // Generated by swift-java

        import SwiftRuntimeFunctions

        """
      )
      self.lookupContext.symbolTable.printImportedModules(&printer)

      for ty in importedTypesForThisFile {
        log.info("Printing Swift thunks for type: \(ty.qualifiedName.bold)")
        printer.printSeparator("Thunks for \(ty.qualifiedName)")

        do {
          try printSwiftThunkSources(&printer, ty: ty)
        } catch {
          log.warning("Failed to print to Swift thunks for type'\(ty.qualifiedName)' to '\(filename)', error: \(error)")
        }

      }

      log.warning("Write Swift thunks file: \(filename.bold)")
      do {
        if let outputFile = try printer.writeContents(
          outputDirectory: self.swiftOutputDirectory,
          javaPackagePath: nil,
          filename: filename,
        ) {
          log.info("Done writing Swift thunks to: \(outputFile.absoluteString)")
          self.expectedOutputSwiftFileNames.remove(filename)
        }
      } catch {
        log.warning("Failed to write to Swift thunks: \(filename), error: \(error)")
      }
    }
  }

  public func printGlobalSwiftThunkSources(_ printer: inout CodePrinter) throws {
    let stt = SwiftThunkTranslator(self)

    printer.print(
      """
      // Generated by swift-java

      import SwiftRuntimeFunctions

      """
    )

    self.lookupContext.symbolTable.printImportedModules(&printer)

    self.currentJavaIdentifiers = JavaIdentifierFactory(
      self.analysis.importedGlobalFuncs + self.analysis.importedGlobalVariables
    )

    for thunk in stt.renderGlobalThunks() {
      printer.print(thunk)
      printer.println()
    }
  }

  public func printSwiftThunkSources(_ printer: inout CodePrinter, decl: ImportedFunc) {
    let stt = SwiftThunkTranslator(self)

    for thunk in stt.render(forFunc: decl) {
      printer.print(thunk)
      printer.println()
    }
  }

  package func printSwiftThunkSources(_ printer: inout CodePrinter, ty: ImportedNominalType) throws {
    let stt = SwiftThunkTranslator(self)

    self.currentJavaIdentifiers = JavaIdentifierFactory(
      ty.initializers + ty.variables + ty.methods
    )

    for thunk in stt.renderThunks(forType: ty) {
      printer.print("\(thunk)")
      printer.print("")
    }
  }
}

struct SwiftThunkTranslator {

  let st: FFMSwift2JavaGenerator

  init(_ st: FFMSwift2JavaGenerator) {
    self.st = st
  }

  func renderGlobalThunks() -> [DeclSyntax] {
    var decls: [DeclSyntax] = []
    decls.reserveCapacity(
      st.analysis.importedGlobalVariables.count + st.analysis.importedGlobalFuncs.count
    )

    for decl in st.analysis.importedGlobalVariables {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in st.analysis.importedGlobalFuncs {
      decls.append(contentsOf: render(forFunc: decl))
    }

    return decls
  }

  /// Render all the thunks that make Swift methods accessible to Java.
  func renderThunks(forType nominal: ImportedNominalType) -> [DeclSyntax] {
    var decls: [DeclSyntax] = []
    decls.reserveCapacity(
      1 + nominal.initializers.count + nominal.variables.count + nominal.methods.count
    )

    decls.append(renderSwiftTypeAccessor(nominal))

    for decl in nominal.initializers {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in nominal.variables {
      decls.append(contentsOf: render(forFunc: decl))
    }

    for decl in nominal.methods {
      decls.append(contentsOf: render(forFunc: decl))
    }

    // Add special thunks for known types (e.g. Data)
    decls.append(contentsOf: renderSpecificTypeThunks(nominal))

    return decls
  }

  /// Accessor to get the `T.self` of the Swift type, without having to rely on mangled name lookups.
  func renderSwiftTypeAccessor(_ nominal: ImportedNominalType) -> DeclSyntax {
    let funcName = SwiftKitPrinting.Names.getType(
      module: st.swiftModuleName,
      nominal: nominal,
    )

    return
      """
      @_cdecl("\(raw: funcName)")
      public func \(raw: funcName)() -> UnsafeMutableRawPointer /* Any.Type */ {
        return unsafeBitCast(\(raw: nominal.swiftNominal.qualifiedName).self, to: UnsafeMutableRawPointer.self)
      }
      """
  }

  func render(forFunc decl: ImportedFunc) -> [DeclSyntax] {
    st.log.trace("Rendering thunks for: \(decl.displayName)")

    let thunkName = st.thunkNameRegistry.functionThunkName(decl: decl)
    guard let translated = st.translatedDecl(for: decl) else {
      return []
    }

    let thunkFunc = translated.loweredSignature.cdeclThunk(
      cName: thunkName,
      swiftAPIName: decl.name,
      as: decl.apiKind,
    )
    return [DeclSyntax(thunkFunc)]
  }

  /// Render special thunks for known types like Foundation.Data
  func renderSpecificTypeThunks(_ nominal: ImportedNominalType) -> [DeclSyntax] {
    guard let knownType = nominal.swiftNominal.knownTypeKind else {
      return []
    }

    switch knownType {
    case .foundationData, .essentialsData:
      return renderFoundationDataThunks(nominal)
    default:
      return []
    }
  }

  /// Render Swift thunks for Foundation.Data helper methods
  private func renderFoundationDataThunks(_ nominal: ImportedNominalType) -> [DeclSyntax] {
    let thunkName = "swiftjava_\(st.swiftModuleName)_\(nominal.swiftNominal.name)_copyBytes__"
    let qualifiedName = nominal.swiftNominal.qualifiedName

    let copyBytesThunk: DeclSyntax =
      """
      @_cdecl("\(raw: thunkName)")
      public func \(raw: thunkName)(
          selfPointer: UnsafeRawPointer,
          destinationPointer: UnsafeMutableRawPointer,
          count: Int
      ) {
          let data = selfPointer.assumingMemoryBound(to: \(raw: qualifiedName).self).pointee
          data.withUnsafeBytes { buffer in
              destinationPointer.copyMemory(from: buffer.baseAddress!, byteCount: count)
          }
      }
      """

    return [copyBytesThunk]
  }
}