[CodeCompletion] Support type checking attributes even if they are not part of the AST

The code completion might occur inside an attriubte that isn’t part of the AST because it’s missing a `VarDecl` that it could be attached to. In these cases, record the `CustomAttr` and type check it standalone, pretending it was part of a `DeclContext`.

This also fixes a few issues where code completion previously wouldn’t find the attribute constructor call and thus wasn’t providing code completion inside the property wrapper.

rdar://92842803

Author: ahoppen

include/swift/AST/ASTNode.h

@@ -78,6 +78,10 @@ namespace swift {
 
     /// Whether the AST node is implicit.
     bool isImplicit() const;
+
+    friend llvm::hash_code hash_value(ASTNode N) {
+      return llvm::hash_value(N.getOpaqueValue());
+    }
   };
 } // namespace swift
 

include/swift/AST/TypeCheckRequests.h

@@ -18,6 +18,7 @@
 
 #include "swift/AST/ActorIsolation.h"
 #include "swift/AST/AnyFunctionRef.h"
+#include "swift/AST/ASTNode.h"
 #include "swift/AST/ASTTypeIDs.h"
 #include "swift/AST/Effects.h"
 #include "swift/AST/GenericParamList.h"
@@ -1544,12 +1545,82 @@ class TypeCheckFunctionBodyRequest
   readDependencySource(const evaluator::DependencyRecorder &) const;
 };
 
+/// Describes the context in which the AST node to type check in a
+/// \c TypeCheckASTNodeAtLocRequest should be searched. This can be either of
+/// two cases:
+///  1. A \c DeclContext that contains the node representing the location to
+///     type check
+///  2. If the node that should be type checked that might not be part of the
+///     AST (e.g. because it is a dangling property attribute), an \c ASTNode
+///     that contains the location to type check in together with a DeclContext
+///     in which we should pretend that node occurs.
+class TypeCheckASTNodeAtLocContext {
+  DeclContext *DC;
+  ASTNode Node;
+
+  /// Memberwise initializer
+  TypeCheckASTNodeAtLocContext(DeclContext *DC, ASTNode Node)
+      : DC(DC), Node(Node) {
+    assert(DC != nullptr);
+  }
+
+public:
+  static TypeCheckASTNodeAtLocContext declContext(DeclContext *DC) {
+    return TypeCheckASTNodeAtLocContext(DC, /*Node=*/nullptr);
+  }
+
+  static TypeCheckASTNodeAtLocContext node(DeclContext *DC, ASTNode Node) {
+    assert(!Node.isNull());
+    return TypeCheckASTNodeAtLocContext(DC, Node);
+  }
+
+  DeclContext *getDeclContext() const { return DC; }
+
+  bool isForUnattachedNode() const { return !Node.isNull(); }
+
+  ASTNode getUnattachedNode() const {
+    assert(isForUnattachedNode());
+    return Node;
+  }
+
+  ASTNode &getUnattachedNode() {
+    assert(isForUnattachedNode());
+    return Node;
+  }
+
+  friend llvm::hash_code hash_value(const TypeCheckASTNodeAtLocContext &ctx) {
+    return llvm::hash_combine(ctx.DC, ctx.Node);
+  }
+
+  friend bool operator==(const TypeCheckASTNodeAtLocContext &lhs,
+                         const TypeCheckASTNodeAtLocContext &rhs) {
+    return lhs.DC == rhs.DC && lhs.Node == rhs.Node;
+  }
+
+  friend bool operator!=(const TypeCheckASTNodeAtLocContext &lhs,
+                         const TypeCheckASTNodeAtLocContext &rhs) {
+    return !(lhs == rhs);
+  }
+
+  friend SourceLoc
+  extractNearestSourceLoc(const TypeCheckASTNodeAtLocContext &ctx) {
+    if (!ctx.Node.isNull()) {
+      return ctx.Node.getStartLoc();
+    } else {
+      return extractNearestSourceLoc(ctx.DC);
+    }
+  }
+};
+
+void simple_display(llvm::raw_ostream &out,
+                    const TypeCheckASTNodeAtLocContext &ctx);
+
 /// Request to typecheck a function body element at the given source location.
 ///
 /// Produces true if an error occurred, false otherwise.
 class TypeCheckASTNodeAtLocRequest
     : public SimpleRequest<TypeCheckASTNodeAtLocRequest,
-                           bool(DeclContext *, SourceLoc),
+                           bool(TypeCheckASTNodeAtLocContext, SourceLoc),
                            RequestFlags::Uncached> {
 public:
   using SimpleRequest::SimpleRequest;
@@ -1558,7 +1629,8 @@ class TypeCheckASTNodeAtLocRequest
   friend SimpleRequest;
 
   // Evaluation.
-  bool evaluate(Evaluator &evaluator, DeclContext *DC, SourceLoc Loc) const;
+  bool evaluate(Evaluator &evaluator, TypeCheckASTNodeAtLocContext,
+                SourceLoc Loc) const;
 };
 
 /// Request to obtain a list of stored properties in a nominal type.
