Enable returning classes from MacroAnnotations (part 3)

Enable the addition of classes from a `MacroAnnotation`:
 * Can add new `class` definitions next to the annotated definition

Special cases:
 * An annotated top-level `def`, `val`, `var`, `lazy val` can return a `class`
   definition that is owned by the package or package object.

Related PRs:
 * Follows #16454

Author: nicolasstucki

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

@@ -63,6 +63,12 @@ class Inlining extends MacroTransform with IdentityDenotTransformer {
   }
 
   private class InliningTreeMap extends TreeMapWithImplicits {
+
+    /** List of top level classes added by macro annotation in a package object.
+     *  These are added the PackageDef that owns this particular package object.
+     */
+    private val topClasses = new collection.mutable.ListBuffer[Tree]
+
     override def transform(tree: Tree)(using Context): Tree = {
       tree match
         case tree: MemberDef =>
@@ -74,7 +80,15 @@ class Inlining extends MacroTransform with IdentityDenotTransformer {
             && MacroAnnotations.hasMacroAnnotation(tree.symbol)
           then
             val trees = new MacroAnnotations(thisPhase).expandAnnotations(tree)
-            flatTree(trees.map(super.transform))
+            val trees1 = trees.map(super.transform)
+
+            // Find classes added to the top level from a package object
+            val (topClasses0, trees2) =
+              if ctx.owner.isPackageObject then trees1.partition(_.symbol.owner == ctx.owner.owner)
+              else (Nil, trees1)
+            topClasses ++= topClasses0
+
+            flatTree(trees2)
           else super.transform(tree)
         case _: Typed | _: Block =>
           super.transform(tree)
@@ -86,6 +100,14 @@ class Inlining extends MacroTransform with IdentityDenotTransformer {
           super.transform(tree)(using StagingContext.quoteContext)
         case _: GenericApply if tree.symbol.isExprSplice =>
           super.transform(tree)(using StagingContext.spliceContext)
+        case _: PackageDef =>
+          super.transform(tree) match
+            case tree1: PackageDef if !topClasses.isEmpty =>
+              topClasses ++= tree1.stats
+              val newStats = topClasses.result()
+              topClasses.clear()
+              cpy.PackageDef(tree1)(tree1.pid, newStats)
+            case tree1 => tree1
         case _ =>
           super.transform(tree)
     }

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

@@ -119,13 +119,11 @@ class MacroAnnotations(thisPhase: DenotTransformer):
     annotInstance.transform(using quotes)(tree.asInstanceOf[quotes.reflect.Definition])
 
   /** Check that this tree can be added by the macro annotation and enter it if needed */
-  private def checkAndEnter(newTree: Tree, annotated: Symbol, annot: Annotation)(using Context) =
+  private def checkAndEnter(newTree: DefTree, annotated: Symbol, annot: Annotation)(using Context) =
     val sym = newTree.symbol
-    if sym.isClass then
-      report.error(i"macro annotation returning a `class` is not yet supported. $annot tried to add $sym", annot.tree)
-    else if sym.isType then
+    if sym.isType && !sym.isClass then
       report.error(i"macro annotation cannot return a `type`. $annot tried to add $sym", annot.tree)
-    else if sym.owner != annotated.owner then
+    else if sym.owner != annotated.owner && !(annotated.owner.isPackageObject && (sym.isClass || sym.is(Module)) && sym.owner == annotated.owner.owner) then
       report.error(i"macro 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 if annotated.isClass && annotated.owner.is(Package) /*&& !sym.isClass*/ then
       report.error(i"macro annotation can not add top-level ${sym.showKind}. $annot tried to add $sym.", annot.tree)

library/src/scala/quoted/Quotes.scala

@@ -3638,6 +3638,7 @@ trait Quotes { self: runtime.QuoteUnpickler & runtime.QuoteMatching =>
        *  @note As a macro can only splice code into the point at which it is expanded, all generated symbols must be
        *        direct or indirect children of the reflection context's owner.
        */
+      // TODO: add flags and privateWithin
       @experimental def newClass(parent: Symbol, name: String, parents: List[TypeRepr], decls: Symbol => List[Symbol], selfType: Option[TypeRepr]): Symbol
 
       /** Generates a new method symbol with the given parent, name and type.
@@ -4217,7 +4218,7 @@ trait Quotes { self: runtime.QuoteUnpickler & runtime.QuoteMatching =>
     //   FLAGS   //
     ///////////////
 
-    /** FlagSet of a Symbol */
+    /** Flags of a Symbol */
     type Flags
 
     /** Module object of `type Flags`  */

tests/neg-macros/annot-mod-top-method-add-top-method/Macro_1.scala

@@ -0,0 +1,13 @@
+import scala.annotation.{experimental, MacroAnnotation}
+import scala.quoted._
+import scala.collection.mutable
+
+@experimental
+// Assumes annotation is on top level def or val
+class addTopLevelMethodOutsidePackageObject extends MacroAnnotation:
+  def transform(using Quotes)(tree: quotes.reflect.Definition): List[quotes.reflect.Definition] =
+    import quotes.reflect._
+    val methType = MethodType(Nil)(_ => Nil, _ => TypeRepr.of[Int])
+    val methSym = Symbol.newMethod(Symbol.spliceOwner.owner, Symbol.freshName("toLevelMethod"), methType, Flags.EmptyFlags, Symbol.noSymbol)
+    val methDef = ValDef(methSym, Some(Literal(IntConstant(1))))
+    List(methDef, tree)

tests/neg-macros/annot-mod-top-method-add-top-method/Test_2.scala

@@ -0,0 +1,5 @@
+@addTopLevelMethodOutsidePackageObject // error
+def foo = 1
+
+@addTopLevelMethodOutsidePackageObject // error
+val bar = 1

tests/neg-macros/annot-mod-top-method-add-top-val/Macro_1.scala

@@ -0,0 +1,12 @@
+import scala.annotation.{experimental, MacroAnnotation}
+import scala.quoted._
+import scala.collection.mutable
+
+@experimental
+// Assumes annotation is on top level def or val
+class addTopLevelValOutsidePackageObject extends MacroAnnotation:
+  def transform(using Quotes)(tree: quotes.reflect.Definition): List[quotes.reflect.Definition] =
+    import quotes.reflect._
+    val valSym = Symbol.newVal(Symbol.spliceOwner.owner, Symbol.freshName("toLevelVal"), TypeRepr.of[Int], Flags.EmptyFlags, Symbol.noSymbol)
+    val valDef = ValDef(valSym, Some(Literal(IntConstant(1))))
+    List(valDef, tree)

tests/neg-macros/annot-mod-top-method-add-top-val/Test_2.scala

@@ -0,0 +1,5 @@
+@addTopLevelValOutsidePackageObject // error
+def foo = 1
+
+@addTopLevelValOutsidePackageObject // error
+val bar = 1

tests/run-macros/annot-add-global-class.check

@@ -0,0 +1,4 @@
+macro generated main
+executed in: Bar$macro$1
+macro generated main
+executed in: Bar$macro$2

tests/run-macros/annot-add-global-class/Macro_1.scala

@@ -0,0 +1,28 @@
+import scala.annotation.{experimental, MacroAnnotation}
+import scala.quoted._
+import scala.collection.mutable
+
+@experimental
+class addClass extends MacroAnnotation:
+  def transform(using Quotes)(tree: quotes.reflect.Definition): List[quotes.reflect.Definition] =
+    import quotes.reflect._
+    tree match
+      case DefDef(name, List(TermParamClause(Nil)), tpt, Some(rhs)) =>
+        val parents = List(TypeTree.of[Object])
+        def decls(cls: Symbol): List[Symbol] =
+          List(Symbol.newMethod(cls, "run", MethodType(Nil)(_ => Nil, _ => TypeRepr.of[Unit]), Flags.EmptyFlags, Symbol.noSymbol))
+
+        val newClassName = Symbol.freshName("Bar")
+        val cls = Symbol.newClass(Symbol.spliceOwner.owner, newClassName, parents = parents.map(_.tpe), decls, selfType = None)
+        val runSym = cls.declaredMethod("run").head
+
+        val runDef = DefDef(runSym, _ => Some(rhs))
+        val clsDef = ClassDef(cls, parents, body = List(runDef))
+
+        val newCls = Apply(Select(New(TypeIdent(cls)), cls.primaryConstructor), Nil)
+
+        val newDef = DefDef.copy(tree)(name, List(TermParamClause(Nil)), tpt, Some(Apply(Select(newCls, runSym), Nil)))
+        List(clsDef, newDef)
+      case _ =>
+        report.error("Annotation only supports `def` with one argument")
+        List(tree)

tests/run-macros/annot-add-global-class/Test_2.scala

@@ -0,0 +1,25 @@
+@addClass def foo(): Unit =
+  println("macro generated main")
+  println("executed in: " + (new Throwable().getStackTrace().head.getClassName))
+//> class Baz$macro$1 extends Object {
+//>   def run() =
+//>     println("macro generated main")
+//>     println("executed in: " + getClass.getName)
+//> }
+//> def foo(): Unit =
+//>   new Baz$macro$1.run
+
+@addClass def bar(): Unit =
+  println("macro generated main")
+  println("executed in: " + (new Throwable().getStackTrace().head.getClassName))
+//> class Baz$macro$2 extends Object {
+//>   def run() =
+//>     println("macro generated main")
+//>     println("executed in: " + getClass.getName)
+//> }
+//> def foo(): Unit =
+//>   new Baz$macro$2.run
+
+@main def Test(): Unit =
+  foo()
+  bar()