Implement quoted Lambda extractor

This extractor returns the lifted representation of the body of an explicit lambda

Author: nicolasstucki

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

@@ -2051,44 +2051,41 @@ class ReflectionCompilerInterface(val rootContext: core.Contexts.Context) extend
     }}
     val argVals = argVals0.reverse
     val argRefs = argRefs0.reverse
-    def rec(fn: Tree, topAscription: Option[TypeTree]): Tree = fn match {
-      case Typed(expr, tpt) =>
-        // we need to retain any type ascriptions we see and:
-        // a) if we succeed, ascribe the result type of the ascription to the inlined body
-        // b) if we fail, re-ascribe the same type to whatever it was we couldn't inline
-        // note: if you see many nested ascriptions, keep the top one as that's what the enclosing expression expects
-        rec(expr, topAscription.orElse(Some(tpt)))
-      case Inlined(call, bindings, expansion) =>
-        // this case must go before closureDef to avoid dropping the inline node
-        cpy.Inlined(fn)(call, bindings, rec(expansion, topAscription))
-      case closureDef(ddef) =>
-        val paramSyms = ddef.vparamss.head.map(param => param.symbol)
-        val paramToVals = paramSyms.zip(argRefs).toMap
-        val result = new TreeTypeMap(
-          oldOwners = ddef.symbol :: Nil,
-          newOwners = ctx.owner :: Nil,
-          treeMap = tree => paramToVals.get(tree.symbol).map(_.withSpan(tree.span)).getOrElse(tree)
-        ).transform(ddef.rhs)
-        topAscription match {
-          case Some(tpt) =>
-            // we assume the ascribed type has an apply that has a MethodType with a single param list (there should be no polys)
-            val methodType = tpt.tpe.member(nme.apply).info.asInstanceOf[MethodType]
-            // result might contain paramrefs, so we substitute them with arg termrefs
-            val resultTypeWithSubst = methodType.resultType.substParams(methodType, argRefs.map(_.tpe))
-            Typed(result, TypeTree(resultTypeWithSubst).withSpan(fn.span)).withSpan(fn.span)
-          case None =>
-            result
-        }
-      case tpd.Block(stats, expr) =>
-        seq(stats, rec(expr, topAscription)).withSpan(fn.span)
-      case _ =>
-        val maybeAscribed = topAscription match {
-          case Some(tpt) => Typed(fn, tpt).withSpan(fn.span)
-          case None => fn
-        }
-        maybeAscribed.select(nme.apply).appliedToArgs(argRefs).withSpan(fn.span)
+    val reducedBody = lambdaExtractor(fn) match {
+      case Some(body) => body(argRefs)
+      case None => fn.select(nme.apply).appliedToArgs(argRefs)
+    }
+    seq(argVals, reducedBody).withSpan(fn.span)
+  }
+
+  def lambdaExtractor(fn: Term)(using ctx: Context): Option[List[Term] => Term] = {
+    def rec(fn: Term, transformBody: Term => Term): Option[List[Term] => Term] = {
+      fn match {
+        case Inlined(call, bindings, expansion) =>
+          // this case must go before closureDef to avoid dropping the inline node
+          rec(expansion, cpy.Inlined(fn)(call, bindings, _))
+        case Typed(expr, tpt) =>
+          val resTpe = tpt.tpe.dropDependentRefinement.argInfos.last
+          rec(expr, Typed(_, TypeTree(resTpe).withSpan(tpt.span)))
+        case closureDef(ddef) =>
+          def replace(body: Term, argRefs: List[Term]): Term = {
+            val paramSyms = ddef.vparamss.head.map(param => param.symbol)
+            val paramToVals = paramSyms.zip(argRefs).toMap
+            new TreeTypeMap(
+              oldOwners = ddef.symbol :: Nil,
+              newOwners = ctx.owner :: Nil,
+              treeMap = tree => paramToVals.get(tree.symbol).map(_.withSpan(tree.span)).getOrElse(tree)
+            ).transform(body)
+          }
+          Some(argRefs => replace(transformBody(ddef.rhs), argRefs))
+        case Block(stats, expr) =>
+          // this case must go after closureDef to avoid matching the closure
+          rec(expr, cpy.Block(fn)(stats, _))
+        case _ =>
+          None
+      }
     }
-    seq(argVals, rec(fn, None))
+    rec(fn, identity)
   }
 
   /////////////

