Add rewrite prototype and couple of fixes

Author: nicolasstucki

compiler/src/dotty/tools/dotc/tastyreflect/ReflectionCompilerInterface.scala

@@ -204,7 +204,7 @@ class ReflectionCompilerInterface(val rootContext: core.Contexts.Context) extend
   def DefDef_apply(symbol: Symbol, rhsFn: List[Type] => List[List[Term]] => Option[Term])(given Context): DefDef =
     withDefaultPos(tpd.polyDefDef(symbol.asTerm, tparams => vparamss => rhsFn(tparams)(vparamss).getOrElse(tpd.EmptyTree)))
 
-  def DefDef_copy(original: DefDef)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given Context): DefDef =
+  def DefDef_copy(original: Tree)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given Context): DefDef =
     tpd.cpy.DefDef(original)(name.toTermName, typeParams, paramss, tpt, rhs.getOrElse(tpd.EmptyTree))
 
   type ValDef = tpd.ValDef

library/src/scala/quoted/Expr.scala

@@ -190,6 +190,11 @@ package internal {
      *  May contain references to code defined outside this TastyTreeExpr instance.
      */
     final class TastyTreeExpr[Tree](val tree: Tree, val scopeId: Int) extends Expr[Any] {
+      override def equals(that: Any): Boolean = that match {
+        case that: TastyTreeExpr[_] => tree == that.tree && scopeId == that.scopeId
+        case _ => false
+      }
+      override def hashCode: Int = tree.hashCode
       override def toString: String = s"Expr(<tasty tree>)"
     }
 

library/src/scala/quoted/Type.scala

@@ -71,6 +71,11 @@ package internal {
 
     /** An Type backed by a tree */
     final class TreeType[Tree](val typeTree: Tree, val scopeId: Int) extends scala.quoted.Type[Any] {
+      override def equals(that: Any): Boolean = that match {
+        case that: TreeType[_] => typeTree == that.typeTree && scopeId == that.scopeId
+        case _ => false
+      }
+      override def hashCode: Int = typeTree.hashCode
       override def toString: String = s"Type(<tasty tree>)"
     }
 

library/src/scala/tasty/reflect/CompilerInterface.scala

@@ -262,7 +262,7 @@ trait CompilerInterface {
   def DefDef_rhs(self: DefDef)(given ctx: Context): Option[Term]
 
   def DefDef_apply(symbol: Symbol, rhsFn: List[Type] => List[List[Term]] => Option[Term])(given ctx: Context): DefDef
-  def DefDef_copy(original: DefDef)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given ctx: Context): DefDef
+  def DefDef_copy(original: Tree)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given ctx: Context): DefDef
 
   /** Tree representing a value definition in the source code This inclues `val`, `lazy val`, `var`, `object` and parameter definitions. */
   type ValDef <: Definition

library/src/scala/tasty/reflect/SourceCodePrinter.scala

@@ -514,6 +514,9 @@ class SourceCodePrinter[R <: Reflection & Singleton](val tasty: R)(syntaxHighlig
       case IsTypeTree(tpt) =>
         printTypeTree(tpt)
 
+      case Closure(meth, _) =>
+        printTree(meth)
+
       case _ =>
         throw new MatchError(tree.showExtractors)
 

library/src/scala/tasty/reflect/TreeOps.scala

@@ -97,7 +97,7 @@ trait TreeOps extends Core {
   object DefDef {
     def apply(symbol: Symbol, rhsFn: List[Type] => List[List[Term]] => Option[Term])(given ctx: Context): DefDef =
       internal.DefDef_apply(symbol, rhsFn)
-    def copy(original: DefDef)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given ctx: Context): DefDef =
+    def copy(original: Tree)(name: String, typeParams: List[TypeDef], paramss: List[List[ValDef]], tpt: TypeTree, rhs: Option[Term])(given ctx: Context): DefDef =
       internal.DefDef_copy(original)(name, typeParams, paramss, tpt, rhs)
     def unapply(tree: Tree)(given ctx: Context): Option[(String, List[TypeDef], List[List[ValDef]], TypeTree, Option[Term])] =
       internal.matchDefDef(tree).map(x => (x.name, x.typeParams, x.paramss, x.returnTpt, x.rhs))

tests/run-macros/flops-rewrite.check

@@ -0,0 +1,9 @@
+scala.Nil.map[scala.Nothing](((x: scala.Nothing) => x))
+scala.Nil
+
+scala.Nil.map[scala.Nothing](((x: scala.Nothing) => x)).++[scala.Nothing](scala.Nil.map[scala.Nothing](((x: scala.Nothing) => x)))
+scala.Nil
+
+scala.Nil.map[scala.Nothing](((x: scala.Nothing) => x)).++[scala.Int](scala.List.apply[scala.Int]((3: scala.<repeated>[scala.Int]))).++[scala.Int](scala.Nil)
+scala.List.apply[scala.Int]((3: scala.<repeated>[scala.Int]))
+

tests/run-macros/flops-rewrite/Macro_1.scala

@@ -0,0 +1,81 @@
+import scala.quoted._
+
+inline def rewrite[T](x: => T): T = ${ rewriteMacro('x) }
+
+private def rewriteMacro[T: Type](x: Expr[T])(given QuoteContext): Expr[T] = {
+  val x2 = Rewriter(
+    postTransform = {
+      case '{ Nil.map[$t]($f) } => '{ Nil }
+      case '{ Nil.filter($f) }  => '{ Nil }
+      case '{ Nil.++[$t]($xs) } => xs
+      case '{ ($xs: List[$t]).++(Nil) } => xs
+      case x => x
+    }
+  ).rewrite(x)
+
+  '{
+    println(${Expr(x.show)})
+    println(${Expr(x2.show)})
+    println()
+    $x2
+  }
+}
+
+private object Rewriter {
+  def apply(preTransform: Expr[Any] => Expr[Any] = identity, postTransform: Expr[Any] => Expr[Any] = identity, fixPoint: Boolean = false): Rewriter =
+    new Rewriter(preTransform, postTransform, fixPoint)
+}
+
+private class Rewriter(preTransform: Expr[Any] => Expr[Any], postTransform: Expr[Any] => Expr[Any], fixPoint: Boolean) {
+  def rewrite[T](e: Expr[T])(given QuoteContext, Type[T]): Expr[T] = {
+    val e2 = checkedTransform(e, preTransform)
+    val e3 = rewriteChildren(e2)
+    val e4 = checkedTransform(e3, postTransform)
+    if fixPoint && e4 != e then rewrite(e4)
+    else e4
+  }
+
+  private def checkedTransform[T: Type](e: Expr[T], transform: Expr[T] => Expr[Any])(given QuoteContext): Expr[T] = {
+    transform(e) match {
+      case '{ $x: T } => x
+      case '{ $x: $t } => throw new Exception(
+        s"""Transformed
+           |${e.show}
+           |into
+           |${x.show}
+           |
+           |Expected type to be
+           |${summon[Type[T]].show}
+           |but was
+           |${t.show}
+         """.stripMargin)
+    }
+  }
+
+  def rewriteChildren[T: Type](e: Expr[T])(given qctx: QuoteContext): Expr[T] = {
+    import qctx.tasty.{_, given}
+    class MapChildren extends TreeMap {
+      override def transformTerm(tree: Term)(given ctx: Context): Term = tree match {
+        case IsClosure(_) =>
+          tree
+        case IsInlined(_) | IsSelect(_) =>
+          transformChildrenTerm(tree)
+        case _ =>
+          tree.tpe match {
+            case IsMethodType(_) | IsPolyType(_) =>
+              transformChildrenTerm(tree)
+            case _ =>
+              tree.seal match {
+                case '{ $x: $t } => rewrite(x).unseal
+              }
+          }
+      }
+      def transformChildrenTerm(tree: Term)(given ctx: Context): Term =
+        super.transformTerm(tree)
+    }
+    (new MapChildren).transformChildrenTerm(e.unseal).seal.cast[T] // Cast will only fail if this implementation has a bug
+  }
+
+}
+
+

tests/run-macros/flops-rewrite/Test_2.scala

@@ -0,0 +1,9 @@
+object Test {
+
+  def main(args: Array[String]): Unit = {
+    rewrite(Nil.map(x => x))
+    rewrite(Nil.map(x => x) ++ Nil.map(x => x))
+    rewrite(Nil.map(x => x) ++ List(3) ++ Nil)
+  }
+
+}