[Macros] Harden discriminator verification and macro expansion mangling

Author: DougGregor

lib/AST/ASTMangler.cpp

@@ -3764,6 +3764,7 @@ void ASTMangler::appendMacroExpansionContext(
     outerExpansionDC = expansion->getDeclContext();
     discriminator = expansion->getDiscriminator();
     role = MacroRole::Declaration;
+    baseName = expansion->getMacroName().getBaseName();
     break;
   }
 
@@ -3799,6 +3800,17 @@ void ASTMangler::appendMacroExpansionContext(
     outerExpansionLoc = decl->getLoc();
     outerExpansionDC = decl->getDeclContext();
     discriminator = decl->getAttachedMacroDiscriminator(role, attr);
+
+    auto *macroDecl = evaluateOrDefault(
+        ctx.evaluator,
+        ResolveMacroRequest{const_cast<CustomAttr *>(attr),
+                            getAttachedMacroRoles(), outerExpansionDC},
+        nullptr);
+    if (macroDecl)
+      baseName = macroDecl->getBaseName();
+    else
+      baseName = ctx.getIdentifier("__unknown_macro__");
+
     break;
   }
 

lib/AST/ASTVerifier.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/Initializer.h"
+#include "swift/AST/MacroDiscriminatorContext.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ParameterList.h"
 #include "swift/AST/Pattern.h"
@@ -233,7 +234,8 @@ class Verifier : public ASTWalker {
       ClosureDiscriminators;
   DeclContext *CanonicalTopLevelSubcontext = nullptr;
 
-  typedef std::pair<DeclContext *, Identifier> MacroExpansionDiscriminatorKey;
+  typedef std::pair</*MacroDiscriminatorContext*/const void *, Identifier>
+      MacroExpansionDiscriminatorKey;
   llvm::DenseMap<MacroExpansionDiscriminatorKey, SmallBitVector>
       MacroExpansionDiscriminators;
 
@@ -2410,9 +2412,8 @@ class Verifier : public ASTWalker {
     }
 
     void verifyChecked(MacroExpansionExpr *expansion) {
-      auto dc = getCanonicalDeclContext(expansion->getDeclContext());
       MacroExpansionDiscriminatorKey key{
-        dc,
+        MacroDiscriminatorContext::getParentOf(expansion).getOpaqueValue(),
         expansion->getMacroName().getBaseName().getIdentifier()
       };
       auto &discriminatorSet = MacroExpansionDiscriminators[key];
@@ -2433,9 +2434,8 @@ class Verifier : public ASTWalker {
     }
 
     void verifyChecked(MacroExpansionDecl *expansion) {
-      auto dc = getCanonicalDeclContext(expansion->getDeclContext());
       MacroExpansionDiscriminatorKey key{
-        dc,
+        MacroDiscriminatorContext::getParentOf(expansion).getOpaqueValue(),
         expansion->getMacroName().getBaseName().getIdentifier()
       };
       auto &discriminatorSet = MacroExpansionDiscriminators[key];