[Macros] Be deliberate about walking macro arguments vs. expansions

Provide ASTWalker with a customization point to specify whether to
check macro arguments (which are type checked but never emitted), the
macro expansion (which is the result of applying the macro and is
actually emitted into the source), or both. Provide answers for the
~115 different ASTWalker visitors throughout the code base.

Fixes rdar://104042945, which concerns checking of effects in
macro arguments---which we shouldn't do.

Author: DougGregor

include/swift/AST/ASTWalker.h

@@ -79,6 +79,24 @@ enum class LazyInitializerWalking {
   InAccessor
 };
 
+/// Specifies the behavior for walking a macro expansion, whether we want to
+/// see the macro arguments, the expansion, or both.
+enum class MacroWalking {
+  /// Walk into the expansion of the macro, to see the semantic effect of
+  /// the macro expansion.
+  Expansion,
+
+  /// Walk into the arguments of the macro as written in the source code.
+  ///
+  /// The actual arguments walked may not make it into the program itself,
+  /// because they can be translated by the macro in arbitrary ways.
+  Arguments,
+
+  /// Walk into both the arguments of the macro as written in the source code
+  /// and also the macro expansion.
+  ArgumentsAndExpansion
+};
+
 /// An abstract class used to traverse an AST.
 class ASTWalker {
 public:
@@ -504,6 +522,24 @@ class ASTWalker {
     return LazyInitializerWalking::InPatternBinding;
   }
 
+  /// This method configures how the walker should walk into uses of macros.
+  virtual MacroWalking getMacroWalkingBehavior() const = 0;
+
+  /// Determine whether we should walk macro arguments (as they appear in
+  /// source) and the expansion (which is semantically part of the program).
+  std::pair<bool, bool> shouldWalkMacroArgumentsAndExpansion() const {
+    switch (getMacroWalkingBehavior()) {
+    case MacroWalking::Expansion:
+      return std::make_pair(false, true);
+
+    case MacroWalking::Arguments:
+      return std::make_pair(true, false);
+
+    case MacroWalking::ArgumentsAndExpansion:
+      return std::make_pair(true, true);
+    }
+  }
+
   /// This method configures whether the walker should visit the body of a
   /// closure that was checked separately from its enclosing expression.
   ///
@@ -541,10 +577,6 @@ class ASTWalker {
   /// TODO: Consider changing this to false by default.
   virtual bool shouldWalkSerializedTopLevelInternalDecls() { return true; }
 
-  /// Whether we should walk into expanded macros, whether it be from a
-  /// \c MacroExpansionExpr or declarations created by attached macros.
-  virtual bool shouldWalkMacroExpansions() { return true; }
-
   /// walkToParameterListPre - This method is called when first visiting a
   /// ParameterList, before walking into its parameters.
   ///

include/swift/AST/TypeDeclFinder.h

@@ -67,6 +67,10 @@ class TypeReprIdentFinder : public ASTWalker {
   /// The function to call when a \c IdentTypeRepr is seen.
   llvm::function_ref<bool(const IdentTypeRepr *)> Callback;
 
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
   PostWalkAction walkToTypeReprPost(TypeRepr *TR) override;
 
 public:

include/swift/IDE/SourceEntityWalker.h

@@ -201,9 +201,10 @@ class SourceEntityWalker {
 
   virtual bool shouldWalkIntoGenericParams() { return true; }
 
-  /// Whether we should walk into expanded macros, whether it be from a
-  /// \c MacroExpansionExpr or declarations created by attached macros.
-  virtual bool shouldWalkMacroExpansions() { return false; }
+  /// Only walk the arguments of a macro, to represent the source as written.
+  virtual MacroWalking getMacroWalkingBehavior() const {
+    return MacroWalking::Arguments;
+  }
 
 protected:
   SourceEntityWalker() = default;

include/swift/Sema/CompletionContextFinder.h

@@ -49,6 +49,10 @@ class CompletionContextFinder : public ASTWalker {
   DeclContext *InitialDC;
 
 public:
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
   /// Finder for completion contexts within the provided initial expression.
   CompletionContextFinder(ASTNode initialNode, DeclContext *DC)
       : InitialExpr(initialNode.dyn_cast<Expr *>()), InitialDC(DC) {

include/swift/Sema/ConstraintSystem.h

@@ -3376,6 +3376,11 @@ class ConstraintSystem {
     CacheExprTypes(Expr *expr, ConstraintSystem &cs, bool excludeRoot)
         : RootExpr(expr), CS(cs), ExcludeRoot(excludeRoot) {}
 
+    /// Walk everything in a macro
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     PostWalkResult<Expr *> walkToExprPost(Expr *expr) override {
       if (ExcludeRoot && expr == RootExpr) {
         assert(!expr->getType() && "Unexpected type in root of expression!");
@@ -6954,6 +6959,10 @@ class OverloadSetCounter : public ASTWalker {
   : NumOverloads(overloads)
   {}
 
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
     if (auto applyExpr = dyn_cast<ApplyExpr>(expr)) {
       // If we've found function application and it's

lib/APIDigester/ModuleAnalyzerNodes.cpp

@@ -2405,6 +2405,11 @@ class ConstExtractor: public ASTWalker {
     }
     return false;
   }
+
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::ArgumentsAndExpansion;
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     if (E->isSemanticallyConstExpr()) {
       record(E, E);

lib/AST/ASTScopeCreation.cpp

@@ -141,6 +141,10 @@ class ScopeCreator final : public ASTAllocated<ScopeCreator> {
       PreWalkAction walkToParameterListPre(ParameterList *PL) override {
         return Action::SkipChildren();
       }
+
+      MacroWalking getMacroWalkingBehavior() const override {
+        return MacroWalking::ArgumentsAndExpansion;
+      }
     };
 
     expr->walk(ClosureFinder(*this, parent));

lib/AST/ASTVerifier.cpp

@@ -271,6 +271,10 @@ class Verifier : public ASTWalker {
     return Verifier(topDC->getParentModule(), DC);
   }
 
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Expansion;
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
     switch (E->getKind()) {
 #define DISPATCH(ID) return dispatchVisitPreExpr(static_cast<ID##Expr*>(E))

lib/AST/ASTWalker.cpp

@@ -430,9 +430,13 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool visitMacroExpansionDecl(MacroExpansionDecl *MED) {
-    if (MED->getArgs() && doIt(MED->getArgs()))
+    bool shouldWalkArguments, shouldWalkExpansion;
+    std::tie(shouldWalkArguments, shouldWalkExpansion) =
+        Walker.shouldWalkMacroArgumentsAndExpansion();
+    if (shouldWalkArguments && MED->getArgs() && doIt(MED->getArgs()))
       return true;
-    if (Walker.shouldWalkMacroExpansions()) {
+
+    if (shouldWalkExpansion) {
       for (auto *decl : MED->getRewritten())
         if (doIt(decl))
           return true;
@@ -1297,20 +1301,24 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitMacroExpansionExpr(MacroExpansionExpr *E) {
-    ArgumentList *args = nullptr;
-    if (E->getArgs()) {
-      args = doIt(E->getArgs());
+    bool shouldWalkArguments, shouldWalkExpansion;
+    std::tie(shouldWalkArguments, shouldWalkExpansion) =
+        Walker.shouldWalkMacroArgumentsAndExpansion();
+
+    if (shouldWalkArguments && E->getArgs()) {
+      ArgumentList *args = doIt(E->getArgs());
       if (!args) return nullptr;
+      E->setArgs(args);
     }
-    if (Walker.shouldWalkMacroExpansions()) {
+
+    if (shouldWalkExpansion) {
       Expr *rewritten = nullptr;
       if (E->getRewritten()) {
         rewritten = doIt(E->getRewritten());
         if (!rewritten) return nullptr;
       }
       E->setRewritten(rewritten);
     }
-    E->setArgs(args);
     return E;
   }
 
@@ -1425,7 +1433,8 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool shouldSkip(Decl *D) {
-    if (!Walker.shouldWalkMacroExpansions() && D->isInGeneratedBuffer())
+    if (!Walker.shouldWalkMacroArgumentsAndExpansion().second &&
+        D->isInGeneratedBuffer())
       return true;
 
     if (auto *VD = dyn_cast<VarDecl>(D)) {

lib/AST/Decl.cpp

@@ -7648,6 +7648,11 @@ swift::findWrappedValuePlaceholder(Expr *init) {
   public:
     PropertyWrapperValuePlaceholderExpr *placeholder = nullptr;
 
+    /// Only walk the arguments of a macro, to represent the source as written.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
       if (placeholder)
         return Action::SkipChildren(E);
@@ -8343,6 +8348,11 @@ AbstractFunctionDecl::getBodyFingerprintIncludingLocalTypeMembers() const {
   public:
     HashCombiner(StableHasher &hasher) : hasher(hasher) {}
 
+    /// Only walk the arguments of a macro, to represent the source as written.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     PreWalkAction walkToDeclPre(Decl *D) override {
       if (D->isImplicit())
         return Action::SkipChildren();

lib/AST/Expr.cpp

@@ -488,7 +488,12 @@ forEachImmediateChildExpr(llvm::function_ref<Expr *(Expr *)> callback) {
   struct ChildWalker : ASTWalker {
     llvm::function_ref<Expr *(Expr *)> callback;
     Expr *ThisNode;
-    
+
+    /// Only walk the arguments of a macro, to represent the source as written.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     ChildWalker(llvm::function_ref<Expr *(Expr *)> callback, Expr *ThisNode)
       : callback(callback), ThisNode(ThisNode) {}
 
@@ -533,6 +538,11 @@ void Expr::forEachChildExpr(llvm::function_ref<Expr *(Expr *)> callback) {
   struct ChildWalker : ASTWalker {
     llvm::function_ref<Expr *(Expr *)> callback;
 
+    /// Only walk the arguments of a macro, to represent the source as written.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     ChildWalker(llvm::function_ref<Expr *(Expr *)> callback)
     : callback(callback) {}
 
@@ -853,6 +863,11 @@ llvm::DenseMap<Expr *, Expr *> Expr::getParentMap() {
   public:
     llvm::DenseMap<Expr *, Expr *> &ParentMap;
 
+    /// Walk everything that's available.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     explicit RecordingTraversal(llvm::DenseMap<Expr *, Expr *> &parentMap)
       : ParentMap(parentMap) { }
 
@@ -1265,6 +1280,11 @@ void PackExpansionExpr::getExpandedPacks(SmallVectorImpl<ASTNode> &packs) {
   struct PackCollector : public ASTWalker {
     llvm::SmallVector<ASTNode, 2> packs;
 
+    /// Walk everything that's available.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
       // Don't walk into nested pack expansions
       if (isa<PackExpansionExpr>(E)) {

lib/AST/InlinableText.cpp

@@ -49,6 +49,11 @@ class IsFeatureCheck : public ASTWalker {
 public:
   bool foundFeature = false;
 
+  /// Walk everything that's available.
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::ArgumentsAndExpansion;
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
     if (auto unresolved = dyn_cast<UnresolvedDeclRefExpr>(expr)) {
       if (unresolved->getName().getBaseName().userFacingName().startswith("$"))
@@ -110,6 +115,11 @@ struct ExtractInactiveRanges : public ASTWalker {
   SmallVector<CharSourceRange, 4> ranges;
   SourceManager &sourceMgr;
 
+  /// Walk everything that's available.
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::ArgumentsAndExpansion;
+  }
+
   explicit ExtractInactiveRanges(SourceManager &sourceMgr)
     : sourceMgr(sourceMgr) {}
 

lib/AST/NameLookup.cpp

@@ -2920,6 +2920,11 @@ CollectedOpaqueReprs swift::collectOpaqueReturnTypeReprs(TypeRepr *r, ASTContext
   public:
     explicit Walker(CollectedOpaqueReprs &reprs, ASTContext &ctx, DeclContext *d) : Reprs(reprs), Ctx(ctx), dc(d) {}
 
+    /// Walk everything that's available.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     PreWalkAction walkToTypeReprPre(TypeRepr *repr) override {
 
       // Don't allow variadic opaque parameter or return types.

lib/AST/Pattern.cpp

@@ -167,6 +167,12 @@ namespace {
     WalkToVarDecls(const std::function<void(VarDecl*)> &fn)
     : fn(fn) {}
 
+    /// Walk everything that's available; there shouldn't be macro expansions
+    /// that matter anyway.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     PostWalkResult<Pattern *> walkToPatternPost(Pattern *P) override {
       // Handle vars.
       if (auto *Named = dyn_cast<NamedPattern>(P))

lib/AST/Stmt.cpp

@@ -245,6 +245,11 @@ ASTNode BraceStmt::findAsyncNode() {
   class FindInnerAsync : public ASTWalker {
     ASTNode AsyncNode;
 
+    /// Walk only the macro arguments.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     PreWalkResult<Expr *> walkToExprPre(Expr *expr) override {
       // If we've found an 'await', record it and terminate the traversal.
       if (isa<AwaitExpr>(expr)) {

lib/AST/TypeRepr.cpp

@@ -83,6 +83,11 @@ bool TypeRepr::findIf(llvm::function_ref<bool(TypeRepr *)> pred) {
     explicit Walker(llvm::function_ref<bool(TypeRepr *)> pred)
         : Pred(pred), FoundIt(false) {}
 
+    /// Walk everything in a macro
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::ArgumentsAndExpansion;
+    }
+
     PreWalkAction walkToTypeReprPre(TypeRepr *ty) override {
       if (Pred(ty)) {
         FoundIt = true;

lib/ConstExtract/ConstExtract.cpp

@@ -47,6 +47,10 @@ class NominalTypeConformanceCollector : public ASTWalker {
       std::vector<NominalTypeDecl *> &ConformanceDecls)
       : Protocols(Protocols), ConformanceTypeDecls(ConformanceDecls) {}
 
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::ArgumentsAndExpansion;
+  }
+
   PreWalkAction walkToDeclPre(Decl *D) override {
     if (auto *NTD = llvm::dyn_cast<NominalTypeDecl>(D))
       if (!isa<ProtocolDecl>(NTD))

lib/IDE/CompletionLookup.cpp

@@ -187,6 +187,11 @@ bool swift::ide::canDeclContextHandleAsync(const DeclContext *DC) {
       const ClosureExpr *Target;
       bool Result = false;
 
+      /// Walk everything in a macro.
+      MacroWalking getMacroWalkingBehavior() const override {
+        return MacroWalking::ArgumentsAndExpansion;
+      }
+
       AsyncClosureChecker(const ClosureExpr *Target) : Target(Target) {}
 
       PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
@@ -3230,6 +3235,11 @@ void CompletionLookup::getStmtLabelCompletions(SourceLoc Loc, bool isContinue) {
   public:
     SmallVector<Identifier, 2> Result;
 
+    /// Walk only the arguments of a macro.
+    MacroWalking getMacroWalkingBehavior() const override {
+      return MacroWalking::Arguments;
+    }
+
     LabelFinder(SourceManager &SM, SourceLoc TargetLoc, bool IsContinue)
         : SM(SM), TargetLoc(TargetLoc), IsContinue(IsContinue) {}