[CS] Delay macro expansion until end of CSApply

Attempting to expand macros in the middle of
CSApply can result in attempting to run
MiscDiagnostics within a closure that hasn't yet
had the solution applied to the AST, which can
crash the implicit-self diagnostic logic. Move
the expansion to the end of CSApply such that
expansions are type-checked along with local
decls, ensuring it's run after the solution has
been applied to the AST.

rdar://138997009

Author: hamishknight

lib/Sema/CSApply.cpp

@@ -410,6 +410,7 @@ namespace {
   class ExprRewriter : public ExprVisitor<ExprRewriter, Expr *> {
     // Delayed items to type-check.
     SmallVector<Decl *, 4> LocalDeclsToTypeCheck;
+    SmallVector<MacroExpansionExpr *, 4> MacrosToExpand;
 
   public:
     ConstraintSystem &cs;
@@ -431,6 +432,10 @@ namespace {
       LocalDeclsToTypeCheck.push_back(D);
     }
 
+    void addMacroToExpand(MacroExpansionExpr *E) {
+      MacrosToExpand.push_back(E);
+    }
+
     /// Coerce the given tuple to another tuple type.
     ///
     /// \param expr The expression we're converting.
@@ -5535,16 +5540,19 @@ namespace {
 
       // FIXME: Expansion should be lazy.
       // i.e. 'ExpandMacroExpansionExprRequest' should be sinked into
-      // 'getRewritten()', and performed on-demand.
+      // 'getRewritten()', and performed on-demand. Unfortunately that requires
+      // auditing every ASTWalker's `getMacroWalkingBehavior` since
+      // `MacroWalking::Expansion` does not actually kick expansion.
       if (!cs.Options.contains(ConstraintSystemFlags::DisableMacroExpansions) &&
           // Do not expand macros inside macro arguments. For example for
           // '#stringify(#assert(foo))' when typechecking `#assert(foo)`,
           // we don't want to expand it.
           llvm::none_of(llvm::ArrayRef(ExprStack).drop_back(1),
                         [](Expr *E) { return isa<MacroExpansionExpr>(E); })) {
-        (void)evaluateOrDefault(cs.getASTContext().evaluator,
-                                ExpandMacroExpansionExprRequest{E},
-                                std::nullopt);
+        // We need to delay the expansion until we're done applying the solution
+        // since running MiscDiagnostics on the expansion may walk up and query
+        // the type of a parent closure, e.g `diagnoseImplicitSelfUseInClosure`.
+        addMacroToExpand(E);
       }
 
       cs.cacheExprTypes(E);
@@ -5615,6 +5623,14 @@ namespace {
       // Type-check any local decls encountered.
       for (auto *D : LocalDeclsToTypeCheck)
         TypeChecker::typeCheckDecl(D);
+
+      // Expand any macros encountered.
+      // FIXME: Expansion should be lazy.
+      auto &eval = cs.getASTContext().evaluator;
+      for (auto *E : MacrosToExpand) {
+        (void)evaluateOrDefault(eval, ExpandMacroExpansionExprRequest{E},
+                                std::nullopt);
+      }
     }
 
     /// Diagnose an optional injection that is probably not what the

test/Macros/macro_misc_diags.swift

@@ -98,3 +98,28 @@ _ = #trailingClosure {
   // CHECK-DIAG: Client.swift:[[@LINE-1]]:27: warning: trailing closure in this context is confusable with the body of the statement; pass as a parenthesized argument to silence this warning
   // CHECK-DIAG: @__swiftmacro_6Client0017Clientswift_yEEFcfMX[[@LINE-4]]{{.*}}trailingClosurefMf_.swift:2:27: warning: trailing closure in this context is confusable with the body of the statement
 }
+
+// rdar://138997009 - Make sure we don't crash in MiscDiagnostics' implicit
+// self diagnosis.
+struct rdar138997009 {
+  func foo() {}
+  func bar() {
+    _ = {
+      _ = #trailingClosure {
+        foo()
+      }
+    }
+  }
+}
+
+class rdar138997009_Class {
+  func foo() {}
+  func bar() {
+    _ = {
+      _ = #trailingClosure {
+        foo()
+        // CHECK-DIAG: @__swiftmacro_6Client0017Clientswift_yEEFcfMX[[@LINE-3]]{{.*}}trailingClosurefMf_.swift:2:9: error: call to method 'foo' in closure requires explicit use of 'self' to make capture semantics explicit
+      }
+    }
+  }
+}