get version of any TASTy file

tastyUtils.scala

//> using scala "3.4.2"

//> using dependency "com.lihaoyi::os-lib:0.10.2"
//> using dependency "com.lihaoyi::sourcecode:0.4.2"
//> using dependency "io.get-coursier:coursier_2.13:2.1.10"

//> using buildInfo

// run scala-cli with --power and -with-compiler flags

import scala.sys.process.*
import java.nio.file.{Files, Paths, Path}
import java.io.File.pathSeparator as cpSep
import scala.annotation.threadUnsafe as tu
import coursier.*
import scala.util.CommandLineParser
import java.net.URLClassLoader
import java.lang.invoke.MethodHandles
import java.lang.invoke.MethodType
import scala.annotation.static
import scala.util.Using
import java.util.jar.JarInputStream
import java.util.jar.JarFile

lazy val out = os.Path(sourcecode.File()) / os.up / "out"

/** Usage example: `scala-cli tastyUtils.scala --power -with-compiler -M printTastyVersion -- 3.1.1` */
@main def printTastyVersion(compiler: String): Unit =
  println(getTastyVersion(compiler).show)

def blue(str: String) = Console.BLUE + str + Console.RESET
def red(str: String) = Console.RED + str + Console.RESET
def green(str: String) = Console.GREEN + str + Console.RESET

/** Usage example: `scala-cli tastyUtils.scala --power -with-compiler -M canReadTasty -- 28.4-experimental-1 28.3` */
@main def canReadTasty(read: TastyVersion, file: TastyVersion): Boolean =
  val res = file <:< read
  val op = if res then "✅" else "❌"
  val verb = if res then green("can read") else red("cannot read")
  println(s"$op ${blue(read.show)} $verb ${blue(file.show)}")
  res

/** Usage example: `scala-cli tastyUtils.scala --power -with-compiler -M cleanup */
@main def cleanup(): Unit =
  os.remove.all(out)

/** Usage example: `scala-cli tastyUtils.scala --power -with-compiler -M canReadCompiled -- 3.3.3 3.4.2` */
@main def canReadCompiled(read: String, file: String): Unit =
  val readTasty = getTastyVersion(read)
  val fileTasty = getTastyVersion(file)
  val res = canReadTasty(readTasty, fileTasty)
  val op = if res then "✅" else "❌"
  val verb = if res then green("can read") else red("cannot read")
  println(s"$op ${blue(read)} $verb ${blue(file)}")

/** Usage example: `scala-cli tastyUtils.scala --power -with-compiler -M printTastyOfLibrary -- com.lihaoyi::utest:0.8.3 3.4.2` */
@main def printTastyOfLibrary(coord: String, compiler: String): Unit =
  val compilerTasty = getTastyVersion(compiler)
  val libTasty = tastyOfLibraryImpl(coord, compilerTasty)
  println(s"TASTy ${green(libTasty.show)}, read by Scala ${blue(compiler)}")

class Invoker

object Invoker:
  @static private val tastyCore = fetchTastyCore(scala.cli.build.BuildInfo.scalaVersion)
  @static private val cp = new URLClassLoader(tastyCore.map(_.toURI().toURL()).toArray)
  @static private val TastyFormat = cp.loadClass("dotty.tools.tasty.TastyFormat")
  @static private val lookup = MethodHandles.lookup()
  @static val compare2 =
    lookup.findStatic(TastyFormat, "isVersionCompatible", MethodType.methodType(classOf[Boolean], classOf[Int], classOf[Int], classOf[Int], classOf[Int], classOf[Int], classOf[Int]))

case class TastyVersion(major: Int, minor: Int, exp: Int):
  def suffix: String = if exp == 0 then "" else s"-experimental-$exp"
  def show: String = s"$major.$minor$suffix"

object TastyVersion:
  val Format = raw"(\d+).(\d+)(?:-experimental-(\d+))?".r

  extension (file: TastyVersion)
    def <:< (compiler: TastyVersion): Boolean =
      Invoker.compare2.invokeExact(file.major, file.minor, file.exp, compiler.major, compiler.minor, compiler.exp): Boolean

  given CommandLineParser.FromString[TastyVersion] = s =>
    parse(s).getOrElse(throw IllegalArgumentException(s"Invalid TastyVersion: $s"))

  def parse(string: String): Option[TastyVersion] = string match
    case Format(maj, min, exp) => Some(TastyVersion(maj.toInt, min.toInt, if exp == null then 0 else exp.toInt))
    case _ => None

