Partially fix ExInlinedCall

frontends/dotty/src/main/scala/stainless/frontends/dotc/ASTExtractors.scala

@@ -529,27 +529,49 @@ trait ASTExtractors {
     }
 
     object ExInlinedCall {
-      /** Extracts an inlined function or method call, returning a 4-tuple
-        * containing the receiver, function or method symbol, type arguments,
-        * and term arguments.
-        * 
+      /** Extracts an inlined function or method call, returning a 5-tuple
+        * containing the inlined call's receiver, function or method symbol,
+        * type arguments, term arguments, and expansion.
+        *
         * Unlift the receiver and arguments to their definitions if they are
         * inline proxies.
         */
-      def unapply(tree: tpd.Tree): Option[(Option[tpd.Tree], Symbol, Seq[tpd.Tree], Seq[tpd.Tree])] = tree match {
-        // TODO(mbovel): should probably use InlineOrProxy instead of
-        // Synthetic. Why is this not working?
-        case Block(stats, Inlined(ExCall(rec, sym, tps, args), _, _)) if stats.forall(_.symbol.is(Synthetic)) =>
-          def unliftArg(arg: tpd.Tree) = arg match {
-            case arg @ Ident(_) if arg.symbol.is(Synthetic) =>
-              stats.find(_.symbol == arg.symbol) match
-                case Some(proxyDef: tpd.ValOrDefDef) => proxyDef.rhs
-                case _ => throw new IllegalStateException(s"Could not find inline proxy definition for $arg")
-            case arg => arg
-          }
-          Some(rec.map(unliftArg), sym, tps, args.map(unliftArg))
-        case Inlined(ExCall(rec, sym, tps, args), _, _) =>
-          Some(rec, sym, tps, args)
+      def unapply(tree: tpd.Tree): Option[(Option[tpd.Tree], Symbol, Seq[tpd.Tree], Seq[tpd.Tree], tpd.Tree)] = tree match {
+        case Block(stats, Inlined(ExCall(rec, sym, tps, args), _, expansion))
+            if stats.forall(_.symbol.isOneOf(InlineProxy | Synthetic)) =>
+
+          val unliftInlineProxies =
+            new tpd.TreeMap:
+              /** Adapted from Dotty's `tpd.MapToUnderlying.transform`. */
+              override def transform(tree: tpd.Tree)(using DottyContext): tpd.Tree =
+                tree match
+                  // The `InlineProxy` flag does not seem to be set for inline
+                  // argument proxies, not sure why. It should be set at:
+                  // https://github.com/scala/scala3/blob/20e6f11f4fe47982259eba949eea78d65765142f/compiler/src/dotty/tools/dotc/inlines/Inliner.scala#L224.
+                  // Including `Synthetic` to make this work, like in Dotty's
+                  // `tpd.TreeOps.underlyingArgument`, but this might be an
+                  // over-approximation.
+                  case tree: tpd.Ident if tree.symbol.isOneOf(InlineProxy | Synthetic) =>
+                    stats.find(_.symbol == tree.symbol) match
+                      case Some(defTree: tpd.ValOrDefDef) =>
+                          val rhs = defTree.rhs
+                          assert(!rhs.isEmpty)
+                          // Contrary to Dotty's `tpd.MapToUnderlying`, do not
+                          // recurse over rhs.
+                          rhs
+                      case defTree =>
+                        throw new IllegalStateException(s"Expected ValOrDefDef for ${tree.symbol}, found ${defTree.map(_.show)}")
+                  case tree =>
+                    super.transform(tree)
+          Some(
+            rec,
+            sym,
+            tps,
+            args,
+            unliftInlineProxies.transform(expansion)
+          )
+        case Inlined(ExCall(rec, sym, tps, args), _, expansion) =>
+          Some(rec, sym, tps, args, expansion)
         case _ => None
       }
     }
@@ -607,7 +629,7 @@ trait ASTExtractors {
       def unapply(tree: tpd.Tree): Option[tpd.Tree] = tree match {
         case Apply(
             ExSymbol("scala", "Predef$", "Ensuring") |
-            ExSymbol("stainless", "lang", "StaticChecks$", "Ensuring"),
+            ExSymbol("stainless", "lang", "StaticChecks$", "ensuring"),
             Seq(arg)) => Some(arg)
 
         case Apply(ExSymbol("stainless", "lang", "package$", "Throwing"), Seq(arg)) => Some(arg)
@@ -1225,8 +1247,11 @@ trait ASTExtractors {
           _,
           ExSymbol("stainless", "lang", "StaticChecks$", "ensuring"),
           _,
-          Seq(rec, contract, message)
-        ) => Some((rec, contract, true))
+          Seq(_, contract, message),
+          expansion
+        ) =>
+          println(s"Static Ensuring:\n${expansion.show},\n${contract.show}")
+          Some((expansion, contract, true))
 
         case _ => None
       }