Merge pull request #63738 from bnbarham/macros-everywhere

[AST] Allow ignoring macro expansions

Author: bnbarham

include/swift/AST/ASTWalker.h

@@ -541,6 +541,10 @@ 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/Attr.h

@@ -58,6 +58,7 @@ class ClassDecl;
 class GenericFunctionType;
 class LazyConformanceLoader;
 class LazyMemberLoader;
+class ModuleDecl;
 class PatternBindingInitializer;
 class TrailingWhereClause;
 class TypeExpr;
@@ -2581,6 +2582,59 @@ class DeclAttributes {
   SourceLoc getStartLoc(bool forModifiers = false) const;
 };
 
+/// Attributes applied directly to the declaration.
+///
+/// We should really just have \c DeclAttributes and \c SemanticDeclAttributes,
+/// but currently almost all callers expect the latter. Instead of changing all
+/// callers of \c getAttrs, instead provide a way to retrieive the original
+/// attributes.
+class OrigDeclAttributes {
+  SmallVector<DeclAttribute *> attributes;
+  using AttrListTy = SmallVectorImpl<DeclAttribute *>;
+
+public:
+  OrigDeclAttributes() = default;
+
+  OrigDeclAttributes(DeclAttributes semanticAttrs, ModuleDecl *mod);
+
+
+  using iterator = AttrListTy::iterator;
+  using const_iterator = AttrListTy::const_iterator;
+
+  iterator begin() { return attributes.begin(); }
+  iterator end() { return attributes.end(); }
+  const_iterator begin() const { return attributes.begin(); }
+  const_iterator end() const { return attributes.end(); }
+
+  template <typename AttrType, bool AllowInvalid>
+  using AttributeKindRange =
+  OptionalTransformRange<iterator_range<const_iterator>,
+  ToAttributeKind<AttrType, AllowInvalid>, const_iterator>;
+
+  template <typename AttrType, bool AllowInvalid = false>
+  AttributeKindRange<AttrType, AllowInvalid> getAttributes() const {
+    return AttributeKindRange<AttrType, AllowInvalid>(make_range(begin(), end()), ToAttributeKind<AttrType, AllowInvalid>());
+  }
+
+  /// Retrieve the first attribute of the given attribute class.
+  template <typename AttrType>
+  const AttrType *getAttribute(bool allowInvalid = false) const {
+    for (auto *attr : attributes) {
+      if (auto *specificAttr = dyn_cast<AttrType>(attr)) {
+        if (specificAttr->isValid() || allowInvalid)
+          return specificAttr;
+      }
+    }
+    return nullptr;
+  }
+
+  /// Determine whether there is an attribute with the given attribute class.
+  template <typename AttrType>
+  bool hasAttribute(bool allowInvalid = false) const {
+    return getAttribute<AttrType>(allowInvalid) != nullptr;
+  }
+};
+
 /// TypeAttributes - These are attributes that may be applied to types.
 class TypeAttributes {
   // Get a SourceLoc for every possible attribute that can be parsed in source.

include/swift/AST/Decl.h

@@ -862,6 +862,11 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   /// with this declaration.
   ArrayRef<CustomAttr *> getRuntimeDiscoverableAttrs() const;
 
+  /// Returns the attributes that were directly attached to this declaration
+  /// as written in source, ie. does not include attributes generated by macro
+  /// expansions.
+  OrigDeclAttributes getOriginalAttrs() const;
+
   /// Returns the semantic attributes attached to this declaration,
   /// including attributes that are generated as the result of member
   /// attribute macro expansion.
@@ -923,6 +928,10 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
 
   SourceLoc TrailingSemiLoc;
 
+  /// Whether this declaration is within a generated buffer, \c false if this
+  /// declaration was constructed from a serialized module.
+  bool isInGeneratedBuffer() const;
+
   /// Returns the appropriate kind of entry point to generate for this class,
   /// based on its attributes.
   ///

include/swift/AST/Module.h

@@ -349,6 +349,10 @@ class ModuleDecl
   /// \c nullptr if the source location isn't in this module.
   SourceFile *getSourceFileContainingLocation(SourceLoc loc);
 
+  /// Whether the given location is inside a generated buffer, \c false if
+  /// the given location isn't in this module.
+  bool isInGeneratedBuffer(SourceLoc loc);
+
   /// Creates a map from \c #filePath strings to corresponding \c #fileID
   /// strings, diagnosing any conflicts.
   ///

include/swift/IDE/SourceEntityWalker.h

@@ -201,6 +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; }
+
 protected:
   SourceEntityWalker() = default;
   SourceEntityWalker(const SourceEntityWalker &) = default;

include/swift/IDE/SyntaxModel.h

