Scala 3 macros trait arguments

Viewed 193

Here is the basic setup:

trait MyProduct[A, B](a: A, b:  B)
class X
class Y

class XxY(x: X, y: Y) extends MyProduct(x,y)

I'm trying to examine the arguments to the MyProduct trait. More specifically I want to extract some information, such as the strings "x" and "y", that indicates which fields of XxY are passed to MyProduct.

Importantly, I need to satisfy the case with multiple fields of the same type, eg class XxYxY(x: X, y1: Y, y2: Y) extends MyProduct3(x, y1, y2), so reasoning from the type arguments isn't sufficient.

I think maybe the issue is I haven't found a way to get a symbol to the extends clause itself. I can find the ClassDef for XxY, from which I can extract parents and get the already constructed type MyProduct[X,Y]. I've also looked at the symbol.declarations of the enclosing module and the XxY.<init> function to see if there might be data that I can use to find the arguments, but none of these seem to have the information I'm looking for.

Looking at the tree for the enclosing module makes me think that perhaps this information is erased and instead needs to be parsed as text from the source code, but I'm hoping someone has a better solution.

EDIT from comment:

As input I have a Type instance of the enclosing object/module. For example:

inline def simpleProductTest[T <: Singleton]: String = ${simpleProductTestImpl[T]}

def simpleProductTestImpl[T](using q: Quotes)(using Type[T]): Expr[String] = ???

As output I would like something that allows me to examine the arguments to the trait used in the extends clause.

1 Answers

If you do

import scala.quoted.*

object Macros {
  inline def simpleProductTest[T]: Unit = ${simpleProductTestImpl[T]}

  def simpleProductTestImpl[T](using Quotes)(using Type[T]): Expr[Unit] = {
    import quotes.reflect.*
    println(TypeRepr.of[T].typeSymbol.tree.asInstanceOf[TypeDef].rhs)
    '{()}
  }
}

you'll see

object App {
  trait MyProduct[A, B](a: A, b:  B)
  class X
  class Y

  class XxY(x: X, y: Y, z: Int) extends MyProduct(x,y)

  Macros.simpleProductTest[XxY]
}

that there is no information about x,y

Template(
  DefDef(
    <init>,
    List(List(
      ValDef(x,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class X)],EmptyTree),
      ValDef(y,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class Y)],EmptyTree),
      ValDef(z,TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class <root>)),object scala),class Int)],EmptyTree)
    )),
    TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class XxY)],
    EmptyTree
  ),
  List(
    TypeTree[TypeRef(ThisType(TypeRef(NoPrefix,module class lang)),class Object)],
    TypeTree[AppliedType(TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),trait MyProduct),
      List(
        TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class X),
        TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class Y)
      )
    )]
  ),
  ValDef(_,EmptyTree,EmptyTree),
  List(
    ValDef(x,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class X)],EmptyTree),
    ValDef(y,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class Y)],EmptyTree),
    ValDef(z,TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class <root>)),object scala),class Int)],EmptyTree)
  )
)

And Template doesn't exist in scala.quoted.*, it's dotty.tools.dotc.ast.Trees.Template.

But if you print .tasty file with TASTy inspection

libraryDependencies += "org.scala-lang" %% "scala3-tasty-inspector" % "3.0.0"

import scala.quoted.*
import scala.tasty.inspector.*

class MyInspector extends Inspector:
  def inspect(using Quotes)(tastys: List[Tasty[quotes.type]]): Unit =
    import quotes.reflect.*
    for tasty <- tastys do
      val tree = tasty.ast
      println(tree)

@main def test: Unit =
  val tastyFiles = List("target/scala-3.0.0/classes/App.tasty")
  TastyInspector.inspectTastyFiles(tastyFiles)(new MyInspector)

you'll see

