get version of any TASTy file

tastyUtils.scala

//> using scala "3.5.2"
//> using dep "com.lihaoyi::os-lib:0.11.2"
//> using dep "com.lihaoyi::sourcecode:0.4.2"
//> using dep "io.get-coursier:coursier_2.13:2.1.14"
//> using buildInfo

// setup with `scala --power setup-ide -with-compiler .`

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
import java.lang.invoke.MethodHandle

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

/** Usage example: `scala . --power -with-compiler -M hasCorrectTasty -- 3.6.0`
  */
@main def hasCorrectTasty(compiler: String): Unit =
  val version = getTastyVersion(compiler)
  def fail = println(
    s"❌ ${red("Scala")} ${blue(compiler)} ${red("has incorrect TASTy version:")} ${blue(version.show)}"
  )
  def succ = println(
    s"✅ ${green("Scala")} ${blue(compiler)} ${green("has correct TASTy version:")} ${blue(version.show)}"
  )
  compiler match
    case s"3.$m.$p" if m.toIntOption.isDefined && p.toIntOption.isDefined =>
      if version.major == 28 && version.minor == m.toInt && version.exp == 0
      then succ
      else fail
    case s"3.$m.$p-$rest"
        if m.toIntOption.isDefined && p.toIntOption.isDefined =>
      if version.major == 28 then
        val (minor, patch) = (m.toInt, p.toInt)
        if version.exp > 0 then
          if patch != 0 && minor == version.minor then fail
          else succ
        else if patch == 0 || version.minor != minor then fail
        else succ
      else fail
    case _ => fail

/** Usage example: `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
def yellow(str: String) = Console.YELLOW + str + Console.RESET

/** Usage example: `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 . --power -with-compiler -M cleanup
  */
@main def cleanup(): Unit =
  os.remove.all(out)

/** Usage example: `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 . --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"${blue(coord)} has TASTy ${blue(libTasty.show)}, can be read by Scala ${blue(compiler)} ${yellow(s"(parsed by Scala ${scala.cli.build.BuildInfo.scalaVersion})")}"
  )

class Invoker

object Invoker:
  def find(): MethodHandle =
    val tastyCore = fetchTastyCore(
      scala.cli.build.BuildInfo.scalaVersion
    )
    val mt = MethodType.methodType(
      classOf[Boolean],
      classOf[Int],
      classOf[Int],
      classOf[Int],
      classOf[Int],
      classOf[Int],
      classOf[Int]
    )
    tastyCore
      .map: files =>
        val cp = new URLClassLoader(
          files.map(_.toURI().toURL()).toArray
        )
        val TastyFormat =
          cp.loadClass("dotty.tools.tasty.TastyFormat")
        lookup.findStatic(TastyFormat, "isVersionCompatible", mt)
      .getOrElse(MethodHandles.empty(mt))

  @static private val lookup = MethodHandles.lookup()
  @static val compare2 = find()

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, null) => Some(TastyVersion(maj.toInt, min.toInt, 0))
    case Format(maj, min, exp) =>
      Some(TastyVersion(maj.toInt, min.toInt, 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 =
  val res =
    for tastyCore <- fetchTastyCore(compiler) yield
      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
    end for
  res.getOrElse:
    Console.err.println(s"Failed to fetch tasty-core for Scala $compiler")
    sys.exit(1)

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) =
  try Some(fetchLibrary("org.scala-lang", "tasty-core_3", compiler))
  catch
    case e: Exception =>
      None

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);
  }
}
"""