[clangd][ObjC] Improve completions for protocols + category names

[dgoldman]

Summary

• Render protocols as interfaces to differentiate them from classes since a protocol and class can have the same name. Take this one step further though, and only recommend protocols in ObjC protocol completions.

• Properly call includeSymbolFromIndex even with a cached speculative fuzzy find request

• Don't use the index to provide completions for category names, symbols there don't make sense

[kadircet]

can you also add test cases for the other two (filtering both for speculative index queries/regular ones, and making sure we don't suggest symbols from index for category names), so that we don't regress in the future?

[dgoldman]

Add tests

[dgoldman]

In D132962#3761463, @kadircet wrote:

can you also add test cases for the other two (filtering both for speculative index queries/regular ones, and making sure we don't suggest symbols from index for category names), so that we don't regress in the future?

Done

[dgoldman]

Run clang format

[kadircet]

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

3289

nit: just completions("Fo^", ...)

3335

we actually need to issue the new request in a context that can show FoodClass, as ATM if we were to do filtering during speculative query response, this test would still pass.

[dgoldman]

Fix strings + test

[dgoldman]

Run clang-format again

[kadircet]

thanks!

clang-tools-extra/clangd/CodeComplete.cpp

1730–1731

nit: while here can you change this to read as:
Incomplete ||= Opts.Index->fuzzyFind(...);

[dgoldman]

Address comment

[dgoldman]

clang-tools-extra/clangd/CodeComplete.cpp

1730–1731

Done, but there's no ||= and we wouldn't want to short circuit, so used |= instead.

clang-tools-extra/clangd/CodeComplete.cpp

case SK::NamespaceAlias:
     return CompletionItemKind::Module;
   case SK::Macro:
     return CompletionItemKind::Text;
   case SK::Enum:
     return CompletionItemKind::Enum;
   case SK::Struct:
     return CompletionItemKind::Struct;
   case SK::Class:
-  case SK::Protocol:
   case SK::Extension:
   case SK::Union:
     return CompletionItemKind::Class;
+  case SK::Protocol:
+    // Use interface instead of class for differentiation of classes and
+    // protocols with the same name (e.g. @interface NSObject vs. @protocol
+    // NSObject).
+    return CompletionItemKind::Interface;
   case SK::TypeAlias:
     // We use the same kind as the VSCode C++ extension.
     // FIXME: pick a better option when we have one.
     return CompletionItemKind::Interface;
   case SK::Using:
     return CompletionItemKind::Reference;
   case SK::Function:
   case SK::ConversionFunction:
bool contextAllowsIndex(enum CodeCompletionContext::Kind K) {
   case CodeCompletionContext::CCC_EnumTag:
   case CodeCompletionContext::CCC_UnionTag:
   case CodeCompletionContext::CCC_ClassOrStructTag:
   case CodeCompletionContext::CCC_ObjCProtocolName:
   case CodeCompletionContext::CCC_Namespace:
   case CodeCompletionContext::CCC_Type:
   case CodeCompletionContext::CCC_ParenthesizedExpression:
   case CodeCompletionContext::CCC_ObjCInterfaceName:
-  case CodeCompletionContext::CCC_ObjCCategoryName:
   case CodeCompletionContext::CCC_Symbol:
   case CodeCompletionContext::CCC_SymbolOrNewName:
     return true;
   case CodeCompletionContext::CCC_OtherWithMacros:
   case CodeCompletionContext::CCC_DotMemberAccess:
   case CodeCompletionContext::CCC_ArrowMemberAccess:
+  case CodeCompletionContext::CCC_ObjCCategoryName:
   case CodeCompletionContext::CCC_ObjCPropertyAccess:
   case CodeCompletionContext::CCC_MacroName:
   case CodeCompletionContext::CCC_MacroNameUse:
   case CodeCompletionContext::CCC_PreprocessorExpression:
   case CodeCompletionContext::CCC_PreprocessorDirective:
   case CodeCompletionContext::CCC_SelectorName:
   case CodeCompletionContext::CCC_TypeQualifiers:
   case CodeCompletionContext::CCC_ObjCInstanceMessage:
bool allowIndex(CodeCompletionContext &CC) {
   case NestedNameSpecifier::TypeSpecWithTemplate:
   // Unresolved inside a template.
   case NestedNameSpecifier::Identifier:
     return false;
   }
   llvm_unreachable("invalid NestedNameSpecifier kind");
 }
 
+// Should we include a symbol from the index given the completion kind?
+// FIXME: Ideally we can filter in the fuzzy find request itself.
+bool includeSymbolFromIndex(CodeCompletionContext::Kind Kind,
+                            const Symbol &Sym) {
+  // Objective-C protocols are only useful in ObjC protocol completions,
+  // in other places they're confusing, especially when they share the same
+  // identifier with a class.
+  if (Sym.SymInfo.Kind == index::SymbolKind::Protocol &&
+      Sym.SymInfo.Lang == index::SymbolLanguage::ObjC)
+    return Kind == CodeCompletionContext::CCC_ObjCProtocolName;
+  else if (Kind == CodeCompletionContext::CCC_ObjCProtocolName)
+    // Don't show anything else in ObjC protocol completions.
+    return false;
+  return true;
+}
+
 std::future<SymbolSlab> startAsyncFuzzyFind(const SymbolIndex &Index,
                                             const FuzzyFindRequest &Req) {
   return runAsync<SymbolSlab>([&Index, Req]() {
     trace::Span Tracer("Async fuzzyFind");
     SymbolSlab::Builder Syms;
     Index.fuzzyFind(Req, [&Syms](const Symbol &Sym) { Syms.insert(Sym); });
     return std::move(Syms).build();
   });
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) {
-          if (includeSymbolFromIndex(Sym))
-            ResultsBuilder.insert(Sym);
-        }))
-      Incomplete = true;
+    Incomplete |= Opts.Index->fuzzyFind(
+        Req, [&](const Symbol &Sym) { ResultsBuilder.insert(Sym); });
     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
   // bundles are scored and top results are returned, best to worst.