def tastyOfLibraryImpl(coord: String, compilerTasty: TastyVersion): TastyVersion =
  import scala.util.boundary, boundary.break
  import dotty.tools.tasty.UnpicklerConfig
  import dotty.tools.dotc.core.tasty.TastyUnpickler.Scala3CompilerConfig
  import dotty.tools.tasty.{TastyHeaderUnpickler, TastyReader}

  val tastyconfig: UnpicklerConfig = new UnpicklerConfig with Scala3CompilerConfig:
    override def minorVersion: Int = compilerTasty.minor
    override def experimentalVersion: Int = compilerTasty.exp
    override def majorVersion: Int = compilerTasty.major


  val s"$org::$name:$version" = coord: @unchecked
  val lib = fetchLibrary(org, s"${name}_3", version)
  val jar = lib.find(f => f.getName().contains(name) && f.getName.endsWith(".jar")).get
  // TODO: look inside jar for first .tasty file and parse it
  val tastyBytes =
    Using.Manager: use =>
      val jarFile = use(JarFile(jar))
      val entries = jarFile.entries()
      // val jarStream = use(JarInputStream(jar.toURI().toURL().openStream()))
      boundary:
        while entries.hasMoreElements() do
          val e = entries.nextElement()
          if e.getName().endsWith(".tasty") then
            val in = use(jarFile.getInputStream(e))
            break(Some(in.readAllBytes()))
        None
    .get.getOrElse(throw IllegalArgumentException(s"No .tasty file found in $jar"))
  end tastyBytes

  val reader = TastyReader(tastyBytes)
  val header = TastyHeaderUnpickler(tastyconfig, reader).readFullHeader()
  TastyVersion(header.majorVersion, header.minorVersion, header.experimentalVersion)
end tastyOfLibraryImpl

def getTastyVersion(compiler: String): TastyVersion =
  @tu lazy val tastyCore = fetchTastyCore(compiler)

  os.makeDir.all(out)

  val PrintTastyClass = out / "PrintTastyFromClasspath.class"

  if !os.exists(PrintTastyClass) then writePrintTastyClass(out, tastyCore)

  val versionStringOut = Java((out +: tastyCore).mkString(cpSep), "PrintTastyFromClasspath")

  val Some(version) = versionStringOut.linesIterator.toList.headOption.flatMap(TastyVersion.parse): @unchecked

  version

def writePrintTastyClass(out: os.Path, tastyCore: Seq[java.io.File]) =
  val temp = Files.createTempDirectory("tasty-version")
  try {
    val srcFile = temp.resolve("PrintTastyFromClasspath.java")
    Files.write(srcFile, PrintTastyFromClasspath.getBytes)
    Javac(out.toString, tastyCore.mkString(cpSep), srcFile.toString)
  } finally {
    def deleteRecursively(f: java.io.File): Unit =
      if (f.isDirectory)
        f.listFiles.foreach(deleteRecursively)
      f.delete
    deleteRecursively(temp.toFile)
  }

def fetchTastyCore(compiler: String) = fetchLibrary("org.scala-lang", "tasty-core_3", compiler)

def fetchLibrary(org: String, name: String, version: String) = Fetch()
  .addDependencies(
    Dependency(
      Module(Organization(org), ModuleName(name)),
      version
    )
  )
  .run()


object Javac:
  def apply(dir: String, classpath: String, source: String): Unit =
    val javac = javax.tools.ToolProvider.getSystemJavaCompiler()
    javac.run(null, null, null, "-d", dir, "-classpath", classpath, source)

object Java:
  def apply(classpath: String, clazz: String): String =
    val javaHome = Paths.get(System.getProperty("java.home"))
    val java = javaHome.resolve("bin").resolve("java")
    Seq(java.toString, "-cp", classpath, clazz).!!

val PrintTastyFromClasspath = """
import dotty.tools.tasty.TastyFormat;

public class PrintTastyFromClasspath {
  public static void main(String[] args) {
    int maj = TastyFormat.MajorVersion();
    int min = TastyFormat.MinorVersion();
    int exp = TastyFormat.ExperimentalVersion();
    String suffix = exp == 0 ? "" : "-experimental-" + exp;
    System.out.println("" + maj + "." + min + suffix);
  }
}
"""