@@ -3605,6 +3677,7 @@ class GetSourceFileAsyncNode
   bool isCached() const { return true; }
 };
 
+void simple_display(llvm::raw_ostream &out, ASTNode node);
 void simple_display(llvm::raw_ostream &out, Type value);
 void simple_display(llvm::raw_ostream &out, const TypeRepr *TyR);
 void simple_display(llvm::raw_ostream &out, ImplicitMemberAction action);

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -345,7 +345,7 @@ SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               BraceStmt *(AbstractFunctionDecl *), SeparatelyCached,
               NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckASTNodeAtLocRequest,
-              bool(DeclContext *, SourceLoc),
+              bool(TypeCheckASTNodeAtLocContext, SourceLoc),
               Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, UnderlyingTypeRequest, Type(TypeAliasDecl *),
               SeparatelyCached, NoLocationInfo)

include/swift/Parse/CodeCompletionCallbacks.h

@@ -118,6 +118,11 @@ class CodeCompletionCallbacks {
   /// Set target decl for attribute if the CC token is in attribute of the decl.
   virtual void setAttrTargetDeclKind(Optional<DeclKind> DK) {}
 
+  /// Set that the code completion token occurred in a custom attribute. This
+  /// allows us to type check the custom attribute even if it is not attached to
+  /// the AST, e.g. because there is no var declaration following it.
+  virtual void setCompletingInAttribute(CustomAttr *Attr){};
+
   /// Complete expr-dot after we have consumed the dot.
   virtual void completeDotExpr(CodeCompletionExpr *E, SourceLoc DotLoc) {};
 

include/swift/Sema/IDETypeChecking.h

@@ -19,7 +19,9 @@
 #ifndef SWIFT_SEMA_IDETYPECHECKING_H
 #define SWIFT_SEMA_IDETYPECHECKING_H
 
+#include "swift/AST/ASTNode.h"
 #include "swift/AST/Identifier.h"
+#include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/SourceLoc.h"
 #include <memory>
 #include <tuple>
@@ -140,8 +142,9 @@ namespace swift {
   /// Typecheck the given expression.
   bool typeCheckExpression(DeclContext *DC, Expr *&parsedExpr);
 
-  /// Type check a function body element which is at \p TargetLoc .
-  bool typeCheckASTNodeAtLoc(DeclContext *DC, SourceLoc TargetLoc);
+  /// Type check a function body element which is at \p TagetLoc.
+  bool typeCheckASTNodeAtLoc(TypeCheckASTNodeAtLocContext TypeCheckCtx,
+                             SourceLoc TargetLoc);
 
   /// Thunk around \c TypeChecker::typeCheckForCodeCompletion to make it
   /// available to \c swift::ide.

lib/AST/TypeCheckRequests.cpp

@@ -62,6 +62,14 @@ void swift::simple_display(llvm::raw_ostream &out,
   }
 }
 
+void swift::simple_display(llvm::raw_ostream &out, ASTNode node) {
+  if (node) {
+    node.dump(out);
+  } else {
+    out << "null";
+  }
+}
+
 void swift::simple_display(llvm::raw_ostream &out, Type type) {
   if (type)
     type.print(out);

lib/IDE/CodeCompletion.cpp

@@ -115,6 +115,12 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
   DeclContext *CurDeclContext = nullptr;
   DeclAttrKind AttrKind;
 
+  /// When the code completion token occurs in a custom attribute, the attribute
+  /// it occurs in. Used so we can complete inside the attribute even if it's
+  /// not attached to the AST, e.g. because there is no var decl it could be
+  /// attached to.
+  CustomAttr *AttrWithCompletion = nullptr;
+
   /// In situations when \c SyntaxKind hints or determines
   /// completions, i.e. a precedence group attribute, this
   /// can be set and used to control the code completion scenario.
@@ -227,6 +233,11 @@ class CodeCompletionCallbacksImpl : public CodeCompletionCallbacks {
       AttTargetDK = DK;
   }
 
+  void setCompletingInAttribute(CustomAttr *Attr) override {
+    AttrWithCompletion = Attr;
+    CurDeclContext = P.CurDeclContext;
+  }
+
   void completeDotExpr(CodeCompletionExpr *E, SourceLoc DotLoc) override;
   void completeStmtOrExpr(CodeCompletionExpr *E) override;
   void completePostfixExprBeginning(CodeCompletionExpr *E) override;
@@ -1330,16 +1341,24 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     ? ParsedExpr->getLoc()
     : CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc();
 
-  switch (Kind) {
-  case CompletionKind::DotExpr: {
-    assert(CodeCompleteTokenExpr);
-    assert(CurDeclContext);
-
-    DotExprTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
-                                              CurDeclContext);
+  auto typeCheckWithLookup = [this, &CompletionLoc](
+                                 TypeCheckCompletionCallback &Lookup) {
     llvm::SaveAndRestore<TypeCheckCompletionCallback*>
       CompletionCollector(Context.CompletionCallback, &Lookup);
-    typeCheckContextAt(CurDeclContext, CompletionLoc);
+    if (AttrWithCompletion) {
+      /// The attribute might not be attached to the AST if there is no var decl
+      /// it could be attached to. Type check it standalone.
+      ASTNode Call = CallExpr::create(
+          CurDeclContext->getASTContext(), AttrWithCompletion->getTypeExpr(),
+          AttrWithCompletion->getArgs(), /*implicit=*/true);
+      typeCheckContextAt(
+          TypeCheckASTNodeAtLocContext::node(CurDeclContext, Call),
+          CompletionLoc);
+    } else {
+      typeCheckContextAt(
+          TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
+          CompletionLoc);
+    }
 
     // This (hopefully) only happens in cases where the expression isn't
     // typechecked during normal compilation either (e.g. member completion in a
@@ -1348,6 +1367,16 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     // tooling in general though.
     if (!Lookup.gotCallback())
       Lookup.fallbackTypeCheck(CurDeclContext);
+  };
+
+  switch (Kind) {
+  case CompletionKind::DotExpr: {
+    assert(CodeCompleteTokenExpr);
+    assert(CurDeclContext);
+
+    DotExprTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
+                                              CurDeclContext);
+    typeCheckWithLookup(Lookup);
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
 
@@ -1362,12 +1391,7 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 
     UnresolvedMemberTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
                                                        CurDeclContext);
-    llvm::SaveAndRestore<TypeCheckCompletionCallback*>
-      CompletionCollector(Context.CompletionCallback, &Lookup);
-    typeCheckContextAt(CurDeclContext, CompletionLoc);
-
-    if (!Lookup.gotCallback())
-      Lookup.fallbackTypeCheck(CurDeclContext);
+    typeCheckWithLookup(Lookup);
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
     Lookup.deliverResults(CurDeclContext, DotLoc, CompletionContext, Consumer);
@@ -1380,9 +1404,7 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     // so we can safely cast the \c ParsedExpr back to a \c KeyPathExpr.
     auto KeyPath = cast<KeyPathExpr>(ParsedExpr);
     KeyPathTypeCheckCompletionCallback Lookup(KeyPath);
-    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
-        Context.CompletionCallback, &Lookup);
-    typeCheckContextAt(CurDeclContext, CompletionLoc);
+    typeCheckWithLookup(Lookup);
 
     Lookup.deliverResults(CurDeclContext, DotLoc, CompletionContext, Consumer);
     return true;
@@ -1392,13 +1414,7 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
     assert(CurDeclContext);
     ArgumentTypeCheckCompletionCallback Lookup(CodeCompleteTokenExpr,
                                                CurDeclContext);
-    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
-        Context.CompletionCallback, &Lookup);
-    typeCheckContextAt(CurDeclContext, CompletionLoc);
-
-    if (!Lookup.gotCallback()) {
-      Lookup.fallbackTypeCheck(CurDeclContext);
-    }
+    typeCheckWithLookup(Lookup);
 
     Lookup.deliverResults(ShouldCompleteCallPatternAfterParen, CompletionLoc,
                           CurDeclContext, CompletionContext, Consumer);
@@ -1425,13 +1441,7 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
       // need to have a TypeCheckCompletionCallback so we can call
       // deliverResults on it to deliver the keyword results from the completion
       // context's result sink to the consumer.
-      llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
-          Context.CompletionCallback, &Lookup);
-      typeCheckContextAt(CurDeclContext, CompletionLoc);
-
-      if (!Lookup.gotCallback()) {
-        Lookup.fallbackTypeCheck(CurDeclContext);
-      }
+      typeCheckWithLookup(Lookup);
     }
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
@@ -1446,13 +1456,7 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 
     AfterPoundExprCompletion Lookup(CodeCompleteTokenExpr, CurDeclContext,
                                     ParentStmtKind);
