Add :asmp to the REPL

Author: griggt

compiler/src/dotty/tools/repl/Disassembler.scala

@@ -575,3 +575,180 @@ class JavapTask(val loader: ClassLoader, repl: DisassemblerRepl) extends Disasse
     inputs.map(runInput).toList
   end apply
 end JavapTask
+
+object Asmp extends Disassembler:
+  import Disassembler.*
+
+  def apply(opts: DisassemblerOptions)(using repl: DisassemblerRepl): List[DisResult] =
+    val tool = AsmpTool()
+    val clazz = DisassemblyClass(repl.classLoader, repl)
+    tool(opts.flags)(opts.targets.map(clazz.bytes(_)))
+
+  val helps = List(
+    "usage"       -> ":asmp [opts] [path or class or -]...",
+    "-help"       -> "Prints this help message",
+    "-verbose/-v" -> "Stack size, number of locals, method args",
+    "-private/-p" -> "Private classes and members",
+    "-package"    -> "Package-private classes and members",
+    "-protected"  -> "Protected classes and members",
+    "-public"     -> "Public classes and members",
+    "-c"          -> "Disassembled code",
+    "-s"          -> "Internal type signatures",
+    "-filter"     -> "Filter REPL machinery from output",
+    "-raw"        -> "Don't post-process output from ASM",  // TODO for debugging
+    "-decls"      -> "Declarations",
+    "-bridges"    -> "Bridges",
+    "-synthetics" -> "Synthetics",
+  )
+
+  override def filters(target: String, opts: DisassemblerOptions): List[String => String] =
+    val commonFilters = super.filters(target, opts)
+    if opts.flags.contains("-decls") then filterCommentsBlankLines :: commonFilters
+    else squashConsectiveBlankLines :: commonFilters
+
+  private def squashConsectiveBlankLines(s: String) = s.replaceAll("\n{3,}", "\n\n")
+
+  private def filterCommentsBlankLines(s: String): String =
+    val comment = raw"\s*// .*".r
+    def isBlankLine(s: String) = s.trim == ""
+    def isComment(s: String) = comment.matches(s)
+    filteredLines(s, t => !isComment(t) && !isBlankLine(t))
+end Asmp
+
+object AsmpOptions extends DisassemblerOptionParser(Asmp.helps):
+  val defaultToolOptions = List("-protected", "-verbose")
+
+class AsmpTool extends DisassemblyTool:
+  import DisassemblyTool.*
+  import Disassembler.splitHashMember
+  import java.io.{PrintWriter, StringWriter}
+  import scala.tools.asm.{Attribute, ClassReader, Label, Opcodes}
+  import scala.tools.asm.util.{Textifier, TraceClassVisitor}
+  import dotty.tools.backend.jvm.ClassNode1
+
+  enum Mode:
+    case Verbose, Code, Signatures
+
+  class FilteringTextifier(opts: Seq[String], mode: Mode, accessFilter: Int => Boolean, nameFilter: Option[String]) extends Textifier(Opcodes.ASM9):
+    private def keep(access: Int, name: String): Boolean =
+      accessFilter(access) && nameFilter.map(_ == name).getOrElse(true)
+
+    override def visitField(access: Int, name: String, descriptor: String, signature: String, value: Any): Textifier =
+      if keep(access, name) then
+        super.visitField(access, name, descriptor, signature, value)
+        addNewTextifier(discard = (mode == Mode.Signatures))
+      else
+        addNewTextifier(discard = true)
+
+    override def visitMethod(access:Int, name: String, descriptor: String, signature: String, exceptions: Array[String]): Textifier =
+      if keep(access, name) then
+        super.visitMethod(access, name, descriptor, signature, exceptions)
+        addNewTextifier(discard = (mode == Mode.Signatures))
+      else
+        addNewTextifier(discard = true)
+
+    override def visitInnerClass(name: String, outerName: String, innerName: String, access: Int): Unit =
+      if mode == Mode.Verbose && keep(access, name) then
+        super.visitInnerClass(name, outerName, innerName, access)
+
+    override def visitClassAttribute(attribute: Attribute): Unit =
+      if mode == Mode.Verbose && nameFilter.isEmpty then
+        super.visitClassAttribute(attribute)
+
+    override def visitClassAnnotation(descriptor: String, visible: Boolean): Textifier =
+      // suppress ScalaSignature unless -raw given. Should we? TODO
+      if mode == Mode.Verbose && nameFilter.isEmpty && descriptor != "Lscala/reflect/ScalaSignature;" then
+        super.visitClassAnnotation(descriptor, visible)
+      else
+        addNewTextifier(discard = true)
+
+    override def visitSource(file: String, debug: String): Unit =
+      if mode == Mode.Verbose && nameFilter.isEmpty then
+        super.visitSource(file, debug)
+
+    override def visitAnnotation(descriptor: String, visible: Boolean): Textifier =
+      if mode == Mode.Verbose then
+        super.visitAnnotation(descriptor, visible)
+      else
+        addNewTextifier(discard = true)
+
+    override def visitLineNumber(line: Int, start: Label): Unit =
+      if mode == Mode.Verbose then
+        super.visitLineNumber(line, start)
+
+    override def visitMaxs(maxStack: Int, maxLocals: Int): Unit =
+      if mode == Mode.Verbose then
+        super.visitMaxs(maxStack, maxLocals)
+
+    override def visitLocalVariable(name: String, descriptor: String, signature: String, start: Label, end: Label, index: Int): Unit =
+      if mode == Mode.Verbose then
+        super.visitLocalVariable(name, descriptor, signature, start, end, index)
+
+    private def isLabel(s: String) = raw"\s*L\d+\s*".r.matches(s)
+
+    // ugly hack to prevent orphaned label when local vars, max stack not displayed (e.g. in -c mode)
+    override def visitMethodEnd(): Unit = text.size match
+      case 0 =>
+      case n =>
+        if isLabel(text.get(n - 1).toString) then
+          try text.remove(n - 1)
+          catch case _: UnsupportedOperationException => ()
+
+    private def addNewTextifier(discard: Boolean = false): Textifier =
+      val tx = FilteringTextifier(opts, mode, accessFilter, nameFilter)
+      if !discard then text.add(tx.getText());
+      tx
+  end FilteringTextifier
+
+  override def apply(options: Seq[String])(inputs: Seq[Input]): List[DisResult] =
+    def parseAccessOption(opts: Seq[String]): Int =
+      if opts.contains("-public") then Opcodes.ACC_PUBLIC
+      else if opts.contains("-protected") then Opcodes.ACC_PROTECTED
+      else if opts.contains("-private") || opts.contains("-p") then Opcodes.ACC_PRIVATE
+      else 0
+
+    def accessFilter(opts: Seq[String]): Int => Boolean =
+      inline def contains(mask: Int) = (a: Int) => (a & mask) != 0
+      inline def excludes(mask: Int) = (a: Int) => (a & mask) == 0
+      def accessible: Int => Boolean = parseAccessOption(opts) match
+        case Opcodes.ACC_PUBLIC    => contains(Opcodes.ACC_PUBLIC)
+        case Opcodes.ACC_PROTECTED => contains(Opcodes.ACC_PUBLIC | Opcodes.ACC_PROTECTED)
+        case Opcodes.ACC_PRIVATE   => _ => true
+        case _ /*package*/         => excludes(Opcodes.ACC_PRIVATE)
+      def included(access: Int): Boolean = mode(opts) match
+        case Mode.Verbose => true
+        case _ =>
+          val isBridge = contains(Opcodes.ACC_BRIDGE)(access)
+          val isSynthetic = contains(Opcodes.ACC_SYNTHETIC)(access)
+          if isSynthetic && opts.contains("-synthetics") then true
+          else if isBridge && opts.contains("-bridges") then true
+          else if isSynthetic || isBridge then false
+          else true
+      (x: Int) => accessible(x) && included(x)
+
+    def mode(opts: Seq[String]): Mode =
+      if opts.contains("-c") then Mode.Code
+      else if opts.contains("-s") || opts.contains("-decls") then Mode.Signatures
+      else Mode.Verbose  // default
+
+    def runInput(input: Input): DisResult = input match
+      case Input(target, actual, Success(bytes)) =>
+        val sw = StringWriter()
+        val pw = PrintWriter(sw)
+        val node = ClassNode1()
+
+        def nameFilter = splitHashMember(target).map(s => if s.isEmpty then "apply" else s)
+        val tx =
+          if options.contains("-raw") then Textifier()
+          else FilteringTextifier(options, mode(options), accessFilter(options), nameFilter)
+
+        ClassReader(bytes).accept(node, 0)
+        node.accept(TraceClassVisitor(null, tx, pw))
+        pw.flush()
+        DisSuccess(target, sw.toString)
+      case Input(_, _, Failure(e)) =>
+        DisError(e.getMessage)
+    end runInput
+
+    inputs.map(runInput).toList
+end AsmpTool

