TASTy macro annotation (part 1)

Add basic support for TASTy annotations.

* Introduce experimental `scala.annotations.TastyAnnotation`
* Macro annotations can analyze or modify definitions
* Macro annotation can add definition around the annotated definition
  * Added members are not visible while typing
  * Added members are not visible to other macro annotations
  * Added definition must have the same owner
* Implement macro annotation expansion
  * Implemented at Inlining phase
    * Can use macro annotations in staged expressions (expanded when at stage 0)
    * Can use staged expression to implement macro annotations
    * Can insert calls to inline methods in macro annotations
  * Current limitations (to be loosened)
    * Can only be used on `def`, `val`, `lazy val` and `var`
    * Can only add `def`, `val`, `lazy val` and `var` definitions

Based on:
  * #15626
  * https://infoscience.epfl.ch/record/294615?ln=en

Author: nicolasstucki

compiler/src/dotty/tools/dotc/CompilationUnit.scala

@@ -16,6 +16,7 @@ import core.Decorators._
 import config.{SourceVersion, Feature}
 import StdNames.nme
 import scala.annotation.internal.sharable
+import transform.TastyAnnotations
 
 class CompilationUnit protected (val source: SourceFile) {
 
@@ -45,6 +46,8 @@ class CompilationUnit protected (val source: SourceFile) {
    */
   var needsInlining: Boolean = false
 
+  var hasTastyAnnotations: Boolean = false
+
   /** Set to `true` if inliner added anonymous mirrors that need to be completed */
   var needsMirrorSupport: Boolean = false
 
@@ -119,6 +122,7 @@ object CompilationUnit {
       force.traverse(unit1.tpdTree)
       unit1.needsStaging = force.containsQuote
       unit1.needsInlining = force.containsInline
+      unit1.hasTastyAnnotations = force.containsTastyAnnotation
     }
     unit1
   }
@@ -147,6 +151,7 @@ object CompilationUnit {
     var containsQuote = false
     var containsInline = false
     var containsCaptureChecking = false
+    var containsTastyAnnotation = false
     def traverse(tree: Tree)(using Context): Unit = {
       if (tree.symbol.isQuote)
         containsQuote = true
@@ -160,6 +165,9 @@ object CompilationUnit {
                 Feature.handleGlobalLanguageImport(prefix, imported)
             case _ =>
         case _ =>
+      for annot <- tree.symbol.annotations do
+        if TastyAnnotations.isTastyAnnotation(annot) then
+          ctx.compilationUnit.hasTastyAnnotations = true
       traverseChildren(tree)
     }
   }

compiler/src/dotty/tools/dotc/config/Printers.scala

@@ -32,6 +32,7 @@ object Printers {
   val init = noPrinter
   val inlining = noPrinter
   val interactiv = noPrinter
+  val macroAnnot = noPrinter
   val matchTypes = noPrinter
   val nullables = noPrinter
   val overload = noPrinter

compiler/src/dotty/tools/dotc/core/Definitions.scala

@@ -895,6 +895,8 @@ class Definitions {
   @tu lazy val QuotedTypeModule: Symbol = QuotedTypeClass.companionModule
     @tu lazy val QuotedTypeModule_of: Symbol = QuotedTypeModule.requiredMethod("of")
 
+  @tu lazy val TastyAnnotationClass: ClassSymbol = requiredClass("scala.annotation.TastyAnnotation")
+
   @tu lazy val CanEqualClass: ClassSymbol = getClassIfDefined("scala.Eql").orElse(requiredClass("scala.CanEqual")).asClass
     def CanEqual_canEqualAny(using Context): TermSymbol =
       val methodName = if CanEqualClass.name == tpnme.Eql then nme.eqlAny else nme.canEqualAny

compiler/src/dotty/tools/dotc/quoted/Interpreter.scala

@@ -32,7 +32,7 @@ import dotty.tools.dotc.reporting.Message
 import dotty.tools.repl.AbstractFileClassLoader
 
 /** Tree interpreter for metaprogramming constructs */
-abstract class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context):
+class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context):
   import Interpreter._
   import tpd._
 
@@ -65,7 +65,7 @@ abstract class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context)
 
     // TODO disallow interpreted method calls as arguments
     case Call(fn, args) =>
-      if (fn.symbol.isConstructor && fn.symbol.owner.owner.is(Package))
+      if (fn.symbol.isConstructor)
         interpretNew(fn.symbol, args.flatten.map(interpretTree))
       else if (fn.symbol.is(Module))
         interpretModuleAccess(fn.symbol)
@@ -185,7 +185,8 @@ abstract class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context)
     loadModule(fn.moduleClass)
 
   private def interpretNew(fn: Symbol, args: => List[Object])(implicit env: Env): Object = {
-    val clazz = loadClass(fn.owner.fullName.toString)
+    val className = fn.owner.fullName.mangledString.replaceAll("\\$\\.", "\\$")
+    val clazz = loadClass(className)
     val constr = clazz.getConstructor(paramsSig(fn): _*)
     constr.newInstance(args: _*).asInstanceOf[Object]
   }
