[CodeCompletion] Mark types that can’t be used as attributes as having an invalid type relation when used after '@'

When completing after `@`, record what kind of attributes are applicable here (property wrapper, result builder, global actor), mark types that are marked as property wrapper etc. as having a 'Convertible' type relation and mark all other types as having an invalid type relation.

rdar://78239501

Author: ahoppen

include/swift/IDE/CodeCompletionResultType.h

@@ -21,6 +21,15 @@
 namespace swift {
 namespace ide {
 
+enum class CustomAttributeKind : uint8_t {
+  /// A type that can be used as a property wrapper.
+  PropertyWrapper = 1 << 0,
+  /// A type that can be used as a result builder.
+  ResultBuilder = 1 << 1,
+  /// A type that can be used as a global actor.
+  GlobalActor = 1 << 2,
+};
+
 /// The expected contextual type(s) for code-completion.
 class ExpectedTypeContext {
   /// Possible types of the code completion expression.
@@ -37,6 +46,11 @@ class ExpectedTypeContext {
   bool IsImplicitSingleExpressionReturn = false;
   bool PreferNonVoid = false;
 
+  /// If not empty, \c PossibleTypes are ignored and types that have an
+  /// attribute kind that is contained in this list will be reported as
+  /// 'Convertible'. All other types are related as 'Invalid'.
+  OptionSet<CustomAttributeKind> ExpectedCustomAttributeKinds;
+
 public:
   ExpectedTypeContext() = default;
 
@@ -82,6 +96,15 @@ class ExpectedTypeContext {
   void setPreferNonVoid(bool PreferNonVoid) {
     this->PreferNonVoid = PreferNonVoid;
   }
+
+  OptionSet<CustomAttributeKind> getExpectedCustomAttributeKinds() const {
+    return ExpectedCustomAttributeKinds;
+  }
+
+  void setExpectedCustomAttributeKinds(
+      OptionSet<CustomAttributeKind> ExpectedAttributeKinds) {
+    this->ExpectedCustomAttributeKinds = ExpectedAttributeKinds;
+  }
 };
 
 /// Describes the relationship between the type of the completion results and
@@ -174,6 +197,9 @@ class USRBasedTypeContext {
 
   SmallVector<ContextualType, 4> ContextualTypes;
 
+  /// See \c ExpectedTypeContext::ExpectedAttributeKinds.
+  OptionSet<CustomAttributeKind> ExpectedCustomAttributeKinds;
+
 public:
   /// Create a USR-based equivalent of the \p TypeContext.
   USRBasedTypeContext(const ExpectedTypeContext *TypeContext,
@@ -203,10 +229,15 @@ class USRBasedType {
   /// declared in the same module that the type was defined.
   ArrayRef<const USRBasedType *> Supertypes;
 
+  /// The kinds of attributes this type can be used as.
+  OptionSet<CustomAttributeKind> CustomAttributeKinds;
+
   /// Memberwise initializer. \p USR and \p Supertypes need to be allocated in
   /// the same arena as this \c USRBasedType.
-  USRBasedType(StringRef USR, ArrayRef<const USRBasedType *> Supertypes)
-      : USR(USR), Supertypes(Supertypes) {}
+  USRBasedType(StringRef USR, ArrayRef<const USRBasedType *> Supertypes,
+               OptionSet<CustomAttributeKind> CustomAttributeKinds)
+      : USR(USR), Supertypes(Supertypes),
+        CustomAttributeKinds(CustomAttributeKinds) {}
 
   /// Implementation of \c typeRelation. \p VisistedTypes keeps track which
   /// types have already been visited.
@@ -225,13 +256,18 @@ class USRBasedType {
 
   /// Construct as \c USRBasedType from a USR and its supertypes. This method
   /// takes care of copying \p USR and \p Supertypes to \p Arena.
-  static const USRBasedType *fromUSR(StringRef USR,
-                                     ArrayRef<const USRBasedType *> Supertypes,
-                                     USRBasedTypeArena &Arena);
+  static const USRBasedType *
+  fromUSR(StringRef USR, ArrayRef<const USRBasedType *> Supertypes,
+          OptionSet<CustomAttributeKind> CustomAttributeKinds,
+          USRBasedTypeArena &Arena);
 
   StringRef getUSR() const { return USR; }
   ArrayRef<const USRBasedType *> getSupertypes() const { return Supertypes; }
 
+  OptionSet<CustomAttributeKind> getCustomAttributeKinds() const {
+    return CustomAttributeKinds;
+  }
+
   /// Compute the relation of \c ResultType when to this contextual type.
   CodeCompletionResultTypeRelation
   typeRelation(const USRBasedType *ResultType,

include/swift/IDE/CompletionLookup.h

@@ -240,13 +240,16 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 
   void setIsStaticMetatype(bool value) { IsStaticMetatype = value; }
 
-  void setExpectedTypes(ArrayRef<Type> Types,
-                        bool isImplicitSingleExpressionReturn,
-                        bool preferNonVoid = false) {
+  void setExpectedTypes(
+      ArrayRef<Type> Types, bool isImplicitSingleExpressionReturn,
+      bool preferNonVoid = false,
+      OptionSet<CustomAttributeKind> expectedCustomAttributeKinds = {}) {
     expectedTypeContext.setIsImplicitSingleExpressionReturn(
         isImplicitSingleExpressionReturn);
     expectedTypeContext.setPreferNonVoid(preferNonVoid);
     expectedTypeContext.setPossibleTypes(Types);
+    expectedTypeContext.setExpectedCustomAttributeKinds(
+        expectedCustomAttributeKinds);
   }
 
   void setIdealExpectedType(Type Ty) { expectedTypeContext.setIdealType(Ty); }

lib/IDE/CodeCompletion.cpp

@@ -1701,6 +1701,35 @@ void CodeCompletionCallbacksImpl::doneParsing() {
 
   case CompletionKind::AttributeBegin: {
     Lookup.getAttributeDeclCompletions(IsInSil, AttTargetDK);
+    OptionSet<CustomAttributeKind> ExpectedCustomAttributeKinds;
+    if (AttTargetDK) {
+      switch (*AttTargetDK) {
+      case DeclKind::Var:
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::GlobalActor;
+        LLVM_FALLTHROUGH;
+      case DeclKind::Param:
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::ResultBuilder;
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::PropertyWrapper;
+        break;
+      case DeclKind::Func:
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::ResultBuilder;
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::GlobalActor;
+        break;
+      default:
+        break;
+      }
+    }
+    if (!ExpectedCustomAttributeKinds) {
+      // If we don't know on which decl kind we are completing, suggest all
+      // attribute kinds.
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::PropertyWrapper;
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::ResultBuilder;
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::GlobalActor;
+    }
+    Lookup.setExpectedTypes(/*Types=*/{},
+                            /*isImplicitSingleExpressionReturn=*/false,
+                            /*preferNonVoid=*/false,
+                            ExpectedCustomAttributeKinds);
 
     // TypeName at attribute position after '@'.
     // - VarDecl: Property Wrappers.

lib/IDE/CodeCompletionCache.cpp

@@ -103,7 +103,7 @@ CodeCompletionCache::~CodeCompletionCache() {}
 ///
 /// This should be incremented any time we commit a change to the format of the
 /// cached results. This isn't expected to change very often.
-static constexpr uint32_t onDiskCompletionCacheVersion = 6; // Store completion item result types in the cache
+static constexpr uint32_t onDiskCompletionCacheVersion = 7; // Store whether a type can be used as attribute
 
 /// Deserializes CodeCompletionResults from \p in and stores them in \p V.
 /// \see writeCacheModule.
@@ -193,8 +193,9 @@ static bool readCachedModule(llvm::MemoryBuffer *in,
       auto supertypeIndex = read32le(p);
       supertypes.push_back(getType(supertypeIndex));
     }
-    const USRBasedType *res =
-        USRBasedType::fromUSR(usr, supertypes, V.USRTypeArena);
+    auto customAttributeKinds = OptionSet<CustomAttributeKind, uint8_t>(*p++);
+    const USRBasedType *res = USRBasedType::fromUSR(
+        usr, supertypes, customAttributeKinds, V.USRTypeArena);
     knownTypes[index] = res;
     return res;
   };
@@ -381,6 +382,9 @@ static void writeCachedModule(llvm::raw_ostream &out,
     for (auto supertypeIndex : supertypeIndicies) {
       LE.write(static_cast<uint32_t>(supertypeIndex));
     }
+    OptionSet<CustomAttributeKind, uint8_t> customAttributeKinds =
+        type->getCustomAttributeKinds();
+    LE.write(static_cast<uint8_t>(customAttributeKinds.toRaw()));
     knownTypes[type] = size;
     return static_cast<uint32_t>(size);
   };

lib/IDE/CodeCompletionResultType.cpp

@@ -22,11 +22,32 @@ using namespace swift;
 using namespace ide;
 using TypeRelation = CodeCompletionResultTypeRelation;
 
+// MARK: - Utilities
+
+/// Returns the kind of attributes \c Ty can be used as.
+static OptionSet<CustomAttributeKind> getCustomAttributeKinds(Type Ty) {
+  OptionSet<CustomAttributeKind> Result;
+  if (auto NominalTy = Ty->getAs<NominalType>()) {
+    auto NominalDecl = NominalTy->getDecl();
+    if (NominalDecl->getAttrs().hasAttribute<PropertyWrapperAttr>()) {
+      Result |= CustomAttributeKind::PropertyWrapper;
+    }
+    if (NominalDecl->getAttrs().hasAttribute<ResultBuilderAttr>()) {
+      Result |= CustomAttributeKind::ResultBuilder;
+    }
+    if (NominalDecl->isGlobalActor()) {
+      Result |= CustomAttributeKind::GlobalActor;
+    }
+  }
+  return Result;
+}
+
 // MARK: - USRBasedTypeContext
 
 USRBasedTypeContext::USRBasedTypeContext(const ExpectedTypeContext *TypeContext,
                                          USRBasedTypeArena &Arena)
-    : Arena(Arena) {
+    : Arena(Arena), ExpectedCustomAttributeKinds(
+                        TypeContext->getExpectedCustomAttributeKinds()) {
 
   for (auto possibleTy : TypeContext->getPossibleTypes()) {
     ContextualTypes.emplace_back(USRBasedType::fromType(possibleTy, Arena),
@@ -65,6 +86,11 @@ USRBasedTypeContext::USRBasedTypeContext(const ExpectedTypeContext *TypeContext,
 
 TypeRelation
 USRBasedTypeContext::typeRelation(const USRBasedType *ResultType) const {
+  if (ExpectedCustomAttributeKinds) {
+    return ResultType->getCustomAttributeKinds() & ExpectedCustomAttributeKinds
+               ? TypeRelation::Convertible
+               : TypeRelation::Unrelated;
+  }
   const USRBasedType *VoidType = Arena.getVoidType();
   if (ResultType == VoidType) {
     // Void is not convertible to anything and we don't report Void <-> Void
@@ -84,7 +110,7 @@ USRBasedTypeContext::typeRelation(const USRBasedType *ResultType) const {
 
 USRBasedTypeArena::USRBasedTypeArena() {
   // '$sytD' is the USR of the Void type.
-  VoidType = USRBasedType::fromUSR("$sytD", {}, *this);
+  VoidType = USRBasedType::fromUSR("$sytD", {}, {}, *this);
 }
 
 const USRBasedType *USRBasedTypeArena::getVoidType() const { return VoidType; }
@@ -124,11 +150,12 @@ TypeRelation USRBasedType::typeRelationImpl(
 }
 
 const USRBasedType *USRBasedType::null(USRBasedTypeArena &Arena) {
-  return USRBasedType::fromUSR(/*USR=*/"", /*Supertypes=*/{}, Arena);
+  return USRBasedType::fromUSR(/*USR=*/"", /*Supertypes=*/{}, {}, Arena);
 }
 
 const USRBasedType *
 USRBasedType::fromUSR(StringRef USR, ArrayRef<const USRBasedType *> Supertypes,
+                      OptionSet<CustomAttributeKind> CustomAttributeKinds,
                       USRBasedTypeArena &Arena) {
   auto ExistingTypeIt = Arena.CanonicalTypes.find(USR);
   if (ExistingTypeIt != Arena.CanonicalTypes.end()) {
@@ -141,7 +168,7 @@ USRBasedType::fromUSR(StringRef USR, ArrayRef<const USRBasedType *> Supertypes,
   Supertypes = Supertypes.copy(Arena.Allocator);
 
   const USRBasedType *Result =
-      new (Arena.Allocator) USRBasedType(USR, Supertypes);
+      new (Arena.Allocator) USRBasedType(USR, Supertypes, CustomAttributeKinds);
   Arena.CanonicalTypes[USR] = Result;
   return Result;
 }
@@ -222,7 +249,8 @@ const USRBasedType *USRBasedType::fromType(Type Ty, USRBasedTypeArena &Arena) {
     return ImpliedSupertypes.contains(Ty);
   });
 
-  return USRBasedType::fromUSR(USR, Supertypes, Arena);
+  return USRBasedType::fromUSR(USR, Supertypes, ::getCustomAttributeKinds(Ty),
+                               Arena);
 }
 
 TypeRelation USRBasedType::typeRelation(const USRBasedType *ResultType,
@@ -284,6 +312,12 @@ calculateMaxTypeRelation(Type Ty, const ExpectedTypeContext &typeContext,
                          const DeclContext &DC) {
   if (Ty->isVoid() && typeContext.requiresNonVoid())
     return TypeRelation::Invalid;
+  if (typeContext.getExpectedCustomAttributeKinds()) {
+    return (getCustomAttributeKinds(Ty) &
+            typeContext.getExpectedCustomAttributeKinds())
+               ? TypeRelation::Convertible
+               : TypeRelation::Unrelated;
+  }
   if (typeContext.empty())
     return TypeRelation::Unknown;