[CodeCompletion] Add completion for @storageRestrictions

Author: ahoppen

include/swift/IDE/CompletionLookup.h

@@ -105,6 +105,9 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     TypeInDeclContext,
     ImportFromModule,
     GenericRequirement,
+
+    /// Look up stored properties within a type.
+    StoredProperty,
   };
 
   LookupKind Kind;
@@ -528,6 +531,9 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   void getValueExprCompletions(Type ExprType, ValueDecl *VD = nullptr,
                                bool IsDeclUnapplied = false);
 
+  /// Add completions for stored properties of \p D.
+  void getStoredPropertyCompletions(const NominalTypeDecl *D);
+
   void collectOperators(SmallVectorImpl<OperatorDecl *> &results);
 
   void addPostfixBang(Type resultType);

include/swift/Parse/IDEInspectionCallbacks.h

@@ -36,6 +36,25 @@ enum class CustomSyntaxAttributeKind {
   Available,
   FreestandingMacro,
   AttachedMacro,
+  StorageRestrictions
+};
+
+/// A bit of a hack. When completing inside the '@storageRestrictions'
+/// attribute, we use the \c ParamIndex field to communicate where inside the
+/// attribute we are performing the completion.
+enum class StorageRestrictionsCompletionKind : int {
+  /// We are completing directly after the '(' and require a 'initializes' or
+  /// 'accesses' label.
+  Label,
+  /// We are completing in a context that only allows arguments (ie. accessed or
+  /// initialized variables) and doesn't permit an argument label.
+  Argument,
+  /// Completion in a context that allows either an argument or the
+  /// 'initializes' label.
+  ArgumentOrInitializesLabel,
+  /// Completion in a context that allows either an argument or the
+  /// 'accesses' label.
+  ArgumentOrAccessesLabel
 };
 
 /// Parser's interface to code completion.

lib/IDE/CompletionLookup.cpp

@@ -1305,6 +1305,7 @@ bool CompletionLookup::isImplicitlyCurriedInstanceMethod(
 
   switch (Kind) {
   case LookupKind::ValueExpr:
+  case LookupKind::StoredProperty:
     return ExprType->is<AnyMetatypeType>();
   case LookupKind::ValueInDeclContext:
     if (InsideStaticMethod)
@@ -2166,6 +2167,13 @@ void CompletionLookup::foundDecl(ValueDecl *D, DeclVisibilityKind Reason,
     }
 
     return;
+  case LookupKind::StoredProperty:
+    if (auto *VD = dyn_cast<VarDecl>(D)) {
+      if (VD->getImplInfo().hasStorage()) {
+        addVarDeclRef(VD, Reason, dynamicLookupInfo);
+      }
+      return;
+    }
   }
 }
 
@@ -2381,6 +2389,17 @@ void CompletionLookup::getValueExprCompletions(Type ExprType, ValueDecl *VD,
                            /*includeProtocolExtensionMembers*/ true);
 }
 
+void CompletionLookup::getStoredPropertyCompletions(const NominalTypeDecl *D) {
+  Kind = LookupKind::StoredProperty;
+  NeedLeadingDot = false;
+
+  lookupVisibleMemberDecls(*this, D->getDeclaredInterfaceType(),
+                           /*DotLoc=*/SourceLoc(), CurrDeclContext,
+                           /*IncludeInstanceMembers*/ true,
+                           /*includeDerivedRequirements*/ false,
+                           /*includeProtocolExtensionMembers*/ false);
+}
+
 void CompletionLookup::collectOperators(
     SmallVectorImpl<OperatorDecl *> &results) {
   assert(CurrDeclContext);
@@ -2969,6 +2988,44 @@ void CompletionLookup::getAttributeDeclParamCompletions(
       break;
     }
     break;
+  case CustomSyntaxAttributeKind::StorageRestrictions:
+    if (!HasLabel) {
+      bool suggestInitializesLabel = false;
+      bool suggestAccessesLabel = false;
+      bool suggestArgument = false;
+      switch (static_cast<StorageRestrictionsCompletionKind>(ParamIndex)) {
+      case StorageRestrictionsCompletionKind::Label:
+        suggestAccessesLabel = true;
+        suggestInitializesLabel = true;
+        break;
+      case StorageRestrictionsCompletionKind::Argument:
+        suggestArgument = true;
+        break;
+      case StorageRestrictionsCompletionKind::ArgumentOrInitializesLabel:
+        suggestArgument = true;
+        suggestInitializesLabel = true;
+        break;
+      case StorageRestrictionsCompletionKind::ArgumentOrAccessesLabel:
+        suggestArgument = true;
+        suggestAccessesLabel = true;
+        break;
+      }
+      if (suggestInitializesLabel) {
+        addDeclAttrParamKeyword("initializes", /*Parameters=*/{},
+                                "Specify variables initialized by the accessor",
+                                true);
+      }
+      if (suggestAccessesLabel) {
+        addDeclAttrParamKeyword("accesses", /*Parameters=*/{},
+                                "Specify variables accessed by the accessor",
+                                true);
+      }
+      if (suggestArgument) {
+        if (auto NT = dyn_cast<NominalTypeDecl>(CurrDeclContext)) {
+          getStoredPropertyCompletions(NT);
+        }
+      }
+    }
   }
 }