Template(
  DefDef(
    <init>,
    List(List(
      ValDef(x,Ident(X),EmptyTree),
      ValDef(y,Ident(Y),EmptyTree),
      ValDef(z,Ident(Int),EmptyTree)
    )),
    TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class <root>)),object scala),Unit)],
    EmptyTree
  ),
  List(
    Apply(Select(New(TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class java)),object lang),Object)]),<init>),List()),
    Apply(
      TypeApply(
        Select(New(Ident(MyProduct)),<init>),
        List(
          TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <root>)),module class <empty>)),module class App$)),class X)],
          TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <root>)),module class <empty>)),module class App$)),class Y)]
        )
      ),
      List(Ident(x), Ident(y))  <--- HERE!!!
    )
  ),
  ValDef(_,EmptyTree,EmptyTree),
  List(
    ValDef(x,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <root>)),module class <empty>)),module class App$)),class X)],EmptyTree),
    ValDef(y,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <root>)),module class <empty>)),module class App$)),class Y)],EmptyTree),
    ValDef(z,TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class <root>)),object scala),Int)],EmptyTree)
  )
)

that the tree is slightly different and information about x,y is there.

In Scaladocs for extension method .tree of Symbol it's written

     *  **Warning**: avoid using this method in macros.
     *
     *  **Caveat**: The tree is not guaranteed to exist unless the compiler
     *  option `-Yretain-trees` is enabled.

Let's switch on -Yretain-trees

scalacOptions += "-Yretain-trees"

Then Macros.simpleProductTest[XxY] prints

Template(
  DefDef(
    <init>,
    List(List(
      ValDef(x,Ident(X),EmptyTree),
      ValDef(y,Ident(Y),EmptyTree),
      ValDef(z,Ident(Int),EmptyTree)
    )),
    TypeTree[TypeRef(ThisType(TypeRef(NoPrefix,module class scala)),class Unit)],
    EmptyTree
  ),
  List(
    Apply(Select(New(TypeTree[TypeRef(ThisType(TypeRef(NoPrefix,module class lang)),class Object)]),<init>),List()),
    Apply(
      TypeApply(
        Select(New(Ident(MyProduct)),<init>),
        List(
          TypeTree[TypeVar(TypeParamRef(A) -> TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class X))],
          TypeTree[TypeVar(TypeParamRef(B) -> TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class Y))]
        )
      ),
      List(Ident(x), Ident(y))   <----  AGAIN HERE !!!!
    )
  ),
  ValDef(_,EmptyTree,EmptyTree),
  List(
    ValDef(x,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class X)],EmptyTree),
    ValDef(y,TypeTree[TypeRef(ThisType(TypeRef(ThisType(TypeRef(NoPrefix,module class <empty>)),module class App$)),class Y)],EmptyTree),
    ValDef(z,TypeTree[TypeRef(TermRef(ThisType(TypeRef(NoPrefix,module class <root>)),object scala),class Int)],EmptyTree)
  )
)

i.e. tree is again slightly different but information about x,y appeared. So the question is now only how to extract this information since it's dotty.tools.dotc.ast.*

case class Template[-T >: Untyped] private[ast](
  constr: DefDef[T], 
  parentsOrDerived: List[Tree[T]], 
  self: ValDef[T], 
  private var preBody: LazyTreeList[T @uncheckedVariance]
)(implicit @constructorOnly src: SourceFile)

So finally

libraryDependencies += "org.scala-lang" %% "scala3-compiler" % "3.0.0"

import scala.quoted.*

inline def simpleProductTest[T]: List[String] = ${simpleProductTestImpl[T]}

def simpleProductTestImpl[T](using Quotes)(using Type[T]): Expr[List[String]] = {
  import quotes.reflect.*
  Expr(
    TypeRepr.of[T].typeSymbol.tree.asInstanceOf[TypeDef].rhs
      .asInstanceOf[dotty.tools.dotc.ast.Trees.Template[?]]
      .parentsOrDerived
      .asInstanceOf[List[Apply]]
      .tail.head.args.map(_.symbol.name)
  )
}

Macros.simpleProductTest[XxY] // List(x, y)
Related