Merge pull request #63985 from DougGregor/expression-macros-effects

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,26 @@ class ASTWalker {
     return LazyInitializerWalking::InPatternBinding;
   }
 
+  /// This method configures how the walker should walk into uses of macros.
+  virtual MacroWalking getMacroWalkingBehavior() const {
+    return MacroWalking::ArgumentsAndExpansion;
+  }
+
+  /// 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 +579,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/IDE/Utils.h

@@ -368,6 +368,10 @@ class NameMatcher: public ASTWalker {
   bool handleCustomAttrs(Decl *D);
   ArgumentList *getApplicableArgsFor(Expr* E);
 
+  MacroWalking getMacroWalkingBehavior() const override {
+    return MacroWalking::Arguments;
+  }
+
   PreWalkResult<Expr *> walkToExprPre(Expr *E) override;
   PostWalkResult<Expr *> walkToExprPost(Expr *E) override;
   PreWalkAction walkToDeclPre(Decl *D) override;

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)) {