[Macros] Rework the way we assign closure and local discriminators

Setting closure and local discriminators depends on an in-order walk
of the AST. For macros, it was walking into both macro expansions and
arguments. However, this doesn't work well with lazy macro expansions,
and could result in some closures/local variables not getting
discriminators set at all.

Make the assignment of discriminators only walk macro arguments, and
then lazily assign discriminators for anything within a macro
expansion or in ill-formed code. This replaces the single global "next
autoclosure discriminator" scheme with a per-DeclContext scheme, that
is more reliable/robust, although it does mean that discriminators
of closures and locals within macro expansions are dependent on
ordering. That shouldn't matter, because these are local values.

Fixes rdar://108682196.

Author: DougGregor

include/swift/AST/ASTContext.h

@@ -342,10 +342,6 @@ class ASTContext final {
   /// The # of times we have performed typo correction.
   unsigned NumTypoCorrections = 0;
 
-  /// The next auto-closure discriminator.  This needs to be preserved
-  /// across invocations of both the parser and the type-checker.
-  unsigned NextAutoClosureDiscriminator = 0;
-
   /// Cached mapping from types to their associated tangent spaces.
   llvm::DenseMap<Type, Optional<TangentSpace>> AutoDiffTangentSpaces;
 
@@ -1121,6 +1117,13 @@ class ASTContext final {
   unsigned getNextMacroDiscriminator(MacroDiscriminatorContext context,
                                      DeclBaseName baseName);
 
+  /// Get the next discriminator within the given declaration context.
+  unsigned getNextDiscriminator(const DeclContext *dc);
+
+  /// Set the maximum assigned discriminator within the given declaration context.
+  void setMaxAssignedDiscriminator(
+      const DeclContext *dc, unsigned discriminator);
+
   /// Retrieve the Clang module loader for this ASTContext.
   ///
   /// If there is no Clang module loader, returns a null pointer.

lib/AST/ASTContext.cpp

@@ -403,6 +403,9 @@ struct ASTContext::Implementation {
   llvm::DenseMap<std::pair<const void *, Identifier>, unsigned>
       NextMacroDiscriminator;
 
+  /// Local and closure discriminators per context.
+  llvm::DenseMap<const DeclContext *, unsigned> NextDiscriminator;
+
   /// Structure that captures data that is segregated into different
   /// arenas.
   struct Arena {
@@ -2185,6 +2188,18 @@ unsigned ASTContext::getNextMacroDiscriminator(
   return getImpl().NextMacroDiscriminator[key]++;
 }
 
+/// Get the next discriminator within the given declaration context.
+unsigned ASTContext::getNextDiscriminator(const DeclContext *dc) {
+  return getImpl().NextDiscriminator[dc];
+}
+
+/// Set the maximum assigned discriminator within the given declaration context.
+void ASTContext::setMaxAssignedDiscriminator(
+    const DeclContext *dc, unsigned discriminator) {
+  assert(discriminator >= getImpl().NextDiscriminator[dc]);
+  getImpl().NextDiscriminator[dc] = discriminator;
+}
+
 void ASTContext::verifyAllLoadedModules() const {
 #ifndef NDEBUG
   FrontendStatsTracer tracer(Stats, "verify-all-loaded-modules");

lib/AST/Decl.cpp

@@ -2934,10 +2934,26 @@ unsigned ValueDecl::getLocalDiscriminator() const {
     return 0;
 
   // Assign local discriminators in this context.
+  ASTContext &ctx = getASTContext();
   evaluateOrDefault(
-      getASTContext().evaluator,
+      ctx.evaluator,
       LocalDiscriminatorsRequest{getDeclContext()}, InvalidDiscriminator);
 
+  // If we don't have a discriminator, and either
+  //   1. We have ill-formed code and we're able to assign a discriminator, or
+  //   2. We are in a macro expansion buffer
+  //
+  // then assign the next discriminator now.
+  if (LocalDiscriminator == InvalidDiscriminator &&
+      (ctx.Diags.hadAnyError() ||
+       (getLoc().isValid() &&
+        getModuleContext()->getSourceFileContainingLocation(getLoc())
+          ->getFulfilledMacroRole() != None))) {
+    auto discriminator = ctx.getNextDiscriminator(getDeclContext());
+    ctx.setMaxAssignedDiscriminator(getDeclContext(), discriminator + 1);
+    const_cast<ValueDecl *>(this)->LocalDiscriminator = discriminator;
+  }
+
   assert(LocalDiscriminator != InvalidDiscriminator);
 
   return LocalDiscriminator;

lib/AST/Expr.cpp

@@ -1901,16 +1901,22 @@ unsigned AbstractClosureExpr::getDiscriminator() const {
   if (raw != InvalidDiscriminator)
     return raw;
 
+  ASTContext &ctx = getASTContext();
   evaluateOrDefault(
-      getASTContext().evaluator, LocalDiscriminatorsRequest{getParent()}, 0);
+      ctx.evaluator, LocalDiscriminatorsRequest{getParent()}, 0);
 
-  // Ill-formed code might not be able to assign discriminators, so assign
-  // a new one now.
+  // If we don't have a discriminator, and either
+  //   1. We have ill-formed code and we're able to assign a discriminator, or
+  //   2. We are in a macro expansion buffer
+  //
+  // then assign the next discriminator now.
   if (getRawDiscriminator() == InvalidDiscriminator &&
-      getASTContext().Diags.hadAnyError()) {
+      (ctx.Diags.hadAnyError() ||
+       getParentSourceFile()->getFulfilledMacroRole() != None)) {
+    auto discriminator = ctx.getNextDiscriminator(getParent());
+    ctx.setMaxAssignedDiscriminator(getParent(), discriminator + 1);
     const_cast<AbstractClosureExpr *>(this)->
-        Bits.AbstractClosureExpr.Discriminator =
-          getASTContext().NextAutoClosureDiscriminator++;
+        Bits.AbstractClosureExpr.Discriminator = discriminator;
   }
 
   assert(getRawDiscriminator() != InvalidDiscriminator);

lib/Sema/TypeCheckStmt.cpp

@@ -255,7 +255,7 @@ namespace {
     }
 
     MacroWalking getMacroWalkingBehavior() const override {
-      return MacroWalking::ArgumentsAndExpansion;
+      return MacroWalking::Arguments;
     }
 
     PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
@@ -410,7 +410,6 @@ unsigned LocalDiscriminatorsRequest::evaluate(
                              LocalDiscriminatorsRequest{dc->getParent()}, 0);
   }
 
-  Optional<unsigned> expectedNextAutoclosureDiscriminator = None;
   ASTNode node;
   ParameterList *params = nullptr;
   ParamDecl *selfParam = nullptr;
@@ -436,7 +435,7 @@ unsigned LocalDiscriminatorsRequest::evaluate(
     params = closure->getParameters();
   } else if (auto topLevel = dyn_cast<TopLevelCodeDecl>(dc)) {
     node = topLevel->getBody();
-    expectedNextAutoclosureDiscriminator = ctx.NextAutoClosureDiscriminator;
+    dc = topLevel->getParentModule();
   } else if (auto patternBindingInit = dyn_cast<PatternBindingInitializer>(dc)){
     auto patternBinding = patternBindingInit->getBinding();
     node = patternBinding->getInit(patternBindingInit->getBindingIndex());
@@ -472,12 +471,11 @@ unsigned LocalDiscriminatorsRequest::evaluate(
     params = getParameterList(dc);
   }
 
+  auto startDiscriminator = ctx.getNextDiscriminator(dc);
   if (!node && !params && !selfParam)
-    return 0;
+    return startDiscriminator;
 
-  SetLocalDiscriminators visitor(
-      expectedNextAutoclosureDiscriminator.value_or(0)
-  );
+  SetLocalDiscriminators visitor(startDiscriminator);
 
   // Set local discriminator for the 'self' parameter.
   if (selfParam)
@@ -494,9 +492,7 @@ unsigned LocalDiscriminatorsRequest::evaluate(
     node.walk(visitor);
 
   unsigned nextDiscriminator = visitor.maxAssignedDiscriminator();
-  if (expectedNextAutoclosureDiscriminator) {
-    ctx.NextAutoClosureDiscriminator = nextDiscriminator;
-  }
+  ctx.setMaxAssignedDiscriminator(dc, nextDiscriminator);
 
   // Return the next discriminator.
   return nextDiscriminator;

test/Macros/macro_expand.swift

@@ -400,3 +400,12 @@ func testFreestandingWithClosure(i: Int) {
     return x
   }
 }
+
+// Nested macros with closures
+@freestanding(expression) macro coerceToInt<T>(_ value: T) -> Int = #externalMacro(module: "MacroDefinition", type: "CoerceToIntMacro")
+
+func testFreestandingClosureNesting() {
+  _ = #stringify({ () -> Int in
+    #coerceToInt(2)
+  })
+}