@@ -218,10 +219,6 @@ abstract class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context)
   private def loadClass(name: String): Class[?] =
     try classLoader.loadClass(name)
     catch {
-      case _: ClassNotFoundException if ctx.compilationUnit.isSuspendable  =>
-        if (ctx.settings.XprintSuspension.value)
-          report.echo(i"suspension triggered by a dependency on $name", pos)
-        ctx.compilationUnit.suspend()
       case MissingClassDefinedInCurrentRun(sym) if ctx.compilationUnit.isSuspendable =>
         if (ctx.settings.XprintSuspension.value)
           report.echo(i"suspension triggered by a dependency on $sym", pos)
@@ -276,13 +273,15 @@ abstract class Interpreter(pos: SrcPos, classLoader: ClassLoader)(using Context)
     }
 
   private object MissingClassDefinedInCurrentRun {
-    def unapply(targetException: NoClassDefFoundError)(using Context): Option[Symbol] = {
-      val className = targetException.getMessage
-      if (className eq null) None
-      else {
-        val sym = staticRef(className.toTypeName).symbol
-        if (sym.isDefinedInCurrentRun) Some(sym) else None
-      }
+    def unapply(targetException: Throwable)(using Context): Option[Symbol] = {
+      targetException match
+        case _: NoClassDefFoundError | _: ClassNotFoundException =>
+          val className = targetException.getMessage
+          if className eq null then None
+          else
+            val sym = staticRef(className.toTypeName).symbol
+            if (sym.isDefinedInCurrentRun) Some(sym) else None
+        case _ => None
     }
   }
 

compiler/src/dotty/tools/dotc/transform/Inlining.scala

@@ -7,14 +7,18 @@ import Contexts._
 import Symbols._
 import SymUtils._
 import dotty.tools.dotc.ast.tpd
-
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.quoted._
 import dotty.tools.dotc.core.StagingContext._
 import dotty.tools.dotc.inlines.Inlines
 import dotty.tools.dotc.ast.TreeMapWithImplicits
+import dotty.tools.dotc.core.DenotTransformers.IdentityDenotTransformer
 
 
 /** Inlines all calls to inline methods that are not in an inline method or a quote */
-class Inlining extends MacroTransform {
+class Inlining extends MacroTransform with IdentityDenotTransformer {
+  thisPhase =>
+
   import tpd._
 
   override def phaseName: String = Inlining.name
@@ -23,8 +27,10 @@ class Inlining extends MacroTransform {
 
   override def allowsImplicitSearch: Boolean = true
 
+  override def changesMembers: Boolean = true
+
   override def run(using Context): Unit =
-    if ctx.compilationUnit.needsInlining then
+    if ctx.compilationUnit.needsInlining || ctx.compilationUnit.hasTastyAnnotations then
       try super.run
       catch case _: CompilationUnit.SuspendException => ()
 
@@ -59,8 +65,16 @@ class Inlining extends MacroTransform {
   private class InliningTreeMap extends TreeMapWithImplicits {
     override def transform(tree: Tree)(using Context): Tree = {
       tree match
-        case tree: DefTree =>
+        case tree: MemberDef =>
           if tree.symbol.is(Inline) then tree
+          else if tree.symbol.is(Param) then super.transform(tree)
+          else if
+            !tree.symbol.isPrimaryConstructor
+            && StagingContext.level == 0
+            && TastyAnnotations.hasTastyAnnotation(tree.symbol)
+          then
+            val trees = new TastyAnnotations(thisPhase).expandAnnotations(tree)
+            flatTree(trees.map(super.transform))
           else super.transform(tree)
         case _: Typed | _: Block =>
           super.transform(tree)

compiler/src/dotty/tools/dotc/transform/PostTyper.scala

@@ -375,21 +375,25 @@ class PostTyper extends MacroTransform with IdentityDenotTransformer { thisPhase
             )
           }
         case tree: ValDef =>
+          registerIfHasTastyAnnotations(tree)
           checkErasedDef(tree)
           val tree1 = cpy.ValDef(tree)(rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           if tree1.removeAttachment(desugar.UntupledParam).isDefined then
             checkStableSelection(tree.rhs)
           processValOrDefDef(super.transform(tree1))
         case tree: DefDef =>
+          registerIfHasTastyAnnotations(tree)
           checkErasedDef(tree)
           annotateContextResults(tree)
           val tree1 = cpy.DefDef(tree)(rhs = normalizeErasedRhs(tree.rhs, tree.symbol))
           processValOrDefDef(superAcc.wrapDefDef(tree1)(super.transform(tree1).asInstanceOf[DefDef]))
         case tree: TypeDef =>
+          registerIfHasTastyAnnotations(tree)
           val sym = tree.symbol
           if (sym.isClass)
             VarianceChecker.check(tree)
             annotateExperimental(sym)
+            checkTastyAnnotation(sym)
             tree.rhs match
               case impl: Template =>
                 for parent <- impl.parents do
@@ -483,6 +487,16 @@ class PostTyper extends MacroTransform with IdentityDenotTransformer { thisPhase
     private def normalizeErasedRhs(rhs: Tree, sym: Symbol)(using Context) =
       if (sym.isEffectivelyErased) dropInlines.transform(rhs) else rhs
 
+    /** Check if the definition has macro annotation and sets `compilationUnit.hasTastyAnnotations` if needed. */
+    private def registerIfHasTastyAnnotations(tree: DefTree)(using Context) =
+      if !Inlines.inInlineMethod && TastyAnnotations.hasTastyAnnotation(tree.symbol) then
+        ctx.compilationUnit.hasTastyAnnotations = true
+
+    /** Check macro annotations implementations  */
+    private def checkTastyAnnotation(sym: Symbol)(using Context) =
+      if sym.derivesFrom(defn.TastyAnnotationClass) && !sym.isStatic then
+        report.error("classes that extend TastyAnnotation must not be inner/local classes", sym.srcPos)
+
     private def checkErasedDef(tree: ValOrDefDef)(using Context): Unit =
       if tree.symbol.is(Erased, butNot = Macro) then
         val tpe = tree.rhs.tpe

compiler/src/dotty/tools/dotc/transform/TastyAnnotations.scala

@@ -0,0 +1,117 @@
+package dotty.tools.dotc
+package transform
+
+import scala.language.unsafeNulls
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.ast.Trees.*
+import dotty.tools.dotc.config.Printers.{macroAnnot => debug}
+import dotty.tools.dotc.core.Annotations.*
+import dotty.tools.dotc.core.Contexts.*
+import dotty.tools.dotc.core.Decorators.*
+import dotty.tools.dotc.core.DenotTransformers.DenotTransformer
+import dotty.tools.dotc.core.Flags.*
+import dotty.tools.dotc.core.MacroClassLoader
+import dotty.tools.dotc.core.Symbols.*
+import dotty.tools.dotc.quoted.*
+import dotty.tools.dotc.util.SrcPos
+import scala.quoted.runtime.impl.{QuotesImpl, SpliceScope}
+
+import scala.quoted.Quotes
+
+class TastyAnnotations(thisPhase: DenotTransformer):
+  import tpd.*
+  import TastyAnnotations.*
+
+  /** Expands every TASTy annotation that is on this tree.
+   *  Returns a list with transformed definition and any added definitions.
+   */
+  def expandAnnotations(tree: MemberDef)(using Context): List[DefTree] =
+    if !hasTastyAnnotation(tree.symbol) then
+      List(tree)
+    else if tree.symbol.is(Module) then
+      if tree.symbol.isClass then // error only reported on module class
+        report.error("TASTy annotations are not supported on object", tree)
+      List(tree)
+    else if tree.symbol.isClass then
+      report.error("TASTy annotations are not supported on class", tree)
+      List(tree)
+    else if tree.symbol.isType then
+      report.error("TASTy annotations are not supported on type", tree)
+      List(tree)
+    else
+      debug.println(i"Expanding TASTy annotations of:\n$tree")
+
+      val macroInterpreter = new Interpreter(tree.srcPos, MacroClassLoader.fromContext)
+
+      val allTrees = List.newBuilder[DefTree]
+      var insertedAfter: List[List[DefTree]] = Nil
+
+      // Apply all TASTy annotation to `tree` and collect new definitions in order
+      val transformedTree: DefTree = tree.symbol.annotations.foldLeft(tree) { (tree, annot) =>
+        if isTastyAnnotation(annot) then
+          debug.println(i"Expanding TASTy annotation: ${annot}")
+
+          // Interpret call to `new myAnnot(..).transform(using <Quotes>)(<tree>)`
+          val transformedTrees = callMacro(macroInterpreter, tree, annot)
+          transformedTrees.span(_.symbol != tree.symbol) match
+            case (prefixed, newTree :: suffixed) =>
+              allTrees ++= prefixed
+              insertedAfter = suffixed :: insertedAfter
+              prefixed.foreach(checkAndEnter(_, tree.symbol, annot))
+              suffixed.foreach(checkAndEnter(_, tree.symbol, annot))
+              newTree
+            case (Nil, Nil) =>
+              report.error(i"Unexpected `Nil` returned by `(${annot.tree}).transform(..)` during macro expansion", annot.tree.srcPos)
+              tree
+            case (_, Nil) =>
+              report.error(i"Transformed tree for ${tree} was not return by `(${annot.tree}).transform(..)` during macro expansion", annot.tree.srcPos)
+              tree
+        else
+          tree
+      }
+
+      allTrees += transformedTree
+      insertedAfter.foreach(allTrees.++=)
+
+      val result = allTrees.result()
+      debug.println(result.map(_.show).mkString("expanded to:\n", "\n", ""))
+      result
+
+  /** Interpret the code `new annot(..).transform(using <Quotes(ctx)>)(<tree>)` */
+  private def callMacro(interpreter: Interpreter, tree: MemberDef, annot: Annotation)(using Context): List[MemberDef] =
+    // TODO: Remove when scala.annaotaion.TastyAnnotation is no longer experimental
+    import scala.reflect.Selectable.reflectiveSelectable
+    type TastyAnnotation = {
+      def transform(using Quotes)(tree: Object/*Erased type of quotes.refelct.Definition*/): List[MemberDef /*quotes.refelct.Definition known to be MemberDef in QuotesImpl*/]
+    }
+
+    // Interpret TASTy annotation instantiation `new myAnnot(..)`
+    val annotInstance = interpreter.interpret[TastyAnnotation](annot.tree).get
+    // TODO: Remove when scala.annaotaion.TastyAnnotation is no longer experimental
+    assert(annotInstance.getClass.getClassLoader.loadClass("scala.annotation.TastyAnnotation").isInstance(annotInstance))
+
+    val quotes = QuotesImpl()(using SpliceScope.contextWithNewSpliceScope(tree.symbol.sourcePos)(using MacroExpansion.context(tree)).withOwner(tree.symbol))
+    annotInstance.transform(using quotes)(tree.asInstanceOf[quotes.reflect.Definition])
+
+  /** Check that this tree can be added by the TASTy annotation and enter it if needed */
+  private def checkAndEnter(newTree: Tree, annotated: Symbol, annot: Annotation)(using Context) =
+    val sym = newTree.symbol
+    if sym.isClass then
+      report.error("Generating classes is not supported", annot.tree)
+    else if sym.isType then
+      report.error("Generating type is not supported", annot.tree)
+    else if sym.owner != annotated.owner then
+      report.error(i"TASTy annotation $annot added $sym with an inconsistent owner. Expected it to be owned by ${annotated.owner} but was owned by ${sym.owner}.", annot.tree)
+    else
+      sym.enteredAfter(thisPhase)
+
+object TastyAnnotations:
+
+  /** Is this an annotation that implements `scala.annation.TastyAnnotation` */
+  def isTastyAnnotation(annot: Annotation)(using Context): Boolean =
+    annot.tree.symbol.maybeOwner.derivesFrom(defn.TastyAnnotationClass)
+
+  /** Is this symbol annotated with an annotation that implements `scala.annation.TastyAnnotation` */
+  def hasTastyAnnotation(sym: Symbol)(using Context): Boolean =
+    sym.getAnnotation(defn.TastyAnnotationClass).isDefined

compiler/src/dotty/tools/dotc/transform/YCheckPositions.scala

@@ -61,6 +61,7 @@ class YCheckPositions extends Phase {
 
   private def isMacro(call: Tree)(using Context) =
     call.symbol.is(Macro) ||
+    (call.symbol.isClass && call.tpe.derivesFrom(defn.TastyAnnotationClass)) ||
     // The call of a macro after typer is encoded as a Select while other inlines are Ident
     // TODO remove this distinction once Inline nodes of expanded macros can be trusted (also in Inliner.inlineCallTrace)
     (!(ctx.phase <= postTyperPhase) && call.isInstanceOf[Select])