compiler/src/dotty/tools/repl/ParseResult.scala

@@ -52,6 +52,11 @@ object Load {
   val command: String = ":load"
 }
 
+case class AsmpOf(args: String) extends Command
+object AsmpOf {
+  val command: String = ":asmp"
+}
+
 case class JavapOf(args: String) extends Command
 object JavapOf {
   val command: String = ":javap"
@@ -106,6 +111,7 @@ case object Help extends Command {
       |
       |:help                    print this summary
       |:load <path>             interpret lines in a file
+      |:asmp <path|class>       disassemble a file or class name
       |:javap <path|class>      disassemble a file or class name
       |:quit                    exit the interpreter
       |:type <expression>       evaluate the type of the given expression
@@ -135,6 +141,7 @@ object ParseResult {
     Load.command -> (arg => Load(arg)),
     TypeOf.command -> (arg => TypeOf(arg)),
     DocOf.command -> (arg => DocOf(arg)),
+    AsmpOf.command -> (arg => AsmpOf(arg)),
     JavapOf.command -> (arg => JavapOf(arg))
   )
 

compiler/src/dotty/tools/repl/ReplDriver.scala

@@ -394,6 +394,12 @@ class ReplDriver(settings: Array[String],
         state
       }
 
+    case AsmpOf(line) =>
+      given DisassemblerRepl = DisassemblerRepl(this, state)
+      val opts = AsmpOptions.parse(ReplStrings.words(line))
+      disassemble(Asmp, opts)
+      state
+
     case JavapOf(line) =>
       given DisassemblerRepl = DisassemblerRepl(this, state)
       val opts = JavapOptions.parse(ReplStrings.words(line))

compiler/test/dotty/tools/repl/DisassemblerTests.scala

@@ -235,6 +235,56 @@ class JavapTests extends DisassemblerTest:
     }
 end JavapTests
 
+class AsmpTests extends DisassemblerTest:
+  val packageSeparator = "/"
+
+  @Test def `simple end-to-end` =
+    eval("class Foo1").andThen { implicit state =>
+      run(":asmp -c Foo1")
+      assertDisassemblyIncludes(List(
+        s"public class ${line(1, "Foo1")} {",
+        "public <init>()V",
+        "INVOKESPECIAL java/lang/Object.<init> ()V",
+      ))
+    }
+
+  @Test def `multiple classes in prev entry` =
+    eval {
+      """class Foo2
+        |trait Bar2
+        |""".stripMargin
+    } andThen { implicit state =>
+      run(":asmp -c -")
+      assertDisassemblyIncludes(List(
+        s"public class ${line(1, "Foo2")} {",
+        s"public abstract interface ${line(1, "Bar2")} {",
+      ))
+    }
+
+  @Test def `private selected method` =
+    eval {
+      """class Baz1:
+        |  private def one = 1
+        |  private def two = 2
+        |""".stripMargin
+    } andThen { implicit state =>
+      run(":asmp -p -c Baz1#one")
+      val out = storedOutput()
+      assertDisassemblyIncludes("private one()I", out)
+      assertDisassemblyExcludes("private two()I", out)
+    }
+
+  @Test def `java.lang.String signatures` =
+    fromInitialState { implicit state =>
+      run(":asmp -s java.lang.String")
+      val out = storedOutput()
+      assertDisassemblyIncludes("public static varargs format(Ljava/lang/String;[Ljava/lang/Object;)Ljava/lang/String;", out)
+      assertDisassemblyIncludes("public static join(Ljava/lang/CharSequence;Ljava/lang/Iterable;)Ljava/lang/String;", out)
+      assertDisassemblyIncludes("public concat(Ljava/lang/String;)Ljava/lang/String;", out)
+      assertDisassemblyIncludes("public trim()Ljava/lang/String;", out)
+    }
+end AsmpTests
+
 // Test option parsing
 class JavapOptionTests extends ReplTest:
   private def assertFlags(expected: Seq[String], input: Seq[String])(implicit s: State): Unit =