From 0f1d24dbb551b95a514cbad0f81e2bf3af0e3535 Mon Sep 17 00:00:00 2001
From: Nicolas Stucki <nicolas.stucki@gmail.com>
Date: Mon, 18 Jan 2021 10:14:07 +0100
Subject: [PATCH] Track levels in Inlining phase

---
 .../dotty/tools/dotc/transform/Inlining.scala | 21 +++++++++++++++----
 tests/run-staging/inline-at-level-2.scala     | 20 ++++++++++++++++++
 2 files changed, 37 insertions(+), 4 deletions(-)
 create mode 100644 tests/run-staging/inline-at-level-2.scala

diff --git a/compiler/src/dotty/tools/dotc/transform/Inlining.scala b/compiler/src/dotty/tools/dotc/transform/Inlining.scala
index 163bfb25ff4f..60b5057df33f 100644
--- a/compiler/src/dotty/tools/dotc/transform/Inlining.scala
+++ b/compiler/src/dotty/tools/dotc/transform/Inlining.scala
@@ -39,8 +39,19 @@ class Inlining extends MacroTransform {
 
   override def checkPostCondition(tree: Tree)(using Context): Unit =
     tree match {
-      case tree: RefTree if !Inliner.inInlineMethod && StagingContext.level == 0 =>
-        assert(!tree.symbol.isInlineMethod)
+      case PackageDef(pid, _) if tree.symbol.owner == defn.RootClass =>
+        new TreeTraverser {
+          def traverse(tree: Tree)(using Context): Unit =
+            tree match
+              case _: GenericApply if tree.symbol.isQuote =>
+                traverseChildren(tree)(using StagingContext.quoteContext)
+              case _: GenericApply if tree.symbol.isExprSplice =>
+                traverseChildren(tree)(using StagingContext.spliceContext)
+              case tree: RefTree if !Inliner.inInlineMethod && StagingContext.level == 0 =>
+                assert(!tree.symbol.isInlineMethod, tree.show)
+              case _ =>
+                traverseChildren(tree)
+        }.traverse(tree)
       case _ =>
     }
 
@@ -57,9 +68,11 @@ class Inlining extends MacroTransform {
           val inlined = Inliner.inlineCall(tree1)
           if tree1 eq inlined then inlined
           else transform(inlined)
-        case _: TypeApply if tree.symbol.isQuote =>
+        case _: GenericApply if tree.symbol.isQuote =>
           ctx.compilationUnit.needsStaging = true
-          super.transform(tree)
+          super.transform(tree)(using StagingContext.quoteContext)
+        case _: GenericApply if tree.symbol.isExprSplice =>
+          super.transform(tree)(using StagingContext.spliceContext)
         case _ =>
           super.transform(tree)
 
diff --git a/tests/run-staging/inline-at-level-2.scala b/tests/run-staging/inline-at-level-2.scala
new file mode 100644
index 000000000000..5ec622a66768
--- /dev/null
+++ b/tests/run-staging/inline-at-level-2.scala
@@ -0,0 +1,20 @@
+
+import scala.quoted._
+import scala.quoted.staging._
+
+object Test {
+
+  given Toolbox = Toolbox.make(getClass.getClassLoader)
+
+  def main(args: Array[String]): Unit = run {
+    '{
+      def m(using Quotes) = '{
+        errorOnceInlined() // Should not be inlined while `run` is executed. It would be inlined in the next stage.
+      }
+      ()
+    }
+  }
+
+  inline def errorOnceInlined() = compiletime.error("Error")
+
+}