[clangd] Improve ObjC protocol suggestions from the index

[dgoldman]

Summary

When querying the index during an ObjC protocol name lookup for code
completion, we should only suggest ObjC protocols.

[kadircet]

thanks, lgtm!

clang-tools-extra/clangd/CodeComplete.cpp

for (auto &C : Scored) {
       Output.Completions.back().CompletionTokenRange = ReplacedRange;
     }
     Output.HasMore = Incomplete;
     Output.Context = CCContextKind;
     Output.CompletionRange = ReplacedRange;
     return Output;
   }
 
+  bool includeSymbolFromIndex(const Symbol &Sym) {
+    if (CCContextKind == CodeCompletionContext::CCC_ObjCProtocolName) {
+      return Sym.SymInfo.Lang == index::SymbolLanguage::ObjC &&
+          Sym.SymInfo.Kind == index::SymbolKind::Protocol;
+    }
+    return true;
+  }
+
   SymbolSlab queryIndex() {
     trace::Span Tracer("Query index");
     SPAN_ATTACH(Tracer, "limit", int64_t(Opts.Limit));
 
     // Build the query.
     FuzzyFindRequest Req;
     if (Opts.Limit)
       Req.Limit = Opts.Limit;
if (SpecFuzzyFind && SpecFuzzyFind->Result.valid() && (*SpecReq == Req)) {
       trace::Span WaitSpec("Wait speculative results");
       return SpecFuzzyFind->Result.get();
     }
 
     SPAN_ATTACH(Tracer, "Speculative results", false);
 
     // Run the query against the index.
     SymbolSlab::Builder ResultsBuilder;
-    if (Opts.Index->fuzzyFind(
-            Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); }))
+    if (Opts.Index->fuzzyFind(Req, [&](const Symbol &Sym) {
+          if (includeSymbolFromIndex(Sym))
+            ResultsBuilder.insert(Sym);
+        }))
       Incomplete = true;
     return std::move(ResultsBuilder).build();
   }
 
   // Merges Sema and Index results where possible, to form CompletionCandidates.
   // \p Identifiers is raw identifiers that can also be completion candidates.
   // Identifiers are not merged with results from index or sema.
   // Groups overloads if desired, to form CompletionCandidate::Bundles. The

clang-tools-extra/clangd/unittests/CodeCompleteTests.cpp

)objc",
                              /*Opts=*/{}, "Foo.m");
 
   auto C = Results.Completions;
   EXPECT_THAT(C, ElementsAre(named("secondArgument:")));
   EXPECT_THAT(C, ElementsAre(kind(CompletionItemKind::Method)));
   EXPECT_THAT(C, ElementsAre(signature("(id)object")));
 }
 
+TEST(CompletionTest, ObjectiveCProtocolFromIndex) {
+  Symbol FoodClass = objcClass("FoodClass");
+  Symbol SymFood = objcProtocol("Food");
+  Symbol SymFooey = objcProtocol("Fooey");
+  auto Results = completions(R"objc(
+      id<Foo^>
+    )objc",
+                             {SymFood, FoodClass, SymFooey},
+                             /*Opts=*/{}, "Foo.m");
+
+  auto C = Results.Completions;
+  EXPECT_THAT(C, UnorderedElementsAre(named("Food"), named("Fooey")));
+}
+
 TEST(CompletionTest, CursorInSnippets) {
   clangd::CodeCompleteOptions Options;
   Options.EnableSnippets = true;
   auto Results = completions(
       R"cpp(
     void while_foo(int a, int b);
     void test() {
       whil^

clang-tools-extra/clangd/unittests/TestIndex.h

 Symbol enm(llvm::StringRef Name);
 // Creates a variable symbol.
 Symbol var(llvm::StringRef Name);
 // Creates a namespace symbol.
 Symbol ns(llvm::StringRef Name);
 // Create a C++20 concept symbol.
 Symbol conceptSym(llvm::StringRef Name);
 
+// Create an Objective-C symbol.
+Symbol objcSym(llvm::StringRef Name, index::SymbolKind Kind,
+               llvm::StringRef USRPrefix);
+// Create an @interface or @implementation.
+Symbol objcClass(llvm::StringRef Name);
+// Create an @protocol.
+Symbol objcProtocol(llvm::StringRef Name);
+
 // Create a slab of symbols with the given qualified names as IDs and names.
 SymbolSlab generateSymbols(std::vector<std::string> QualifiedNames);
 
 // Create a slab of symbols with IDs and names [Begin, End].
 SymbolSlab generateNumSymbols(int Begin, int End);
 
 // Returns fully-qualified name out of given symbol.
 std::string getQualifiedName(const Symbol &Sym);

clang-tools-extra/clangd/unittests/TestIndex.cpp

 Symbol ns(llvm::StringRef Name) {
   return sym(Name, index::SymbolKind::Namespace, "@N@\\0");
 }
 
 Symbol conceptSym(llvm::StringRef Name) {
   return sym(Name, index::SymbolKind::Concept, "@CT@\\0");
 }
 
+Symbol objcSym(llvm::StringRef Name, index::SymbolKind Kind,
+               llvm::StringRef USRPrefix) {
+  Symbol Sym;
+  std::string USR = USRPrefix.str() + Name.str();
+  Sym.Name = Name;
+  Sym.Scope = "";
+  Sym.ID = SymbolID(USR);
+  Sym.SymInfo.Kind = Kind;
+  Sym.SymInfo.Lang = index::SymbolLanguage::ObjC;
+  Sym.Flags |= Symbol::IndexedForCodeCompletion;
+  Sym.Origin = SymbolOrigin::Static;
+  return Sym;
+}
+
+Symbol objcClass(llvm::StringRef Name) {
+  return objcSym(Name, index::SymbolKind::Class, "objc(cs)");
+}
+
+Symbol objcProtocol(llvm::StringRef Name) {
+  return objcSym(Name, index::SymbolKind::Protocol, "objc(pl)");
+}
+
 SymbolSlab generateSymbols(std::vector<std::string> QualifiedNames) {
   SymbolSlab::Builder Slab;
   for (llvm::StringRef QName : QualifiedNames)
     Slab.insert(symbol(QName));
   return std::move(Slab).build();
 }
 
 SymbolSlab generateNumSymbols(int Begin, int End) {