[Completion] Only provide macro completions when they are valid

Only return macros that are valid in their current position, ie. an
attached macro is not valid on a nominal.

Also return freestanding expression macros in code block item position
and handle the new freestanding code item macros.

Resolves rdar://105563583.

Author: bnbarham

include/swift/IDE/CodeCompletionCache.h

@@ -116,9 +116,7 @@ class OnDiskCodeCompletionCache {
 struct RequestedCachedModule {
   CodeCompletionCache::Key Key;
   const ModuleDecl *TheModule;
-  bool OnlyTypes;
-  bool OnlyPrecedenceGroups;
-  bool OnlyMacros;
+  CodeCompletionFilter Filter;
 };
 
 } // end namespace ide

include/swift/IDE/CodeCompletionResult.h

@@ -316,6 +316,39 @@ enum class CodeCompletionResultKind : uint8_t {
   MAX_VALUE = BuiltinOperator
 };
 
+enum class CodeCompletionMacroRole : uint8_t {
+  Expression = 1 << 0,
+  Declaration = 1 << 1,
+  CodeItem = 1 << 2,
+  AttachedVar = 1 << 3,
+  AttachedContext = 1 << 4,
+  AttachedDecl = 1 << 5,
+};
+using CodeCompletionMacroRoles = OptionSet<CodeCompletionMacroRole>;
+
+enum class CodeCompletionFilterFlag : uint16_t {
+  Expr = 1 << 0,
+  Type = 1 << 1,
+  PrecedenceGroup = 1 << 2,
+  Module = 1 << 3,
+  ExpressionMacro = 1 << 4,
+  DeclarationMacro = 1 << 5,
+  CodeItemMacro = 1 << 6,
+  AttachedVarMacro = 1 << 7,
+  AttachedContextMacro = 1 << 8,
+  AttachedDeclMacro = 1 << 9,
+};
+using CodeCompletionFilter = OptionSet<CodeCompletionFilterFlag>;
+
+CodeCompletionMacroRoles getCompletionMacroRoles(const Decl *D);
+
+CodeCompletionMacroRoles
+getCompletionMacroRoles(OptionSet<CustomAttributeKind> kinds);
+
+CodeCompletionMacroRoles getCompletionMacroRoles(CodeCompletionFilter filter);
+
+CodeCompletionFilter getCompletionFilter(CodeCompletionMacroRoles roles);
+
 /// The parts of a \c CodeCompletionResult that are not dependent on the context
 /// it appears in and can thus be cached.
 class ContextFreeCodeCompletionResult {
@@ -334,6 +367,8 @@ class ContextFreeCodeCompletionResult {
   CodeCompletionOperatorKind KnownOperatorKind : 6;
   static_assert(int(CodeCompletionOperatorKind::MAX_VALUE) < 1 << 6, "");
 
+  CodeCompletionMacroRoles DeclMacroRoles;
+
   bool IsSystem : 1;
   bool IsAsync : 1;
   /// Whether the result has been annotated as having an async alternative that
@@ -359,17 +394,18 @@ class ContextFreeCodeCompletionResult {
   NullTerminatedStringRef NameForDiagnostics;
 
 public:
-  /// Memberwise initializer. \p AssociatedKInd is opaque and will be
+  /// Memberwise initializer. \p AssociatedKind is opaque and will be
   /// interpreted based on \p Kind. If \p KnownOperatorKind is \c None and the
   /// completion item is an operator, it will be determined based on the
-  /// compleiton string.
+  /// completion string.
   ///
-  /// \note The caller must ensure that the \p CompleitonString and all the
+  /// \note The caller must ensure that the \p CompletionString and all the
   /// \c Ref types outlive this result, typically by storing them in the same
   /// \c CodeCompletionResultSink as the result itself.
   ContextFreeCodeCompletionResult(
       CodeCompletionResultKind Kind, uint8_t AssociatedKind,
-      CodeCompletionOperatorKind KnownOperatorKind, bool IsSystem, bool IsAsync,
+      CodeCompletionOperatorKind KnownOperatorKind,
+      CodeCompletionMacroRoles DeclMacroRoles, bool IsSystem, bool IsAsync,
       bool HasAsyncAlternative, CodeCompletionString *CompletionString,
       NullTerminatedStringRef ModuleName,
       NullTerminatedStringRef BriefDocComment,
@@ -380,8 +416,9 @@ class ContextFreeCodeCompletionResult {
       NullTerminatedStringRef DiagnosticMessage,
       NullTerminatedStringRef FilterName,
       NullTerminatedStringRef NameForDiagnostics)
-      : Kind(Kind), KnownOperatorKind(KnownOperatorKind), IsSystem(IsSystem),
-        IsAsync(IsAsync), HasAsyncAlternative(HasAsyncAlternative),
+      : Kind(Kind), KnownOperatorKind(KnownOperatorKind),
+        DeclMacroRoles(DeclMacroRoles), IsSystem(IsSystem), IsAsync(IsAsync),
+        HasAsyncAlternative(HasAsyncAlternative),
         CompletionString(CompletionString), ModuleName(ModuleName),
         BriefDocComment(BriefDocComment), AssociatedUSRs(AssociatedUSRs),
         ResultType(ResultType), NotRecommended(NotRecommended),
@@ -488,6 +525,10 @@ class ContextFreeCodeCompletionResult {
     return KnownOperatorKind;
   }
 
+  CodeCompletionMacroRoles getMacroRoles() const {
+    return DeclMacroRoles;
+  }
+
   bool isSystem() const { return IsSystem; };
 
   bool isAsync() const { return IsAsync; };

include/swift/IDE/CodeCompletionResultType.h

@@ -28,8 +28,12 @@ enum class CustomAttributeKind : uint8_t {
   ResultBuilder = 1 << 1,
   /// A type that can be used as a global actor.
   GlobalActor = 1 << 2,
-  /// A macro can be used as a custom attribute.
-  Macro = 1 << 3,
+  /// A macro that can be used on variables or subscripts.
+  VarMacro = 1 << 3,
+  /// A macro that can be used on any type context.
+  ContextMacro = 1 << 4,
+  /// A macro that can be used on any declaration.
+  DeclMacro = 1 << 5,
 };
 
 /// The expected contextual type(s) for code-completion.

include/swift/IDE/CompletionLookup.h

@@ -137,7 +137,6 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   bool IsUnwrappedOptional = false;
   SourceLoc DotLoc;
   bool NeedLeadingDot = false;
-  bool NeedLeadingMacroPound = false;
 
   bool NeedOptionalUnwrap = false;
   unsigned NumBytesToEraseForOptionalUnwrap = 0;
@@ -186,53 +185,27 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 public:
   struct RequestedResultsTy {
     const ModuleDecl *TheModule;
-    bool OnlyTypes;
-    bool OnlyPrecedenceGroups;
-    bool OnlyMacros;
+    CodeCompletionFilter Filter;
     bool NeedLeadingDot;
-    bool NeedPound;
-    bool IncludeModuleQualifier;
 
-    static RequestedResultsTy fromModule(const ModuleDecl *TheModule) {
-      return {TheModule, false, false, false, false, false, true};
+    static RequestedResultsTy fromModule(const ModuleDecl *mod,
+                                         CodeCompletionFilter filter) {
+      return {mod, filter, /*NeedLeadingDot=*/false};
     }
 
-    RequestedResultsTy onlyTypes() const {
-      return {TheModule, true, false, false, NeedLeadingDot, false,
-              IncludeModuleQualifier};
+    static RequestedResultsTy topLevelResults(CodeCompletionFilter filter) {
+      return {nullptr, filter, /*NeedLeadingDot=*/false};
     }
 
-    RequestedResultsTy onlyPrecedenceGroups() const {
-      assert(!OnlyTypes && "onlyTypes() already includes precedence groups");
-      return {TheModule, false, true, false, false, false, true};
-    }
-
-    RequestedResultsTy onlyMacros(bool needPound) const {
-      return {TheModule, false, false, true, false, needPound, false};
-    }
-
-    RequestedResultsTy needLeadingDot(bool NeedDot) const {
-      return {TheModule, OnlyTypes, OnlyPrecedenceGroups, OnlyMacros, NeedDot,
-              NeedPound, IncludeModuleQualifier};
-    }
-
-    RequestedResultsTy withModuleQualifier(bool IncludeModule) const {
-      return {TheModule, OnlyTypes, OnlyPrecedenceGroups, OnlyMacros,
-              NeedLeadingDot, NeedPound, IncludeModule};
-    }
-
-    static RequestedResultsTy toplevelResults() {
-      return {nullptr, false, false, false, false, true, true};
+    RequestedResultsTy needLeadingDot(bool needDot) const {
+      return {TheModule, Filter, needDot};
     }
 
     friend bool operator==(const RequestedResultsTy &LHS,
                            const RequestedResultsTy &RHS) {
-      return LHS.TheModule == RHS.TheModule && LHS.OnlyTypes == RHS.OnlyTypes &&
-             LHS.OnlyPrecedenceGroups == RHS.OnlyPrecedenceGroups &&
-             LHS.OnlyMacros == RHS.OnlyMacros &&
-             LHS.NeedLeadingDot == RHS.NeedLeadingDot &&
-             LHS.NeedPound == RHS.NeedPound &&
-             LHS.IncludeModuleQualifier == RHS.IncludeModuleQualifier;
+      return LHS.TheModule == RHS.TheModule &&
+             LHS.Filter.containsOnly(RHS.Filter) &&
+             LHS.NeedLeadingDot == RHS.NeedLeadingDot;
     }
   };
 
@@ -524,7 +497,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
       Type ExprType, const ValueDecl *VD,
       Optional<SemanticContextKind> SemanticContext = None);
 
-  bool tryModuleCompletions(Type ExprType, bool TypesOnly = false);
+  bool tryModuleCompletions(Type ExprType, CodeCompletionFilter Filter);
 
   /// If the given ExprType is optional, this adds completions for the unwrapped
   /// type.
@@ -562,7 +535,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
 
   void addObjCPoundKeywordCompletions(bool needPound);
 
-  void getMacroCompletions(bool needPound);
+  void getMacroCompletions(CodeCompletionMacroRoles roles);
 
   struct FilteredDeclConsumer : public swift::VisibleDeclConsumer {
     swift::VisibleDeclConsumer &Consumer;
@@ -626,7 +599,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
   void getTypeCompletionsInDeclContext(SourceLoc Loc,
                                        bool ModuleQualifier = true);
 
-  void getToplevelCompletions(bool OnlyTypes, bool OnlyMacros);
+  void getToplevelCompletions(CodeCompletionFilter Filter);
 
   void lookupExternalModuleDecls(const ModuleDecl *TheModule,
                                  ArrayRef<std::string> AccessPath,
@@ -645,18 +618,15 @@ using RequestedResultsTy = swift::ide::CompletionLookup::RequestedResultsTy;
 template <>
 struct DenseMapInfo<RequestedResultsTy> {
   static inline RequestedResultsTy getEmptyKey() {
-    return {DenseMapInfo<swift::ModuleDecl *>::getEmptyKey(), false, false,
-            false, false, false, false};
+    return {DenseMapInfo<swift::ModuleDecl *>::getEmptyKey(), {}, false};
   }
   static inline RequestedResultsTy getTombstoneKey() {
-    return {DenseMapInfo<swift::ModuleDecl *>::getTombstoneKey(), false, false,
-            false, false, false, false};
+    return {DenseMapInfo<swift::ModuleDecl *>::getTombstoneKey(), {}, false};
   }
   static unsigned getHashValue(const RequestedResultsTy &Val) {
     return hash_combine(
         DenseMapInfo<swift::ModuleDecl *>::getHashValue(Val.TheModule),
-        Val.OnlyTypes, Val.OnlyPrecedenceGroups, Val.OnlyMacros,
-        Val.NeedLeadingDot, Val.NeedPound, Val.IncludeModuleQualifier);
+        Val.Filter.toRaw(), Val.NeedLeadingDot);
   }
   static bool isEqual(const RequestedResultsTy &LHS,
                       const RequestedResultsTy &RHS) {

lib/IDE/AfterPoundExprCompletion.cpp

@@ -52,7 +52,7 @@ void AfterPoundExprCompletion::deliverResults(
                             /*expectsNonVoid=*/true);
     Lookup.addPoundAvailable(ParentStmtKind);
     Lookup.addObjCPoundKeywordCompletions(/*needPound=*/false);
-    Lookup.getMacroCompletions(/*needPound=*/false);
+    Lookup.getMacroCompletions(CodeCompletionMacroRole::Expression);
   }
 
   deliverCompletionResults(CompletionCtx, Lookup, DC, Consumer);

lib/IDE/CodeCompletion.cpp

@@ -1130,22 +1130,6 @@ static void addPoundDirectives(CodeCompletionResultSink &Sink) {
     Builder.addSimpleTypedParameter("Int");
     Builder.addRightParen();
   });
-  addWithName("warning", CodeCompletionKeywordKind::pound_warning,
-              [&] (CodeCompletionResultBuilder &Builder) {
-    Builder.addLeftParen();
-    Builder.addTextChunk("\"");
-    Builder.addSimpleNamedParameter("message");
-    Builder.addTextChunk("\"");
-    Builder.addRightParen();
-  });
-  addWithName("error", CodeCompletionKeywordKind::pound_error,
-              [&] (CodeCompletionResultBuilder &Builder) {
-    Builder.addLeftParen();
-    Builder.addTextChunk("\"");
-    Builder.addSimpleNamedParameter("message");
-    Builder.addTextChunk("\"");
-    Builder.addRightParen();
-  });
 
   addWithName("if ", CodeCompletionKeywordKind::pound_if,
               [&] (CodeCompletionResultBuilder &Builder) {
@@ -1352,11 +1336,10 @@ void swift::ide::deliverCompletionResults(
         std::pair<PairType, bool> Result = ImportsSeen.insert(K);
         if (!Result.second)
           return; // already handled.
-        RequestedModules.push_back({std::move(K), TheModule,
-          Request.OnlyTypes, Request.OnlyPrecedenceGroups, Request.OnlyMacros});
+        RequestedModules.push_back({std::move(K), TheModule, Request.Filter});
 
         auto TheModuleName = TheModule->getName();
-        if (Request.IncludeModuleQualifier &&
+        if (Request.Filter.contains(CodeCompletionFilterFlag::Module) &&
             (!Lookup.isHiddenModuleName(TheModuleName) ||
              explictlyImportedModules.contains(TheModule)) &&
             seenModuleNames.insert(TheModuleName).second)
@@ -1371,11 +1354,11 @@ void swift::ide::deliverCompletionResults(
       }
     } else {
       // Add results from current module.
-      Lookup.getToplevelCompletions(Request.OnlyTypes, Request.OnlyMacros);
+      Lookup.getToplevelCompletions(Request.Filter);
 
       // Add the qualifying module name
       auto curModule = SF.getParentModule();
-      if (Request.IncludeModuleQualifier &&
+      if (Request.Filter.contains(CodeCompletionFilterFlag::Module) &&
           seenModuleNames.insert(curModule->getName()).second)
         Lookup.addModuleName(curModule);
 
@@ -1790,15 +1773,35 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
       default:
         break;
       }
-    }
-    if (!ExpectedCustomAttributeKinds) {
+
+      switch (*AttTargetDK) {
+      case DeclKind::Var:
+      case DeclKind::Subscript:
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::VarMacro;
+        break;
+      case DeclKind::Struct:
+      case DeclKind::Class:
+      case DeclKind::Protocol:
+      case DeclKind::Enum:
+      case DeclKind::Extension:
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::ContextMacro;
+        break;
+      default:
+        break;
+      }
+      if (*AttTargetDK != DeclKind::Param) {
+        ExpectedCustomAttributeKinds |= CustomAttributeKind::DeclMacro;
+      }
+    } else {
       // 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;
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::VarMacro;
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::ContextMacro;
+      ExpectedCustomAttributeKinds |= CustomAttributeKind::DeclMacro;
     }
-    ExpectedCustomAttributeKinds |= CustomAttributeKind::Macro;
 
     Lookup.setExpectedTypes(/*Types=*/{},
                             /*isImplicitSingleExpressionReturn=*/false,
@@ -1814,8 +1817,11 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
           P.Context.SourceMgr.getIDEInspectionTargetLoc());
 
     // Macro name at attribute position after '@'.
-    Lookup.getToplevelCompletions(
-        /*OnlyTypes=*/false, /*OnlyMacros=*/true);
+    CodeCompletionMacroRoles macroRoles =
+        getCompletionMacroRoles(ExpectedCustomAttributeKinds);
+    if (macroRoles) {
+      Lookup.getMacroCompletions(macroRoles);
+    }
     break;
   }
   case CompletionKind::AttributeDeclParen: {
@@ -1968,6 +1974,15 @@ void CodeCompletionCallbacksImpl::doneParsing(SourceFile *SrcFile) {
 
   case CompletionKind::AfterPoundDirective: {
     addPoundDirectives(CompletionContext.getResultSink());
+
+    CodeCompletionMacroRoles roles;
+    if (!CurDeclContext || !CurDeclContext->isTypeContext()) {
+      roles |= CodeCompletionMacroRole::Expression;
+      roles |= CodeCompletionMacroRole::CodeItem;
+    }
+    roles |= CodeCompletionMacroRole::Declaration;
+    Lookup.getMacroCompletions(roles);
+
     // FIXME: Add pound expressions (e.g. '#selector()') if it's at statements
     // position.
     break;