[IDE][Index] Property wrapper rename support

include/swift/AST/PrintOptions.h

@@ -299,6 +299,9 @@ struct PrintOptions {
   /// Empty means allow all.
   std::vector<AnyAttrKind> ExclusiveAttrList;
 
+  /// List of decls that should be printed even if they are implicit and \c SkipImplicit is set to true.
+  std::vector<const Decl*> TreatAsExplicitDeclList;
+
   /// Whether to print function @convention attribute on function types.
   bool PrintFunctionRepresentationAttrs = true;
 

include/swift/IDE/Utils.h

@@ -268,6 +268,10 @@ class NameMatcher: public ASTWalker {
   std::vector<UnresolvedLoc> LocsToResolve;
   std::vector<ResolvedLoc> ResolvedLocs;
   ArrayRef<Token> TokensToCheck;
+
+  /// The \c Expr argument of a parent \c CustomAttr (if one exists) and
+  /// the \c SourceLoc of the type name it applies to.
+  llvm::Optional<std::pair<SourceLoc, Expr *>> CustomAttrArg;
   unsigned InactiveConfigRegionNestings = 0;
   unsigned SelectorNestings = 0;
 
@@ -287,6 +291,7 @@ class NameMatcher: public ASTWalker {
                   bool checkParentForLabels = false);
   bool tryResolve(ASTWalker::ParentTy Node, SourceLoc NameLoc, LabelRangeType RangeType,
                   ArrayRef<CharSourceRange> LabelLocs);
+  bool handleCustomAttrs(Decl *D);
 
   std::pair<bool, Expr*> walkToExprPre(Expr *E) override;
   Expr* walkToExprPost(Expr *E) override;

lib/AST/ASTPrinter.cpp

@@ -1572,8 +1572,11 @@ bool ShouldPrintChecker::shouldPrint(const Decl *D,
     return false;
   }
 
-  if (Options.SkipImplicit && D->isImplicit())
-    return false;
+  if (Options.SkipImplicit && D->isImplicit()) {
+    const auto &IgnoreList = Options.TreatAsExplicitDeclList;
+    if (std::find(IgnoreList.begin(), IgnoreList.end(), D) == IgnoreList.end())
+        return false;
+  }
 
   if (Options.SkipUnavailable &&
       D->getAttrs().isUnavailable(D->getASTContext()))

lib/AST/ASTWalker.cpp

@@ -122,31 +122,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     return inherited::visit(PL);
   }
 
-  //===--------------------------------------------------------------------===//
-  //                               Attributes
-  //===--------------------------------------------------------------------===//
-  bool visitCustomAttributes(Decl *D) {
-    for (auto *customAttr : D->getAttrs().getAttributes<CustomAttr, true>()) {
-      CustomAttr *mutableCustomAttr = const_cast<CustomAttr *>(customAttr);
-      if (doIt(mutableCustomAttr->getTypeLoc()))
-        return true;
-
-      if (auto semanticInit = customAttr->getSemanticInit()) {
-        if (auto newSemanticInit = doIt(semanticInit))
-          mutableCustomAttr->setSemanticInit(newSemanticInit);
-        else
-          return true;
-      } else if (auto arg = customAttr->getArg()) {
-        if (auto newArg = doIt(arg))
-          mutableCustomAttr->setArg(newArg);
-        else
-          return true;
-      }
-    }
-
-    return false;
-  }
-
   //===--------------------------------------------------------------------===//
   //                                 Decls
   //===--------------------------------------------------------------------===//
@@ -179,9 +154,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
     // If there is a single variable, walk it's attributes.
     bool isPropertyWrapperBackingProperty = false;
     if (auto singleVar = PBD->getSingleVar()) {
-      if (visitCustomAttributes(singleVar))
-        return true;
-
       isPropertyWrapperBackingProperty =
         singleVar->getOriginalWrappedProperty() != nullptr;
     }
@@ -1153,9 +1125,6 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       if (doIt(P))
         return true;
 
-      // Visit any custom attributes on the parameter.
-      visitCustomAttributes(P);
-
       // Don't walk into the type if the decl is implicit, or if the type is
       // implicit.
       if (!P->isImplicit() && !P->isTypeLocImplicit() &&
@@ -1168,7 +1137,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
         P->setDefaultValue(res);
       }
     }
