Implement diagram macros

Author: liufengyun

scalatest/src/main/scala/org/scalatest/DiagrammedExpr.scala

@@ -30,7 +30,7 @@ private[org] case class AnchorValue(anchor: Int, value: Any)
  * so that the generated code can be compiled.  It is expected that ScalaTest users would ever need to use <code>DiagrammedExpr</code>
  * directly.
  */
-trait DiagrammedExpr[T] {
+trait DiagrammedExpr[+T] {
   val anchor: Int
   def anchorValues: List[AnchorValue]
   def value: T

scalatest/src/main/scala/org/scalatest/DiagrammedExprMacro.scala

@@ -422,41 +422,121 @@ import scala.quoted._
 
 
 object DiagrammedExprMacro {
+  def let[S: Type, T](expr: Expr[S])(body: Expr[S] => Expr[T]): Expr[T] =
+    '{
+      val x = ~expr
+      ~body('(x))
+    }
+
+  def lets[S: Type, T](xs: List[Expr[S]])(body: List[Expr[S]] => Expr[T]): Expr[T] = {
+    def rec(xs: List[Expr[S]], acc: List[Expr[S]]): Expr[T] = xs match {
+      case Nil => body(acc)
+      case x :: xs => let(x) { (x: Expr[S]) => rec(xs, x :: acc) }
+    }
+    rec(xs, Nil)
+  }
 
   // Transform the input expression by parsing out the anchor and generate expression that can support diagram rendering
-  def parse(condition: Expr[Boolean])(implicit refl: Reflection): Expr[DiagrammedExpr[Boolean]] = {
+  def parse[T:Type](expr: Expr[T])(implicit refl: Reflection): Expr[DiagrammedExpr[T]] = {
     import refl._
     import quoted.Toolbox.Default._
     import Term._
 
     def isXmlSugar(apply: Apply): Boolean = apply.tpe <:< typeOf[scala.xml.Elem]
     def isJavaStatic(tree: Tree): Boolean = tree.symbol.flags.is(Flags.Static)
 
-    condition.unseal match {
-      case Apply(Select(New(_), _), _) => simpleExpr(condition)  // delegate to simpleExpr if it is a New expression
-      case IsApply(apply) if isXmlSugar(apply) => simpleExpr(condition)
-      case IsApply(apply) if isJavaStatic(apply) => simpleExpr(condition)
-      case IsApply(apply) => applyExpr(condition) // delegate to applyExpr if it is Apply
-      case IsTypeApply(apply) => applyExpr(condition) // delegate to applyExpr if it is Apply
-      case Select(This(_), _) => simpleExpr(condition) // delegate to simpleExpr if it is a Select for this, e.g. referring a to instance member.
-      case IsSelect(x) if x.symbol.flags.is(Flags.Object) => simpleExpr(condition) // don't traverse packages
-      case IsSelect(x) if isJavaStatic(x) => simpleExpr(condition)
-      case IsSelect(select) => selectExpr(condition) // delegate to selectExpr if it is a Select
+    expr.unseal match {
+      case Apply(Select(New(_), _), _) => simpleExpr(expr)  // delegate to simpleExpr if it is a New expression
+      case IsApply(apply) if isXmlSugar(apply) => simpleExpr(expr)
+      case IsApply(apply) if isJavaStatic(apply) => simpleExpr(expr)
+      case IsApply(apply) => applyExpr(expr) // delegate to applyExpr if it is Apply
+      case IsTypeApply(apply) => applyExpr(expr) // delegate to applyExpr if it is Apply
+      case Select(This(_), _) => simpleExpr(expr) // delegate to simpleExpr if it is a Select for this, e.g. referring a to instance member.
+      case IsSelect(x) if x.symbol.flags.is(Flags.Object) => simpleExpr(expr) // don't traverse packages
+      case IsSelect(x) if isJavaStatic(x) => simpleExpr(expr)
+      case IsSelect(select) => selectExpr(expr) // delegate to selectExpr if it is a Select
       case Block(stats, expr) =>
-        Block(stats, parse(expr.seal[Boolean]).unseal).seal[DiagrammedExpr[Boolean]] // call parse recursively using the expr argument if it is a block
-      case _ => simpleExpr(condition) // for others, just delegate to simpleExpr
+        Block(stats, parse(expr.seal[T]).unseal).seal[DiagrammedExpr[T]] // call parse recursively using the expr argument if it is a block
+      case _ => simpleExpr(expr) // for others, just delegate to simpleExpr
     }
   }
 
-  def applyExpr(condition: Expr[Boolean])(implicit refl: Reflection): Expr[DiagrammedExpr[Boolean]] = ???
-  def selectExpr(condition: Expr[Boolean])(implicit refl: Reflection): Expr[DiagrammedExpr[Boolean]] = ???
+  def applyExpr[T:Type](expr: Expr[T])(implicit refl: Reflection): Expr[DiagrammedExpr[T]] = {
+    import refl._
+    import quoted.Toolbox.Default._
+    import Term._
+
+    def apply(l: Expr[_], name: String, r: List[Expr[_]]): Expr[T] = ???
+
+    expr.unseal.underlyingArgument match {
+      case Term.Apply(Term.Select(lhs, op), rhs :: Nil) =>
+        op match {
+          case "||" | "|" =>
+            val left = parse(lhs.seal[T & Boolean])
+            val right = parse(rhs.seal[T & Boolean])
+            '{
+              val l = ~left
+              val r = ~right
+              if (l.value) l
+              else DiagrammedExpr.applyExpr(l, r :: Nil, r.value, ~getAnchor(expr))
+            }
+          case "&&" | "&" =>
+            val left = parse(lhs.seal[T & Boolean])
+            val right = parse(rhs.seal[T & Boolean])
+            '{
+              val l = ~left
+              val r = ~right
+              if (l.value) DiagrammedExpr.applyExpr(l, r :: Nil, r.value, ~getAnchor(expr))
+              else l
+            }
+          case _ =>
+            val left = parse(lhs.seal[Any])
+            val right = parse(rhs.seal[Any])
+            '{
+              val l = ~left
+              val r = ~right
+              val res = ~apply('(l.value), op, '(r.value) :: Nil)
+              DiagrammedExpr.applyExpr(l, r :: Nil, res, ~getAnchor(expr))
+            }
+        }
+      case Term.Apply(Term.Select(lhs, op), args) =>
+        val left = parse(lhs.seal[Any])
+        val rights = args.map(arg => parse(arg.seal[Any]))
+
+        let(left) { (l: Expr[DiagrammedExpr[_]]) =>
+          lets(rights) { (rs: List[Expr[DiagrammedExpr[_]]]) =>
+            val res = apply('((~l).value), op, rs)
+            '{ DiagrammedExpr.applyExpr(~l, ~rs.toExprOfList, ~res, ~getAnchor(expr)) }
+          }
+        }
+      case _ =>
+        simpleExpr(expr)
+    }
+  }
+
+  def selectExpr[T:Type](expr: Expr[T])(implicit refl: Reflection): Expr[DiagrammedExpr[T]] = {
+    import refl._
+    import quoted.Toolbox.Default._
+    import Term._
+
+    def selectField(o: Expr[_], name: String): Expr[T] = ???
+
+    expr.unseal match {
+      case Select(qual, name) =>
+        val obj = parse(qual.seal[Any])
+
+        '{
+          val o = ~obj
+          DiagrammedExpr.selectExpr(o, ~selectField('(o.value), name), ~getAnchor(expr))
+        }
+    }
+  }
 
   def transform(
-    helper:Expr[(DiagrammedExpr[Boolean], Any, String, source.Position) => Assertion],
+    helper: Expr[(DiagrammedExpr[Boolean], Any, String, source.Position) => Assertion],
     condition: Expr[Boolean], prettifier: Expr[Prettifier],
     pos: Expr[source.Position], clue: Expr[Any], sourceText: String
-  )
-  (implicit refl: Reflection): Expr[Assertion] = ???
+  )(implicit refl: Reflection): Expr[Assertion] = ???
 
 
   /**

scalatest/src/main/scala/org/scalatest/matchers/TypeMatcherMacro.scala

@@ -182,15 +182,20 @@ private[scalatest] object TypeMatcherMacro {
   // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAType, used by 'shouldBe a [type]' syntax
   def shouldBeATypeImpl(self: Expr[org.scalatest.Matchers#AnyShouldWrapper[_]], aType: Expr[ResultOfATypeInvocation[_]])(implicit refl: Reflection): Expr[org.scalatest.Assertion] = {
     import refl._
-    checkTypeParameter(refl)(aType.unseal, "an")
+    checkTypeParameter(refl)(aType.unseal, "a")
     '{
       TypeMatcherHelper.assertAType((~self).leftSideValue, ~aType, (~self).prettifier, (~self).pos)
     }
   }
 
-//   // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAType, used by 'mustBe a [type]' syntax
-//   def mustBeATypeImpl(context: Context)(aType: context.Expr[ResultOfATypeInvocation[_]]): context.Expr[org.scalatest.Assertion] =
-//     assertTypeImpl(context)(aType.tree, "mustBe a", "assertAType")
+  // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAType, used by 'mustBe a [type]' syntax
+  def mustBeATypeImpl(self: Expr[org.scalatest.MustMatchers#AnyMustWrapper[_]], aType: Expr[ResultOfATypeInvocation[_]])(implicit refl: Reflection): Expr[org.scalatest.Assertion] = {
+    import refl._
+    checkTypeParameter(refl)(aType.unseal, "a")
+    '{
+      TypeMatcherHelper.assertAType((~self).leftSideValue, ~aType, (~self).prettifier, (~self).pos)
+    }
+  }
 
   // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAType, used by 'shouldBe an [type]' syntax
   def shouldBeAnTypeImpl(self: Expr[org.scalatest.Matchers#AnyShouldWrapper[_]], anType: Expr[ResultOfAnTypeInvocation[_]])(implicit refl: Reflection): Expr[org.scalatest.Assertion] = {
@@ -201,9 +206,14 @@ private[scalatest] object TypeMatcherMacro {
     }
   }
 
-//   // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAnType, used by 'mustBe an [type]' syntax
-//   def mustBeAnTypeImpl(context: Context)(anType: context.Expr[ResultOfAnTypeInvocation[_]]): context.Expr[org.scalatest.Assertion] =
-//     assertTypeImpl(context)(anType.tree, "mustBe an", "assertAnType")
+  // Do checking on type parameter and generate AST to call TypeMatcherHelper.checkAnType, used by 'mustBe an [type]' syntax
+  def mustBeAnTypeImpl(self: Expr[org.scalatest.MustMatchers#AnyMustWrapper[_]], anType: Expr[ResultOfAnTypeInvocation[_]])(implicit refl: Reflection): Expr[org.scalatest.Assertion] = {
+    import refl._
+    checkTypeParameter(refl)(anType.unseal, "an")
+    '{
+      TypeMatcherHelper.assertAnType((~self).leftSideValue, ~anType, (~self).prettifier, (~self).pos)
+    }
+  }
 
 //   /*def expectTypeImpl(context: Context)(tree: context.Tree, beMethodName: String, assertMethodName: String): context.Expr[org.scalatest.Fact] = {
 //     import context.universe._