-    llvm::SaveAndRestore<TypeCheckCompletionCallback *> CompletionCollector(
-        Context.CompletionCallback, &Lookup);
-    typeCheckContextAt(CurDeclContext, CompletionLoc);
-
-    if (!Lookup.gotCallback()) {
-      Lookup.fallbackTypeCheck(CurDeclContext);
-    }
+    typeCheckWithLookup(Lookup);
 
     addKeywords(CompletionContext.getResultSink(), MaybeFuncBody);
 
@@ -1522,7 +1526,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   undoSingleExpressionReturn(CurDeclContext);
   typeCheckContextAt(
-      CurDeclContext,
+      TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
       ParsedExpr
           ? ParsedExpr->getLoc()
           : CurDeclContext->getASTContext().SourceMgr.getCodeCompletionLoc());

lib/IDE/ConformingMethodList.cpp

@@ -68,7 +68,8 @@ void ConformingMethodListCallbacks::doneParsing() {
   if (!ParsedExpr)
     return;
 
-  typeCheckContextAt(CurDeclContext, ParsedExpr->getLoc());
+  typeCheckContextAt(TypeCheckASTNodeAtLocContext::declContext(CurDeclContext),
+                     ParsedExpr->getLoc());
 
   Type T = ParsedExpr->getType();