mergeResults(const std::vector<CodeCompletionResult> &SemaResults,
     };
     // Emit all Sema results, merging them with Index results if possible.
     for (auto &SemaResult : SemaResults)
       AddToBundles(&SemaResult, CorrespondingIndexResult(SemaResult), nullptr);
     // Now emit any Index-only results.
     for (const auto &IndexResult : IndexResults) {
       if (UsedIndexResults.count(&IndexResult))
         continue;
+      if (!includeSymbolFromIndex(CCContextKind, IndexResult))
+        continue;
       AddToBundles(/*SemaResult=*/nullptr, &IndexResult, nullptr);
     }
     // Emit identifier results.
     for (const auto &Ident : IdentifierResults)
       AddToBundles(/*SemaResult=*/nullptr, /*IndexResult=*/nullptr, &Ident);
     // We only keep the best N results at any time, in "native" format.
     TopN<ScoredBundle, ScoredBundleGreater> Top(
         Opts.Limit == 0 ? std::numeric_limits<size_t>::max() : Opts.Limit);

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

auto Results =
                     TU.inputs(FS), MarkupKind::PlainText);
   EXPECT_THAT(Results.signatures, ElementsAre(sig("foo([[int x]]) -> int")));
   EXPECT_EQ(0, Results.activeSignature);
   EXPECT_EQ(0, Results.activeParameter);
 }
 
 class IndexRequestCollector : public SymbolIndex {
 public:
+  IndexRequestCollector(std::vector<Symbol> Syms = {}) : Symbols(Syms) {}
+
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override {
     std::unique_lock<std::mutex> Lock(Mut);
     Requests.push_back(Req);
     ReceivedRequestCV.notify_one();
+    for (const auto &Sym : Symbols)
+      Callback(Sym);
     return true;
   }
 
   void lookup(const LookupRequest &,
               llvm::function_ref<void(const Symbol &)>) const override {}
 
   bool refs(const RefsRequest &,
             llvm::function_ref<void(const Ref &)>) const override {
const std::vector<FuzzyFindRequest> consumeRequests(size_t Num) const {
     EXPECT_TRUE(wait(Lock, ReceivedRequestCV, timeoutSeconds(30),
                      [this, Num] { return Requests.size() == Num; }));
     auto Reqs = std::move(Requests);
     Requests = {};
     return Reqs;
   }
 
 private:
+  std::vector<Symbol> Symbols;
   // We need a mutex to handle async fuzzy find requests.
   mutable std::condition_variable ReceivedRequestCV;
   mutable std::mutex Mut;
   mutable std::vector<FuzzyFindRequest> Requests;
 };
 
 // Clients have to consume exactly Num requests.
 std::vector<FuzzyFindRequest> captureIndexRequests(llvm::StringRef Code,
TEST(CompletionTest, ObjectiveCMethodDeclarationFromMiddle) {
   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("id<Foo^>", {SymFood, FoodClass, SymFooey},
+                             /*Opts=*/{}, "Foo.m");
+
+  // Should only give protocols for ObjC protocol completions.
+  EXPECT_THAT(Results.Completions,
+              UnorderedElementsAre(
+                  AllOf(named("Food"), kind(CompletionItemKind::Interface)),
+                  AllOf(named("Fooey"), kind(CompletionItemKind::Interface))));
+
+  Results = completions("Fo^", {SymFood, FoodClass, SymFooey},
+                        /*Opts=*/{}, "Foo.m");
+  // Shouldn't give protocols for non protocol completions.
+  EXPECT_THAT(
+      Results.Completions,
+      ElementsAre(AllOf(named("FoodClass"), kind(CompletionItemKind::Class))));
+}
+
+TEST(CompletionTest, ObjectiveCProtocolFromIndexSpeculation) {
+  MockFS FS;
+  MockCompilationDatabase CDB;
+  ClangdServer Server(CDB, FS, ClangdServer::optsForTest());
+
+  auto File = testPath("Foo.m");
+  Annotations Test(R"cpp(
+      @protocol Food
+      @end
+      id<Foo$1^> foo;
+      Foo$2^ bar;
+  )cpp");
+  runAddDocument(Server, File, Test.code());
+  clangd::CodeCompleteOptions Opts = {};
+
+  Symbol FoodClass = objcClass("FoodClass");
+  IndexRequestCollector Requests({FoodClass});
+  Opts.Index = &Requests;
+
+  auto CompleteAtPoint = [&](StringRef P) {
+    return cantFail(runCodeComplete(Server, File, Test.point(P), Opts))
+        .Completions;
+  };
+
+  auto C = CompleteAtPoint("1");
+  auto Reqs1 = Requests.consumeRequests(1);
+  ASSERT_EQ(Reqs1.size(), 1u);
+  EXPECT_THAT(C, ElementsAre(AllOf(named("Food"),
+                                   kind(CompletionItemKind::Interface))));
+
+  C = CompleteAtPoint("2");
+  auto Reqs2 = Requests.consumeRequests(1);
+  // Speculation succeeded. Used speculative index result, but filtering now to
+  // now include FoodClass.
+  ASSERT_EQ(Reqs2.size(), 1u);
+  EXPECT_EQ(Reqs2[0], Reqs1[0]);
+  EXPECT_THAT(C, ElementsAre(AllOf(named("FoodClass"),
+                                   kind(CompletionItemKind::Class))));
+}
+
+TEST(CompletionTest, ObjectiveCCategoryFromIndexIgnored) {
+  Symbol FoodCategory = objcCategory("FoodClass", "Extension");
   auto Results = completions(R"objc(
-      id<Foo^>
+      @interface Foo
+      @end
+      @interface Foo (^)
+      @end
     )objc",
-                             {SymFood, FoodClass, SymFooey},
+                             {FoodCategory},
                              /*Opts=*/{}, "Foo.m");
-
-  auto C = Results.Completions;
-  EXPECT_THAT(C, UnorderedElementsAre(named("Food"), named("Fooey")));
+  EXPECT_THAT(Results.Completions, IsEmpty());
 }
 
 TEST(CompletionTest, CursorInSnippets) {
   clangd::CodeCompleteOptions Options;
   Options.EnableSnippets = true;
   auto Results = completions(
       R"cpp(
     void while_foo(int a, int b);
     void test() {
       whil^
     })cpp",
       /*IndexSymbols=*/{}, Options);
 
   // Last placeholder in code patterns should be $0 to put the cursor there.
const char *Source(R"cpp(
     })cpp");
   auto Index = TestTU::withHeaderCode(Header).index();
   clangd::CodeCompleteOptions Opts;
   Opts.Index = Index.get();
   Opts.AllScopes = true;
   auto R = completions(Source, {}, Opts);
   EXPECT_THAT(R.Completions,
               ElementsAre(AllOf(scope("std::"), named("foo"),
                                 kind(CompletionItemKind::Reference))));
 }
 
 TEST(CompletionTest, ScopeIsUnresolved) {
   clangd::CodeCompleteOptions Opts = {};
   Opts.AllScopes = true;
 
   auto Results = completions(R"cpp(
     namespace a {

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

 // 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 @interface or @implementation category.
+Symbol objcCategory(llvm::StringRef Name, llvm::StringRef CategoryName);
 // 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);

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

Symbol objcSym(llvm::StringRef Name, index::SymbolKind Kind,
   Sym.Origin = SymbolOrigin::Static;
   return Sym;
 }
 
 Symbol objcClass(llvm::StringRef Name) {
   return objcSym(Name, index::SymbolKind::Class, "objc(cs)");
 }
 
+Symbol objcCategory(llvm::StringRef Name, llvm::StringRef CategoryName) {
+  std::string USRPrefix = ("objc(cy)" + Name + "@").str();
+  return objcSym(CategoryName, index::SymbolKind::Extension, USRPrefix);
+}
+
 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));