[Macros] Ensure that macro expansion parsing has appropriate context.

The parser needs context to do completely reasonable things like... parse
closures.

Author: DougGregor

include/swift/Parse/LocalContext.h

@@ -58,7 +58,12 @@ class LocalContext {
   /// True if we saw any anonymous closures.  This is useful when
   /// parsing an initializer context, because such contexts only
   /// need to exist if the initializer contains closures.
-  bool hasClosures() const { return NextClosureDiscriminator != 0; }    
+  bool hasClosures() const { return NextClosureDiscriminator != 0; }
+
+  /// Override the next closure discriminator value.
+  void overrideNextClosureDiscriminator(unsigned discriminator) {
+    NextClosureDiscriminator = discriminator;
+  }
 };
 
 /// Information associated with parsing the top-level context.

lib/Sema/DebuggerTestingTransform.cpp

@@ -32,36 +32,31 @@
 
 using namespace swift;
 
-namespace {
+ASTWalker::PostWalkResult<Expr *>
+DiscriminatorFinder::walkToExprPost(Expr *E)  {
+  auto *ACE = dyn_cast<AbstractClosureExpr>(E);
+  if (!ACE)
+    return Action::Continue(E);
 
-/// Find available closure discriminators.
-///
-/// The parser typically takes care of assigning unique discriminators to
-/// closures, but the parser is unavailable during semantic analysis.
-class DiscriminatorFinder : public ASTWalker {
-  unsigned NextDiscriminator = 0;
+  unsigned Discriminator = ACE->getDiscriminator();
+  assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
+         "Existing closures should have valid discriminators");
+  if (Discriminator >= NextDiscriminator)
+    NextDiscriminator = Discriminator + 1;
+  if (FirstDiscriminator == AbstractClosureExpr::InvalidDiscriminator ||
+      Discriminator < FirstDiscriminator)
+    FirstDiscriminator = Discriminator;
 
-public:
-  PostWalkResult<Expr *> walkToExprPost(Expr *E) override {
-    auto *ACE = dyn_cast<AbstractClosureExpr>(E);
-    if (!ACE)
-      return Action::Continue(E);
-
-    unsigned Discriminator = ACE->getDiscriminator();
-    assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
-           "Existing closures should have valid discriminators");
-    if (Discriminator >= NextDiscriminator)
-      NextDiscriminator = Discriminator + 1;
-    return Action::Continue(E);
-  }
+  return Action::Continue(E);
+}
 
-  // Get the next available closure discriminator.
-  unsigned getNextDiscriminator() {
-    if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
-      llvm::report_fatal_error("Out of valid closure discriminators");
-    return NextDiscriminator++;
-  }
-};
+unsigned DiscriminatorFinder::getNextDiscriminator() {
+  if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
+    llvm::report_fatal_error("Out of valid closure discriminators");
+  return NextDiscriminator++;
+}
+
+namespace {
 
 /// Instrument decls with sanity-checks which the debugger can evaluate.
 class DebuggerTestingTransform : public ASTWalker {

lib/Sema/TypeCheckMacros.cpp

@@ -217,6 +217,26 @@ Expr *swift::expandMacroExpr(
   // Parse the expression.
   Parser parser(macroBufferID, *macroSourceFile, &ctx.Diags, nullptr, nullptr);
   parser.consumeTokenWithoutFeedingReceiver();
+
+  // Set up a "local context" for parsing, so that we have a source of
+  // closure and local-variable discriminators.
+  LocalContext tempContext{};
+  parser.CurDeclContext = dc;
+  parser.CurLocalContext = &tempContext;
+  {
+    DiscriminatorFinder finder;
+    expr->walk(finder);
+
+    unsigned closureDiscriminator;
+    if (finder.getFirstDiscriminator() ==
+          AbstractClosureExpr::InvalidDiscriminator)
+      closureDiscriminator = 0;
+    else
+      closureDiscriminator = finder.getFirstDiscriminator() + 1;
+
+    tempContext.overrideNextClosureDiscriminator(closureDiscriminator);
+  }
+
   auto parsedResult = parser.parseExpr(diag::expected_macro_expansion_expr);
   if (parsedResult.isParseError() || parsedResult.isNull()) {
     // Tack on a note to say where we expanded the macro from?

lib/Sema/TypeChecker.h

@@ -1499,6 +1499,22 @@ diagnoseAndRemoveAttr(const Decl *D, const DeclAttribute *attr,
   return diagnoseAttrWithRemovalFixIt(D, attr, std::forward<ArgTypes>(Args)...);
 }
 
+/// Look for closure discriminators within an AST.
+class DiscriminatorFinder : public ASTWalker {
+  unsigned FirstDiscriminator = AbstractClosureExpr::InvalidDiscriminator;
+  unsigned NextDiscriminator = 0;
+
+public:
+  PostWalkResult<Expr *> walkToExprPost(Expr *E) override;
+
+  // Get the next available closure discriminator.
+  unsigned getNextDiscriminator();
+
+  unsigned getFirstDiscriminator() const {
+    return FirstDiscriminator;
+  }
+};
+
 } // end namespace swift
 
 #endif

test/Macros/macros.swift

@@ -7,4 +7,9 @@ func test(a: Int, b: Int) {
   // CHECK: macro_expansion_expr type='(Int, String)'{{.*}}name=stringify
   // CHECK-NEXT: argument_list
   // CHECK: tuple_expr type='(Int, String)' location=Macro expansion of #stringify
+
+  let (b, s2) = #stringify({ () -> Bool in return true })
+  // CHECK: macro_expansion_expr type='(() -> Bool, String)'{{.*}}name=stringify
+  // CHECK-NEXT: argument_list
+  // CHECK: tuple_expr type='(() -> Bool, String)' location=Macro expansion of #stringify
 }