library/src/scala/quoted/matching/Lambda.scala

@@ -0,0 +1,29 @@
+package scala.quoted
+package matching
+
+/** Lambda expression extractor */
+object Lambda {
+
+  /** `case Lambda(body)` matche a lambda and extract the body.
+   *   As the body may (will) contain references to the paramter, `body` is a function that recieves those arguments as `Expr`.
+   *   Once this function is applied the result will be the body of the lambda with all references to the parameters replaced.
+   *   If `body` is of type `(T1, T2, ...) => R` then body will be of type `(Expr[T1], Expr[T2], ...) => Expr[R]`.
+   *
+   *   ```
+   *   '{ (x: Int) => println(x) } match
+   *     case Lambda(body) =>
+   *       // where `body` is: (x: Expr[Int]) => '{ println($x) }
+   *       body('{3}) // returns '{ println(3) }
+   *   ```
+   */
+  def unapply[F, Args <: Tuple, Res, G](expr: Expr[F])(using qctx: QuoteContext, tf: TupledFunction[F, Args => Res], tg: TupledFunction[G, Tuple.Map[Args, Expr] => Expr[Res]]): Option[/*QuoteContext ?=>*/ G] = {
+    import qctx.tasty.{_, given _ }
+    qctx.tasty.internal.lambdaExtractor(expr.unseal).map { fn =>
+      def f(args: Tuple.Map[Args, Expr]): Expr[Res] =
+        fn(args.toArray.map(_.asInstanceOf[Expr[Any]].unseal).toList).seal.asInstanceOf[Expr[Res]]
+      tg.untupled(f)
+    }
+
+  }
+
+}

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

@@ -1543,4 +1543,6 @@ trait CompilerInterface {
    */
   def betaReduce(f: Term, args: List[Term])(using ctx: Context): Term
 
+  def lambdaExtractor(term: Term)(using ctx: Context): Option[List[Term] => Term]
+
 }

tests/run-macros/lambda-extractor-1.check

@@ -0,0 +1,6 @@
+scala.Predef.identity[scala.Int](1)
+1
+{
+  scala.Predef.println(1)
+  1
+}

tests/run-macros/lambda-extractor-1/Macro_1.scala

@@ -0,0 +1,11 @@
+import scala.quoted._
+import scala.quoted.matching._
+
+inline def test(inline f: Int => Int): String = ${ impl('f) }
+
+def impl(using QuoteContext)(f: Expr[Int => Int]): Expr[String] = {
+  Expr(f match {
+    case Lambda(body) => body('{1}).show
+    case _ => f.show
+  })
+}

tests/run-macros/lambda-extractor-1/Test_2.scala

@@ -0,0 +1,6 @@
+
+@main def Test = {
+  println(test(identity))
+  println(test(x => x))
+  println(test(x => { println(x); x }))
+}

tests/run-macros/lambda-extractor-2.check

@@ -0,0 +1,5 @@
+1.+(2)
+{
+  scala.Predef.println(1)
+  2
+}

tests/run-macros/lambda-extractor-2/Macro_1.scala

@@ -0,0 +1,11 @@
+import scala.quoted._
+import scala.quoted.matching._
+
+inline def test(inline f: (Int, Int) => Int): String = ${ impl('f) }
+
+def impl(using QuoteContext)(f: Expr[(Int, Int) => Int]): Expr[String] = {
+  Expr(f match {
+    case Lambda(body) => body('{1}, '{2}).show
+    case _ => f.show
+  })
+}

tests/run-macros/lambda-extractor-2/Test_2.scala

@@ -0,0 +1,5 @@
+
+@main def Test = {
+  println(test((x, y) => x + y))
+  println(test((x, y) => { println(x); y }))
+}

tests/run-staging/i3876-c.check

@@ -6,5 +6,5 @@
 
   (f: scala.Function1[scala.Int, scala.Int] {
     def apply(x: scala.Int): scala.Int
-  }).apply(3)
-}
+  })
+}.apply(3)