-    
+
     return Walker.walkToParameterListPost(PL);
   }
 
@@ -1257,7 +1226,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
     return P;
   }
-  
+
   bool doIt(const StmtCondition &C) {
     for (auto &elt : C) {
       switch (elt.getKind()) {

lib/IDE/SourceEntityWalker.cpp

@@ -34,6 +34,7 @@ class SemaAnnotator : public ASTWalker {
   SmallVector<ConstructorRefCallExpr *, 2> CtorRefs;
   SmallVector<ExtensionDecl *, 2> ExtDecls;
   llvm::SmallDenseMap<OpaqueValueExpr *, Expr *, 4> OpaqueValueMap;
+  llvm::SmallPtrSet<Expr *, 16> ExprsToSkip;
   bool Cancelled = false;
   Optional<AccessKind> OpAccess;
 
@@ -61,6 +62,7 @@ class SemaAnnotator : public ASTWalker {
   std::pair<bool, Pattern *> walkToPatternPre(Pattern *P) override;
 
   bool handleImports(ImportDecl *Import);
+  bool handleCustomAttributes(Decl *D);
   bool passModulePathElements(ArrayRef<ImportDecl::AccessPathElement> Path,
                               const clang::Module *ClangMod);
 
@@ -95,6 +97,11 @@ bool SemaAnnotator::walkToDeclPre(Decl *D) {
   if (shouldIgnore(D, ShouldVisitChildren))
     return ShouldVisitChildren;
 
+  if (!handleCustomAttributes(D)) {
+    Cancelled = true;
+    return false;
+  }
+
   SourceLoc Loc = D->getLoc();
   unsigned NameLen = 0;
   bool IsExtension = false;
@@ -252,6 +259,9 @@ std::pair<bool, Expr *> SemaAnnotator::walkToExprPre(Expr *E) {
   if (isDone())
     return { false, nullptr };
 
+  if (ExprsToSkip.count(E) != 0)
+    return { false, E };
+
   if (!SEWalker.walkToExprPre(E))
     return { false, E };
 
@@ -541,6 +551,41 @@ std::pair<bool, Pattern *> SemaAnnotator::walkToPatternPre(Pattern *P) {
   return { false, P };
 }
 
+bool SemaAnnotator::handleCustomAttributes(Decl *D) {
+  // CustomAttrs of non-param VarDecls are handled when this method is called
+  // on their containing PatternBindingDecls (see below).
+  if (isa<VarDecl>(D) && !isa<ParamDecl>(D))
+    return true;
+
+  if (auto *PBD = dyn_cast<PatternBindingDecl>(D)) {
+    if (auto *SingleVar = PBD->getSingleVar()) {
+      D = SingleVar;
+    } else {
+      return true;
+    }
+  }
+  for (auto *customAttr : D->getAttrs().getAttributes<CustomAttr, true>()) {
+    if (auto *Repr = customAttr->getTypeLoc().getTypeRepr()) {
+      if (!Repr->walk(*this))
+        return false;
+    }
+    if (auto *SemaInit = customAttr->getSemanticInit()) {
+      if (!SemaInit->isImplicit()) {
+        assert(customAttr->getArg());
+        if (!SemaInit->walk(*this))
+          return false;
+        // Don't walk this again via the associated PatternBindingDecl's
+        // initializer
+        ExprsToSkip.insert(SemaInit);
+      }
+    } else if (auto *Arg = customAttr->getArg()) {
+      if (!Arg->walk(*this))
+        return false;
+    }
+  }
+  return true;
+}
+
 bool SemaAnnotator::handleImports(ImportDecl *Import) {
   auto Mod = Import->getModule();
   if (!Mod)

lib/IDE/SwiftSourceDocInfo.cpp

@@ -86,8 +86,10 @@ bool CursorInfoResolver::tryResolve(ValueDecl *D, TypeDecl *CtorTyRef,
     // Handle references to the implicitly generated vars in case statements
     // matching multiple patterns
     if (VD->isImplicit()) {
-      if (auto * Parent = VD->getParentVarDecl())
+      if (auto * Parent = VD->getParentVarDecl()) {
         D = Parent;
+        VD = Parent;
+      }
     }
   }
   CursorInfo.setValueRef(D, CtorTyRef, ExtTyRef, IsRef, Ty, ContainerType);
@@ -135,7 +137,7 @@ ResolvedCursorInfo CursorInfoResolver::resolve(SourceLoc Loc) {
 }
 
 bool CursorInfoResolver::walkToDeclPre(Decl *D, CharSourceRange Range) {
-  if (!rangeContainsLoc(D->getSourceRange()))
+  if (!rangeContainsLoc(D->getSourceRangeIncludingAttrs()))
     return false;
 
   if (isa<ExtensionDecl>(D))
@@ -354,6 +356,42 @@ static std::vector<CharSourceRange> getEnumParamListInfo(SourceManager &SM,
   return LabelRanges;
 }
 
+bool NameMatcher::handleCustomAttrs(Decl *D) {
+  // CustomAttrs of non-param VarDecls are handled when this method is called
+  // on their containing PatternBindingDecls (see below).
+  if (isa<VarDecl>(D) && !isa<ParamDecl>(D))
+    return true;
+
+  if (auto *PBD = dyn_cast<PatternBindingDecl>(D)) {
+    if (auto *SingleVar = PBD->getSingleVar()) {
+      D = SingleVar;
+    } else {
+        return true;
+    }
+  }
+
+  for (auto *customAttr : D->getAttrs().getAttributes<CustomAttr, true>()) {
+    if (shouldSkip(customAttr->getRangeWithAt()))
+      continue;
+    auto *Arg = customAttr->getArg();
+    if (auto *Repr = customAttr->getTypeLoc().getTypeRepr()) {
+      // Note the associated call arguments of the semantic initializer call
+      // in case we're resolving an explicit initializer call within the
+      // CustomAttr's type, e.g. on `Wrapper` in `@Wrapper(initialValue: 10)`.
+      SWIFT_DEFER { CustomAttrArg = None; };
+      if (Arg && !Arg->isImplicit())
+        CustomAttrArg = {Repr->getLoc(), Arg};
+      if (!Repr->walk(*this))
+        return false;
+    }
+    if (Arg && !Arg->isImplicit()) {
+      if (!Arg->walk(*this))
+        return false;
+    }
+  }
+  return !isDone();
+}
+
 bool NameMatcher::walkToDeclPre(Decl *D) {
   // Handle occurrences in any preceding doc comments
   RawComment R = D->getRawComment();
@@ -369,9 +407,12 @@ bool NameMatcher::walkToDeclPre(Decl *D) {
   if (D->isImplicit())
     return !isDone();
 
-  if (shouldSkip(D->getSourceRange()))
+  if (shouldSkip(D->getSourceRangeIncludingAttrs()))
     return false;
 
+  if (!handleCustomAttrs(D))
+    return false;
+
   if (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
     for (auto Clause : ICD->getClauses()) {
       if (!Clause.isActive)
@@ -575,8 +616,16 @@ bool NameMatcher::walkToTypeReprPre(TypeRepr *T) {
   if (isDone() || shouldSkip(T->getSourceRange()))
     return false;
 
-  if (isa<ComponentIdentTypeRepr>(T))
-    tryResolve(ASTWalker::ParentTy(T), T->getLoc());
+  if (isa<ComponentIdentTypeRepr>(T)) {
+    // If we're walking a CustomAttr's type we may have an associated call
+    // argument to resolve with from its semantic initializer.
+    if (CustomAttrArg.hasValue() && CustomAttrArg->first == T->getLoc()) {
+      tryResolve(ASTWalker::ParentTy(T), T->getLoc(), LabelRangeType::CallArg,
+                 getCallArgLabelRanges(getSourceMgr(), CustomAttrArg->second, LabelRangeEndAt::BeforeElemStart));
+    } else {
+      tryResolve(ASTWalker::ParentTy(T), T->getLoc());
+    }
+  }
   return !isDone();
 }
 
@@ -726,13 +775,32 @@ bool NameMatcher::tryResolve(ASTWalker::ParentTy Node, SourceLoc NameLoc,
   CharSourceRange Range = Lexer::getCharSourceRangeFromSourceRange(getSourceMgr(),
                                                                    NameLoc);
   UnresolvedLoc &Next = LocsToResolve.back();
-  if (Range.isValid() && NameLoc == Next.Loc) {
-    LocsToResolve.pop_back();
-    ResolvedLocs.push_back({Node, Range, LabelRanges, RangeType,
-      isActive(), isInSelector()});
-    return true;
+  bool WasResolved = false;
+  if (Range.isValid()) {
+    if (NameLoc == Next.Loc) {
+      LocsToResolve.pop_back();
+      ResolvedLocs.push_back({Node, Range, LabelRanges, RangeType,
+        isActive(), isInSelector()});
+      if (isDone())
+        return true;
+      WasResolved = true;
+    }
+
+    if (Range.getByteLength() > 1 &&
+        (Range.str().front() == '_' || Range.str().front() == '$')) {
+      // Also try after any leading _ or $ for name references of wrapped
+      // properties, e.g. 'foo' in '_foo' and '$foo' occurrences.
+      auto NewRange = CharSourceRange(Range.getStart().getAdvancedLoc(1),
+                                      Range.getByteLength() - 1);
+      if (NewRange.getStart() == Next.Loc) {
+        LocsToResolve.pop_back();
+        ResolvedLocs.push_back({Node, NewRange, {}, LabelRangeType::None,
+          isActive(), isInSelector()});
+        WasResolved = true;
+      }
+    }
   }
-  return false;
+  return WasResolved;
 };
 
 void ResolvedRangeInfo::print(llvm::raw_ostream &OS) {

lib/IDE/SyntaxModel.cpp

@@ -1026,22 +1026,28 @@ bool ModelASTWalker::handleSpecialDeclAttribute(const DeclAttribute *D,
                                                 ArrayRef<Token> Toks) {
   if (!D)
     return false;
-  if (isa<AvailableAttr>(D) || isa<CustomAttr>(D)) {
-    unsigned I = 0;
-    for (; I < TokenNodes.size(); ++ I) {
-      auto Node = TokenNodes[I];
-      if (SM.isBeforeInBuffer(D->getRange().End, Node.Range.getStart()))
-        break;
-      if (Node.Range.contains(D->AtLoc)) {
-        if (!passNode({SyntaxNodeKind::AttributeBuiltin, Node.Range}))
-          break;
-        continue;
+  if (isa<CustomAttr>(D) || isa<AvailableAttr>(D)) {
+    if (!passTokenNodesUntil(D->getRangeWithAt().Start,
+                        PassNodesBehavior::ExcludeNodeAtLocation))
+      return false;
+    if (auto *CA = dyn_cast<CustomAttr>(D)) {
+      if (auto *Repr = CA->getTypeLoc().getTypeRepr()) {
+        if (!Repr->walk(*this))
+          return false;
       }
-
-      if (!passNode(Node))
-        break;
+      if (auto *Arg = CA->getArg()) {
+        if (!Arg->walk(*this))
+          return false;
+      }
+    } else {
+      auto Next = TokenNodes.front();
+      TokenNodes = TokenNodes.drop_front();
+      assert(Next.Range.getStart() == D->getRangeWithAt().Start);
+      if (!passNode({SyntaxNodeKind::AttributeBuiltin, Next.Range}))
+        return false;
     }
-    TokenNodes = TokenNodes.slice(I);
+    if (!passTokenNodesUntil(D->getRange().End, PassNodesBehavior::IncludeNodeAtLocation))
+      return false;
     return true;
   }
   if (isa<RethrowsAttr>(D))
@@ -1078,6 +1084,13 @@ bool ModelASTWalker::handleAttrRanges(ArrayRef<DeclAttributeAndRange> DeclRanges
             [&](DeclAttributeAndRange LHS, DeclAttributeAndRange RHS) {
     return SM.isBeforeInBuffer(LHS.second.Start, RHS.second.End);
   });
+  // Handle duplicate synthesized attributes due to * in @available
+  auto NewEnd = std::unique(SortedRanges.begin(), SortedRanges.end(),
+                            [&](DeclAttributeAndRange LHS, DeclAttributeAndRange RHS) {
+                              return LHS.second == RHS.second;
+                            });
+  if (NewEnd != SortedRanges.end())
+    SortedRanges.erase(NewEnd, SortedRanges.end());
   DeclRanges = SortedRanges;
 
   SourceLoc BeginLoc = DeclRanges.front().second.Start;