Merge pull request #3697 from dotty-staging/quote-constant-extraction

Quote constant extraction

Author: nicolasstucki

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

@@ -1,15 +1,18 @@
 package dotty.tools.dotc.quoted
 
-import java.io.PrintStream
-
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.Phases.Phase
 
-/** Compiler that takes the contents of a quoted expression `expr` and produces outputs
- *  the pretty printed code.
- */
-class ExprDecompiler(out: PrintStream) extends ExprCompiler(null) {
+/** Compiler that takes the contents of a quoted expression `expr` and outputs it's tree. */
+class ExprDecompiler(output: tpd.Tree => Context => Unit) extends ExprCompiler(null) {
   override def phases: List[List[Phase]] = List(
     List(new ExprFrontend(putInClass = false)), // Create class from Expr
-    List(new QuotePrinter(out)) // Print all loaded classes
+    List(new QuoteTreeOutput(output))
   )
+
+  class QuoteTreeOutput(output: tpd.Tree => Context => Unit) extends Phase {
+    override def phaseName: String = "quotePrinter"
+    override def run(implicit ctx: Context): Unit = output(ctx.compilationUnit.tpdTree)(ctx)
+  }
 }

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

@@ -1,17 +1,17 @@
 package dotty.tools.dotc.quoted
 
+import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.Driver
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.StdNames._
 import dotty.tools.io.{AbstractFile, Directory, PlainDirectory, VirtualDirectory}
 import dotty.tools.repl.AbstractFileClassLoader
+import dotty.tools.dotc.printing.DecompilerPrinter
 
 import scala.quoted.Expr
-import java.io.ByteArrayOutputStream
-import java.io.PrintStream
-import java.nio.charset.StandardCharsets
 
 class QuoteDriver extends Driver {
+  import tpd._
 
   def run[T](expr: Expr[T], settings: Runners.RunSettings): T = {
     val ctx: Context = initCtx.fresh
@@ -39,18 +39,24 @@ class QuoteDriver extends Driver {
   }
 
   def show(expr: Expr[_]): String = {
+    def show(tree: Tree, ctx: Context): String = {
+      val printer = new DecompilerPrinter(ctx)
+      val pageWidth = ctx.settings.pageWidth.value(ctx)
+      printer.toText(tree).mkString(pageWidth, false)
+    }
+    withTree(expr, show)
+  }
+
+  def withTree[T](expr: Expr[_], f: (Tree, Context) => T): T = {
     val ctx: Context = initCtx.fresh
     ctx.settings.color.update("never")(ctx) // TODO support colored show
-    val baos = new ByteArrayOutputStream
-    var ps: PrintStream = null
-    try {
-      ps = new PrintStream(baos, true, "utf-8")
-
-      new ExprDecompiler(ps).newRun(ctx).compileExpr(expr)
-
-      new String(baos.toByteArray, StandardCharsets.UTF_8)
+    var output: Option[T] = None
+    def registerTree(tree: tpd.Tree)(ctx: Context): Unit = {
+      assert(output.isEmpty)
+      output = Some(f(tree, ctx))
     }
-    finally if (ps != null) ps.close()
+    new ExprDecompiler(registerTree).newRun(ctx).compileExpr(expr)
+    output.getOrElse(throw new Exception("Could not extact " + expr))
   }
 
   override def initCtx: Context = {

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

@@ -1,7 +1,9 @@
 package dotty.tools.dotc.quoted
 
-import dotty.tools.dotc.ast.Trees.Literal
-import dotty.tools.dotc.core.Constants.Constant
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Constants._
+import dotty.tools.dotc.core.Contexts._
 import dotty.tools.dotc.printing.RefinedPrinter
 
 import scala.quoted.Expr
@@ -10,19 +12,34 @@ import scala.runtime.quoted._
 
 /** Default runners for quoted expressions */
 object Runners {
+  import tpd._
 
   implicit def runner[T]: Runner[T] = new Runner[T] {
 
     def run(expr: Expr[T]): T = Runners.run(expr, RunSettings())
 
     def show(expr: Expr[T]): String = expr match {
       case expr: ConstantExpr[T] =>
-        val ctx = new QuoteDriver().initCtx
-        ctx.settings.color.update("never")(ctx)
+        implicit val ctx = new QuoteDriver().initCtx
+        ctx.settings.color.update("never")
         val printer = new RefinedPrinter(ctx)
         printer.toText(Literal(Constant(expr.value))).mkString(Int.MaxValue, false)
       case _ => new QuoteDriver().show(expr)
     }
+
+    def toConstantOpt(expr: Expr[T]): Option[T] = {
+      def toConstantOpt(tree: Tree): Option[T] = tree match {
+        case Literal(Constant(c)) => Some(c.asInstanceOf[T])
+        case Block(Nil, e) => toConstantOpt(e)
+        case Inlined(_, Nil, e) => toConstantOpt(e)
+        case _ => None
+      }
+      expr match {
+        case expr: ConstantExpr[T] => Some(expr.value)
+        case _ => new QuoteDriver().withTree(expr, (tree, _) => toConstantOpt(tree))
+      }
+    }
+
   }
 
   def run[T](expr: Expr[T], settings: RunSettings): T = expr match {

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

@@ -0,0 +1,7 @@
+package scala.quoted
+
+import scala.runtime.quoted.Runner
+
+object Constant {
+  def unapply[T](expr: Expr[T])(implicit runner: Runner[T]): Option[T] = runner.toConstantOpt(expr)
+}

library/src/scala/quoted/Constant.scala

@@ -3,6 +3,6 @@ package scala.quoted
 import scala.runtime.quoted.Unpickler.Pickled
 
 /** An Expr backed by a pickled TASTY tree */
-final case class TastyExpr[T](tasty: Pickled, args: Seq[Any]) extends Expr[T] with TastyQuoted {
+final class TastyExpr[T](val tasty: Pickled, val args: Seq[Any]) extends Expr[T] with TastyQuoted {
   override def toString(): String = s"Expr(<pickled>)"
 }

library/src/scala/quoted/TastyExpr.scala

@@ -3,6 +3,6 @@ package scala.quoted
 import scala.runtime.quoted.Unpickler.Pickled
 
 /** A Type backed by a pickled TASTY tree */
-final case class TastyType[T](tasty: Pickled, args: Seq[Any]) extends Type[T] with TastyQuoted {
+final class TastyType[T](val tasty: Pickled, val args: Seq[Any]) extends Type[T] with TastyQuoted {
   override def toString(): String = s"Type(<pickled>)"
 }

library/src/scala/quoted/TastyType.scala

@@ -7,4 +7,5 @@ import scala.quoted.Expr
 trait Runner[T] {
   def run(expr: Expr[T]): T
   def show(expr: Expr[T]): String
+  def toConstantOpt(expr: Expr[T]): Option[T]
 }

library/src/scala/runtime/quoted/Runner.scala

@@ -0,0 +1,36 @@
+3
+4
+abc
+null
+OK
+{
+  {
+    val y: Double = 3.0.*(3.0)
+    y
+  }
+}
+9.0
+{
+  {
+    val y: Double = 4.0.*(4.0)
+    y
+  }
+}
+16.0
+{
+  {
+    val y: Double = 5.0.*(5.0)
+    y
+  }
+}
+25.0
+{
+  Test.dynamicPower(
+    {
+      println("foo")
+      2
+    }
+  , 6.0)
+}
+foo
+36.0

tests/run-with-compiler/quote-run-constants-extract.check

@@ -0,0 +1,52 @@
+import scala.quoted._
+
+import dotty.tools.dotc.quoted.Runners._
+
+object Test {
+
+  def main(args: Array[String]): Unit = {
+    (3: Expr[Int]) match { case Constant(n) => println(n) }
+    '(4) match { case Constant(n) => println(n) }
+    '("abc") match { case Constant(n) => println(n) }
+    '(null) match { case Constant(n) => println(n) }
+
+    '(new Object) match { case Constant(n) => println(n); case _ => println("OK") }
+
+
+    // 2 is a lifted constant
+    println(power(2, 3.0).show)
+    println(power(2, 3.0).run)
+
+    // n is a lifted constant
+    val n = 2
+    println(power(n, 4.0).show)
+    println(power(n, 4.0).run)
+
+    // n is a constant in a quote
+    println(power('(2), 5.0).show)
+    println(power('(2), 5.0).run)
+
+    // n2 is clearly not a constant
+    val n2 = '{ println("foo"); 2 }
+    println(power(n2, 6.0).show)
+    println(power(n2, 6.0).run)
+  }
+
+  def power(n: Expr[Int], x: Expr[Double]): Expr[Double] = {
+    n match {
+      case Constant(n1) => powerCode(n1, x)
+      case _ => '{ dynamicPower(~n, ~x) }
+    }
+  }
+
+  private def powerCode(n: Int, x: Expr[Double]): Expr[Double] =
+    if (n == 0) '(1.0)
+    else if (n == 1) x
+    else if (n % 2 == 0) '{ { val y = ~x * ~x; ~powerCode(n / 2, '(y)) } }
+    else '{ ~x * ~powerCode(n - 1, x) }
+
+  def dynamicPower(n: Int, x: Double): Double =
+    if (n == 0) 1.0
+    else if (n % 2 == 0) dynamicPower(n / 2, x * x)
+    else x * dynamicPower(n - 1, x)
+}