@@ -139,7 +139,7 @@ struct SyntaxStructureElement {
 struct SyntaxStructureNode {
   const Decl *Dcl = nullptr;
   SyntaxStructureKind Kind;
-  DeclAttributes Attrs;
+  OrigDeclAttributes Attrs;
   CharSourceRange Range;
   CharSourceRange BodyRange;
   CharSourceRange NameRange;

lib/AST/ASTWalker.cpp

@@ -432,9 +432,11 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   bool visitMacroExpansionDecl(MacroExpansionDecl *MED) {
     if (MED->getArgs() && doIt(MED->getArgs()))
       return true;
-    for (auto *decl : MED->getRewritten())
-      if (doIt(decl))
-        return true;
+    if (Walker.shouldWalkMacroExpansions()) {
+      for (auto *decl : MED->getRewritten())
+        if (doIt(decl))
+          return true;
+    }
     return false;
   }
 
@@ -1295,17 +1297,19 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   Expr *visitMacroExpansionExpr(MacroExpansionExpr *E) {
-    Expr *rewritten = nullptr;
-    if (E->getRewritten()) {
-      rewritten = doIt(E->getRewritten());
-      if (!rewritten) return nullptr;
-    }
     ArgumentList *args = nullptr;
     if (E->getArgs()) {
       args = doIt(E->getArgs());
       if (!args) return nullptr;
     }
-    E->setRewritten(rewritten);
+    if (Walker.shouldWalkMacroExpansions()) {
+      Expr *rewritten = nullptr;
+      if (E->getRewritten()) {
+        rewritten = doIt(E->getRewritten());
+        if (!rewritten) return nullptr;
+      }
+      E->setRewritten(rewritten);
+    }
     E->setArgs(args);
     return E;
   }
@@ -1421,6 +1425,9 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
   }
 
   bool shouldSkip(Decl *D) {
+    if (!Walker.shouldWalkMacroExpansions() && D->isInGeneratedBuffer())
+      return true;
+
     if (auto *VD = dyn_cast<VarDecl>(D)) {
       // VarDecls are walked via their NamedPattern, ignore them if we encounter
       // then in the few cases where they are also pushed outside as members.

lib/AST/Attr.cpp

@@ -851,6 +851,14 @@ SourceLoc DeclAttributes::getStartLoc(bool forModifiers) const {
   return lastAttr ? lastAttr->getRangeWithAt().Start : SourceLoc();
 }
 
+OrigDeclAttributes::OrigDeclAttributes(DeclAttributes allAttributes, ModuleDecl *mod) {
+  for (auto *attr : allAttributes) {
+    if (!mod->isInGeneratedBuffer(attr->AtLoc)) {
+      attributes.emplace_back(attr);
+    }
+  }
+}
+
 static void printAvailableAttr(const AvailableAttr *Attr, ASTPrinter &Printer,
                                const PrintOptions &Options) {
   if (Attr->isLanguageVersionSpecific())

lib/AST/Decl.cpp

@@ -367,6 +367,10 @@ StringRef Decl::getDescriptiveKindName(DescriptiveDeclKind K) {
   llvm_unreachable("bad DescriptiveDeclKind");
 }
 
+OrigDeclAttributes Decl::getOriginalAttrs() const {
+  return OrigDeclAttributes(getAttrs(), getModuleContext());
+}
+
 DeclAttributes Decl::getSemanticAttrs() const {
   auto mutableThis = const_cast<Decl *>(this);
   (void)evaluateOrDefault(getASTContext().evaluator,
@@ -684,7 +688,7 @@ SourceRange Decl::getSourceRangeIncludingAttrs() const {
 
     // Otherwise, include attributes directly attached to the accessor.
     SourceLoc VarLoc = AD->getStorage()->getStartLoc();
-    for (auto Attr : getAttrs()) {
+    for (auto *Attr : getOriginalAttrs()) {
       if (!Attr->getRange().isValid())
         continue;
 
@@ -703,19 +707,23 @@ SourceRange Decl::getSourceRangeIncludingAttrs() const {
   if (auto *PBD = dyn_cast<PatternBindingDecl>(this)) {
     for (auto i : range(PBD->getNumPatternEntries()))
       PBD->getPattern(i)->forEachVariable([&](VarDecl *VD) {
-        for (auto Attr : VD->getAttrs())
+        for (auto *Attr : VD->getOriginalAttrs())
           if (Attr->getRange().isValid())
             Range.widen(Attr->getRangeWithAt());
       });
   }
 
-  for (auto Attr : getAttrs()) {
+  for (auto *Attr : getOriginalAttrs()) {
     if (Attr->getRange().isValid())
       Range.widen(Attr->getRangeWithAt());
   }
   return Range;
 }
 
+bool Decl::isInGeneratedBuffer() const {
+  return getModuleContext()->isInGeneratedBuffer(getStartLoc());
+}
+
 SourceLoc Decl::getLocFromSource() const {
   switch (getKind()) {
 #define DECL(ID, X) \

lib/AST/Module.cpp

@@ -648,6 +648,13 @@ SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   return foundSourceFile;
 }
 
+bool ModuleDecl::isInGeneratedBuffer(SourceLoc loc) {
+  SourceFile *file = getSourceFileContainingLocation(loc);
+  if (!file)
+    return false;
+  return file->Kind == SourceFileKind::MacroExpansion;
+}
+
 ArrayRef<SourceFile *>
 PrimarySourceFilesRequest::evaluate(Evaluator &evaluator,
                                     ModuleDecl *mod) const {