Merge pull request swiftlang#66198 from DougGregor/closure-local-discriminators-in-macro-expansions

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

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)
+  })
+}