Mill javah

build.sc

import mill._, scalalib._

trait JavahModule extends JavaModule {

  def javah = T {
    os.walk(compile().classes.path).filter(_.ext == "class").foreach { path =>
      sgjavah.javah(
        path.toNIO,
        T.dest.toNIO
      )
    }
    PathRef(T.dest)
  }
}

// Adapted from Glavo's gjavah library https://github.com/Glavo/gjavah, released
// under MIT license
object sgjavah {
  import org.objectweb.asm
  import scala.collection.mutable

  // see the following link for a description of the name mappings
  // https://docs.oracle.com/javase/8/docs/technotes/guides/jni/spec/design.html#resolving_native_method_names
  private def mangle(name: String): String = {
    val builder = new StringBuilder(name.length() * 2)
    for (char <- name) char match {
      case '.' => builder.append("_")
      case '_' => builder.append("_1")
      case ';' => builder.append("_2")
      case '[' => builder.append("_3")
      case ch
          if ('0' <= ch && ch <= '9') || ('a' <= ch && ch <= 'z') || ('A' <= ch && ch <= 'Z') =>
        builder.append(ch)
      case ch =>
        builder.append(String.format("_0%04x", ch.toInt))
    }
    builder.result()
  }

  private def mapTypeToNative(tpe: asm.Type) = tpe.toString match {
    case "Z"                        => "jboolean"
    case "B"                        => "jbyte"
    case "C"                        => "jchar"
    case "S"                        => "jshort"
    case "I"                        => "jint"
    case "J"                        => "jlong"
    case "F"                        => "jfloat"
    case "D"                        => "jdouble"
    case "V"                        => "void"
    case "Ljava/lang/Class;"        => "jclass"
    case "Ljava/lang/String;"       => "jstring"
    case "Ljava/lang/Throwable;"    => "jthrowable"
    case "Ljava/lang/Error;"        => "jthrowable"
    case "Ljava/lang/Exception;"    => "jthrowable"
    case "[Z"                       => "jbooleanArray"
    case "[B"                       => "jbyteArray"
    case "[C"                       => "jcharArray"
    case "[S"                       => "jshortArray"
    case "[I"                       => "jintArray"
    case "[J"                       => "jlongArray"
    case "[F"                       => "jfloatArray"
    case "[D"                       => "jdoubleArray"
    case arr if arr.startsWith("[") => "jobjectArray"
    case ref if ref.startsWith("L") => "jobject"
    case other                      => throw new IllegalArgumentException("unknown type: " + other)
  }

  private case class NativeConstant(name: String, value: Object)
  private case class NativeMethod(
      access: Int,
      name: String,
      descriptor: String
  ) {
    val tpe: asm.Type = asm.Type.getType(descriptor)
    def isStatic: Boolean = (access & asm.Opcodes.ACC_STATIC) != 0
  }
  private class ClassVisitor extends asm.ClassVisitor(asm.Opcodes.ASM7) {
    val constants = mutable.ListBuffer.empty[NativeConstant]
    val nativeMethods = mutable.ListBuffer.empty[NativeMethod]
    val methodCounts = mutable.Map.empty[String, Int]

    var className: String = _

    override def visit(
        version: Int,
        access: Int,
        name: String,
        signature: String,
        superName: String,
        interfaces: Array[String]
    ): Unit = {
      className = name
    }

    override def visitMethod(
        access: Int,
        name: String,
        descriptor: String,
        signature: String,
        exceptions: Array[String]
    ): asm.MethodVisitor = {
      methodCounts(name) = methodCounts.getOrElse(name, 0) + 1
      if ((access & asm.Opcodes.ACC_NATIVE) != 0)
        nativeMethods += NativeMethod(access, name, descriptor)
      return null
    }

    override def visitField(
        access: Int,
        name: String,
        descriptor: String,
        signature: String,
        value: Object
    ): asm.FieldVisitor = {
      if (value != null && !value.isInstanceOf[String])
        constants += NativeConstant(name, value)
      return null
    }
  }

  def javah(
      in: java.io.InputStream,
      mkOut: String => java.io.PrintStream
  ): Unit = {
    val visitor = new ClassVisitor()
    val reader = new asm.ClassReader(in)
    reader.accept(
      visitor,
      asm.ClassReader.SKIP_CODE | asm.ClassReader.SKIP_DEBUG | asm.ClassReader.SKIP_FRAMES
    )

    if (visitor.nativeMethods.isEmpty) return

    val className = mangle(visitor.className.replace("/", "."))
    val out = mkOut(className)

    out.println("/* DO NOT EDIT THIS FILE - it is machine generated */")
    out.println("#include <jni.h>")
    out.println(s"/* Header for class ${className} */")
    out.println()
    out.println(s"#ifndef _Included_${className}")
    out.println(s"#define _Included_${className}")
    out.println("#ifdef __cplusplus");
    out.println("extern \"C\" {");
    out.println("#endif");

    for (constant <- visitor.constants) {
      val name = s"${className}_${mangle(constant.name)}"
      val value = constant.value match {
        case x: java.lang.Double    => x.toString
        case x: java.lang.Float     => x.toString + "f"
        case x: java.lang.Long      => x.toString + "i64"
        case x: java.lang.Character => x.toInt.toString + "L"
        case x                      => x.toString + "L"
      }
      out.println(s"#undef $name")
      out.println(s"#define $name $value")
    }

    for (method <- visitor.nativeMethods) {
      val ret = mapTypeToNative(method.tpe.getReturnType())
      val args = mutable.ListBuffer.empty[String]
      args += "JNIEnv *"
      args += (if (method.isStatic) "jclass" else "jobject")
      method.tpe.getArgumentTypes().foreach { tpe =>
        args += mapTypeToNative(tpe)
      }

      val name = if (visitor.methodCounts(method.name) == 1) {
        s"Java_${className}_${mangle(method.name)}"
      } else {
        s"Java_${className}_${mangle(method.name)}__${mangle(method.tpe.getArgumentTypes.mkString(""))}"
      }

      out.println("/*")
      out.println(s" * Class:      ${className}")
      out.println(s" * Method:     ${mangle(method.name)}")
      out.println(s" * Signature:  ${method.tpe.toString()}")
      out.println(" */")
      out.println(s"JNIEXPORT ${ret} JNICALL ${name}")
      out.println("  " + args.mkString("(", ", ", ");"))
      out.println()
    }

    out.println("#ifdef __cplusplus")
    out.println("}")
    out.println("#endif")
    out.println("#endif")
  }

  def javah(
      classFile: java.nio.file.Path,
      outDir: java.nio.file.Path
  ): Unit = {
    var in: java.io.InputStream = null
    try {
      in = java.nio.file.Files.newInputStream(classFile)
      javah(
        in,
        name => {
          val path = outDir.resolve(name + ".h")
          val stream = java.nio.file.Files.newOutputStream(path)
          new java.io.PrintStream(stream)
        }
      )
    } finally {
      if (in != null) in.close()
    }
  }

}