Add a unified prototype

Author: nicolasstucki

tests/run-macros/flops-rewrite-3.check

@@ -0,0 +1,18 @@
+Macro_1$package.plus(1, 4)
+5
+
+Macro_1$package.plus(0, a)
+a
+
+Macro_1$package.plus(a, b)
+a.+(b)
+
+Macro_1$package.plus(Macro_1$package.plus(a, 0), Macro_1$package.plus(0, b))
+0.+(a).+(b)
+
+Macro_1$package.power(4, 5)
+1024
+
+Macro_1$package.power(a, 5)
+a.*(a.*(a.*(a.*(1.*(a)))))
+

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

@@ -0,0 +1,137 @@
+import scala.quoted._
+import scala.quoted.matching._
+
+inline def rewrite[T](x: => T): T = ${ rewriteMacro('x) }
+
+def plus(x: Int, y: Int): Int = x + y
+def times(x: Int, y: Int): Int = x * y
+def power(x: Int, y: Int): Int = if y == 0 then 1 else times(x, power(x, y - 1))
+
+private def rewriteMacro[T: Type](x: Expr[T])(given QuoteContext): Expr[T] = {
+  val rewriter = Rewriter().withFixPoint.withPost(
+    Transformation.safe[Int] {
+      case '{ plus($x, $y) } =>
+        (x, y) match {
+          case (Const(0), _) => y
+          case (Const(a), Const(b)) => Expr(a + b)
+          case (_, Const(_)) =>  '{ $y + $x }
+          case _ => '{ $x + $y }
+        }
+      case '{ times($x, $y) } =>
+        (x, y) match {
+          case (Const(0), _) => '{0}
+          case (Const(1), _) => y
+          case (Const(a), Const(b)) => Expr(a * b)
+          case (_, Const(_)) => '{ $y * $x }
+          case _ => '{ $x * $y }
+        }
+      case '{ power(${Const(x)}, ${Const(y)}) } =>
+        Expr(power(x, y))
+      case '{ power($x, ${Const(y)}) } =>
+        if y == 0 then '{1}
+        else '{ times($x, power($x, ${Expr(y-1)})) }
+    }
+  )
+
+  val x2 = rewriter.rewrite(x)
+
+  '{
+    println(${Expr(x.show)})
+    println(${Expr(x2.show)})
+    println()
+    $x2
+  }
+}
+
+object Transformation {
+  /** A restrictive transformer that is guaranteed to generate type correct code */
+  def safe[T: Type](transform: PartialFunction[Expr[T], Expr[T]]): Transformation =
+    new SafeTransformation(transform)
+
+    /** A general purpose transformer that may fail while transforming.
+     *  It will check the type of the returned Expr and will throw if the type does not conform to the expected type.
+     */
+  def checked(transform: PartialFunction[Expr[Any], Expr[Any]]): Transformation =
+    new CheckedTransformation(transform)
+}
+
+class CheckedTransformation(transform: PartialFunction[Expr[Any], Expr[Any]]) extends Transformation {
+  def apply[T: Type](e: Expr[T])(given QuoteContext): Expr[T] = {
+    transform.applyOrElse(e, identity) match {
+      case '{ $e2: T } => e2
+      case '{ $e2: $t } =>
+        throw new Exception(
+          s"""Transformed
+            |${e.show}
+            |into
+            |${e2.show}
+            |
+            |Expected type to be
+            |${summon[Type[T]].show}
+            |but was
+            |${t.show}
+          """.stripMargin)
+    }
+  }
+}
+
+class SafeTransformation[U: Type](transform: PartialFunction[Expr[U], Expr[U]]) extends Transformation {
+  def apply[T: Type](e: Expr[T])(given QuoteContext): Expr[T] = {
+    e match {
+      case '{ $e: U } => transform.applyOrElse(e, identity) match { case '{ $e2: T } => e2 }
+      case e => e
+    }
+  }
+}
+
+abstract class Transformation {
+  def apply[T: Type](e: Expr[T])(given QuoteContext): Expr[T]
+}
+
+private object Rewriter {
+  def apply(): Rewriter = new Rewriter(Nil, Nil, false)
+}
+
+private class Rewriter private (preTransform: List[Transformation] = Nil, postTransform: List[Transformation] = Nil, fixPoint: Boolean) {
+
+  def withFixPoint: Rewriter =
+    new Rewriter(preTransform, postTransform, fixPoint = true)
+  def withPre(transform: Transformation): Rewriter =
+    new Rewriter(transform :: preTransform, postTransform, fixPoint)
+  def withPost(transform: Transformation): Rewriter =
+    new Rewriter(preTransform, transform :: postTransform, fixPoint)
+
+  def rewrite[T](e: Expr[T])(given QuoteContext, Type[T]): Expr[T] = {
+    val e2 = preTransform.foldLeft(e)((ei, transform) => transform(ei))
+    val e3 = rewriteChildren(e2)
+    val e4 = postTransform.foldLeft(e3)((ei, transform) => transform(ei))
+    if fixPoint && e4 != e then rewrite(e4) else e4
+  }
+
+  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.widen 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-3/Test_2.scala

@@ -0,0 +1,14 @@
+object Test {
+
+  def main(args: Array[String]): Unit = {
+    val a: Int = 5
+    val b: Int = 6
+    rewrite(plus(1, 4))
+    rewrite(plus(0, a))
+    rewrite(plus(a, b))
+    rewrite(plus(plus(a, 0), plus(0, b)))
+    rewrite(power(4, 5))
+    rewrite(power(a, 5))
+  }
+
+}