[clangd] Use Decision Forest to score code completions.

[usaxena95]

Summary

By default clangd will score a code completion item using heuristics model.

Scoring can be done by Decision Forest model by passing --ranking_model=decision_forest to
clangd.

Features omitted from the model:

• NameMatch is excluded because the final score must be multiplicative in NameMatch to allow rescoring by the editor.
• NeedsFixIts is excluded because the generating dataset that needs 'fixits' is non-trivial.

There are multiple ways (heuristics) to combine the above two features with the prediction of the DF:

• NeedsFixIts is used as is with a penalty of 0.5.

Various alternatives of combining NameMatch N and Decision forest Prediction P

• N * scale(P, 0, 1): Linearly scale the output of model to range [0, 1]
• N * a^P:
  • More natural: Prediction of each Decision Tree can be considered as a multiplicative boost (like NameMatch)
  • Ordering is independent of the absolute value of P. Order of two items is proportional to a^{difference in model prediction score}. Higher a gives higher weightage to model output as compared to NameMatch score.

Baseline MRR = 0.619
MRR for various combinations:
N * P = 0.6346, advantage%=2.5768
N * 1.1^P = 0.6600, advantage%=6.6853
N * 1.2^P = 0.6669, advantage%=7.8005
N * 1.3^P = 0.6668, advantage%=7.7795
N * 1.4^P = 0.6659, advantage%=7.6270
N * 1.5^P = 0.6646, advantage%=7.4200
N * 1.6^P = 0.6636, advantage%=7.2671
N * 1.7^P = 0.6629, advantage%=7.1450
N * 2^P = 0.6612, advantage%=6.8673
N * 2.5^P = 0.6598, advantage%=6.6491
N * 3^P = 0.6590, advantage%=6.5242
N * scaled[0, 1] = 0.6465, advantage%=4.5054

[usaxena95]

Moved DF evaluation to Quality.cpp

[adamcz]

Could you add a test that sets this flag? Perhaps we can run CodeCompletionTests.cpp twice, once with this flag, once without? Just to exercise these code paths, I think most expectations there are unordered, so it should work?

clang-tools-extra/clangd/CodeComplete.cpp

1635

Ideally we'd rename the evaluate() here, since SymbolQualitySignals is used for both heuristic and DecisionForest version, but evaluate is heuristic-specific. I think in pefect world this would be out of SymbolQualitySignals class (which would become just storage), but at least it should be renamed to evaluateUsingHeuristic().

1655

Could we make the weight of Relevance.NameMatch configurable, maybe through CodeCompletionOptions or such? I'm worried it may dominate the score too much and being able to configure this would allow us to run experiments easily.

[usaxena95]

Addressed comments.

[usaxena95]

If we just consider ranking, there are not many code paths. Most of the code paths are for gathering and creating code completion items.
In the end all the code paths end up in the addCandidate method to get scored either by heuristics or by decision forest.
So I have added 2 tests for checking sanity and exercising the code path.

clang-tools-extra/clangd/CodeComplete.cpp

1635

Right. I plan to do this but the changes for a rename would include unrelated files. So I planed to do this another patch keeping this one tractable and clean.
Let me know if you want me to do it in this one itself.

1655

Since this is multiplicative, I don't think it is possible to weigh NameMatch.
Although we can weigh the contribution of DF prediction.
Did some maths and realised that the base in pow(base, Evaluate(E)) can effectively represent the weightage given to DF Model.

[adamcz]

clang-tools-extra/clangd/CodeComplete.cpp

1635

Separate change sounds good.

clang-tools-extra/clangd/tool/ClangdMain.cpp

185

typo: Decsion

[usaxena95]

Added fixme for renaming evaluate() and fixed typo.

[adamcz]

Just noticed: typo in commit description - compeltions

clang-tools-extra/clangd/CodeComplete.h

struct CodeCompleteOptions {
   /// before they are ranked (by -Score). Thus the results are yielded in
   /// arbitrary order.
   ///
   /// This callbacks allows capturing various internal structures used by clangd
   /// during code completion. Eg: Symbol quality and relevance signals.
   std::function<void(const CodeCompletion &, const SymbolQualitySignals &,
                      const SymbolRelevanceSignals &, float Score)>
       RecordCCResult;
+
+  /// Model to use for ranking code completion candidates.
+  enum CodeCompletionRankingModel {
+    Heuristics,
+    DecisionForest,
+  } RankingModel = Heuristics;
+
+  /// Weight for combining NameMatch and Prediction of DecisionForest.
+  /// CompletionScore is NameMatch * pow(Base, Prediction).
+  /// The optimal value of Base largely depends on the semantics of the model
+  /// and prediction score (e.g. algorithm used during training, number of
+  /// trees, etc.). Usually if the range of Prediciton is [-20, 20] then a Base
+  /// in [1.2, 1.7] works fine.
+  /// Semantics: E.g. the completion score reduces by 50% if the Prediciton
+  /// score is reduced by 2.6 points for Base = 1.3.
+  float DecisionForestBase = 1.3f;
 };
 
 // Semi-structured representation of a code-complete suggestion for our C++ API.
 // We don't use the LSP structures here (unlike most features) as we want
 // to expose more data to allow for more precise testing and evaluation.
 struct CodeCompletion {
   // The unqualified name of the symbol or other completion item.
   std::string Name;

clang-tools-extra/clangd/CodeComplete.cpp

if (((C.SemaResult &&
           C.SemaResult->Kind == CodeCompletionResult::RK_Macro) ||
          (C.IndexResult &&
           C.IndexResult->SymInfo.Kind == index::SymbolKind::Macro)) &&
         !C.Name.startswith_lower(Filter->pattern()))
       return None;
     return Filter->match(C.Name);
   }
 
+  CodeCompletion::Scores
+  evaluateCompletion(const SymbolQualitySignals &Quality,
+                     const SymbolRelevanceSignals &Relevance) {
+    using RM = CodeCompleteOptions::CodeCompletionRankingModel;
+    CodeCompletion::Scores Scores;
+    switch (Opts.RankingModel) {
+    case RM::Heuristics:
+      Scores.Quality = Quality.evaluate();
+      Scores.Relevance = Relevance.evaluate();
+      Scores.Total =
+          evaluateSymbolAndRelevance(Scores.Quality, Scores.Relevance);
+      // NameMatch is in fact a multiplier on total score, so rescoring is
+      // sound.
+      Scores.ExcludingName = Relevance.NameMatch
+                                 ? Scores.Total / Relevance.NameMatch
+                                 : Scores.Quality;
+      return Scores;
+
+    case RM::DecisionForest:
+      Scores.Quality = 0;
+      Scores.Relevance = 0;
+      // Exponentiating DecisionForest prediction makes the score of each tree a
+      // multiplciative boost (like NameMatch). This allows us to weigh the
+      // prediciton score and NameMatch appropriately.
+      Scores.ExcludingName = pow(Opts.DecisionForestBase,
+                                 evaluateDecisionForest(Quality, Relevance));
+      // NeedsFixIts is not part of the DecisionForest as generating training
+      // data that needs fixits is not-feasible.
+      if (Relevance.NeedsFixIts)
+        Scores.ExcludingName *= 0.5;
+      // NameMatch should be a multiplier on total score to support rescoring.
+      Scores.Total = Relevance.NameMatch * Scores.ExcludingName;
+      return Scores;
+    }
+    llvm_unreachable("Unhandled CodeCompletion ranking model.");
+  }
+
   // Scores a candidate and adds it to the TopN structure.
   void addCandidate(TopN<ScoredBundle, ScoredBundleGreater> &Candidates,
                     CompletionCandidate::Bundle Bundle) {
     SymbolQualitySignals Quality;
     SymbolRelevanceSignals Relevance;
     Relevance.Context = CCContextKind;
     Relevance.Name = Bundle.front().Name;
+    Relevance.FilterLength = HeuristicPrefix.Name.size();
     Relevance.Query = SymbolRelevanceSignals::CodeComplete;
     Relevance.FileProximityMatch = FileProximity.getPointer();
     if (ScopeProximity)
       Relevance.ScopeProximityMatch = ScopeProximity.getPointer();
     if (PreferredType)
       Relevance.HadContextType = true;
     Relevance.ContextWords = &ContextWords;
 
for (const auto &Candidate : Bundle) {
       }
       if (Candidate.IdentifierResult) {
         Quality.References = Candidate.IdentifierResult->References;
         Relevance.Scope = SymbolRelevanceSignals::FileScope;
         Origin |= SymbolOrigin::Identifier;
       }
     }
 
-    CodeCompletion::Scores Scores;
-    Scores.Quality = Quality.evaluate();
-    Scores.Relevance = Relevance.evaluate();
-    Scores.Total = evaluateSymbolAndRelevance(Scores.Quality, Scores.Relevance);
-    // NameMatch is in fact a multiplier on total score, so rescoring is sound.
-    Scores.ExcludingName = Relevance.NameMatch
-                               ? Scores.Total / Relevance.NameMatch
-                               : Scores.Quality;
-
+    CodeCompletion::Scores Scores = evaluateCompletion(Quality, Relevance);
     if (Opts.RecordCCResult)
       Opts.RecordCCResult(toCodeCompletion(Bundle), Quality, Relevance,
                           Scores.Total);
 
     dlog("CodeComplete: {0} ({1}) = {2}\n{3}{4}\n", First.Name,
          llvm::to_string(Origin), Scores.Total, llvm::to_string(Quality),
          llvm::to_string(Relevance));
 

clang-tools-extra/clangd/Quality.h

enum SymbolCategory {
     Namespace,
     Keyword,
     Operator,
   } Category = Unknown;
 
   void merge(const CodeCompletionResult &SemaCCResult);
   void merge(const Symbol &IndexResult);
 
+  // FIXME(usx): Rename to evaluateHeuristics().
   // Condense these signals down to a single number, higher is better.
   float evaluate() const;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                               const SymbolQualitySignals &);
 
 /// Attributes of a symbol-query pair that affect how much we like it.
 struct SymbolRelevanceSignals {
struct SymbolRelevanceSignals {
 
   // Whether clang provided a preferred type in the completion context.
   bool HadContextType = false;
   // Whether a source completion item or a symbol had a type information.
   bool HadSymbolType = false;
   // Whether the item matches the type expected in the completion context.
   bool TypeMatchesPreferred = false;
 
+  /// Length of the unqualified partial name of Symbol typed in
+  /// CompletionPrefix.
+  unsigned FilterLength = 0;
+
   /// Set of derived signals computed by calculateDerivedSignals(). Must not be
   /// set explicitly.
   struct DerivedSignals {
     /// Whether Name contains some word from context.
     bool NameMatchesContext = false;
     /// Min distance between SymbolURI and all the headers included by the TU.
     unsigned FileProximityDistance = FileDistance::Unreachable;
     /// Min distance between SymbolScope and all the available scopes.
struct SymbolRelevanceSignals {
   float evaluate() const;
 };
 llvm::raw_ostream &operator<<(llvm::raw_ostream &,
                               const SymbolRelevanceSignals &);
 
 /// Combine symbol quality and relevance into a single score.
 float evaluateSymbolAndRelevance(float SymbolQuality, float SymbolRelevance);
 
+float evaluateDecisionForest(const SymbolQualitySignals &Quality,
+                             const SymbolRelevanceSignals &Relevance);
 /// TopN<T> is a lossy container that preserves only the "best" N elements.
 template <typename T, typename Compare = std::greater<T>> class TopN {
 public:
   using value_type = T;
   TopN(size_t N, Compare Greater = Compare())
       : N(N), Greater(std::move(Greater)) {}
 
   // Adds a candidate to the set.

clang-tools-extra/clangd/Quality.cpp

 //===--- Quality.cpp ---------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Quality.h"
 #include "AST.h"
+#include "CompletionModel.h"
 #include "FileDistance.h"
 #include "SourceCode.h"
 #include "URI.h"
 #include "index/Symbol.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclTemplate.h"
llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
 
   return OS;
 }
 
 float evaluateSymbolAndRelevance(float SymbolQuality, float SymbolRelevance) {
   return SymbolQuality * SymbolRelevance;
 }
 
+float evaluateDecisionForest(const SymbolQualitySignals &Quality,
+                             const SymbolRelevanceSignals &Relevance) {
+  Example E;
+  E.setIsDeprecated(Quality.Deprecated);
+  E.setIsReservedName(Quality.ReservedName);
+  E.setIsImplementationDetail(Quality.ImplementationDetail);
+  E.setNumReferences(Quality.References);
+  E.setSymbolCategory(Quality.Category);
+
+  SymbolRelevanceSignals::DerivedSignals Derived =
+      Relevance.calculateDerivedSignals();
+  E.setIsNameInContext(Derived.NameMatchesContext);
+  E.setIsForbidden(Relevance.Forbidden);
+  E.setIsInBaseClass(Relevance.InBaseClass);
+  E.setFileProximityDistance(Derived.FileProximityDistance);
+  E.setSemaFileProximityScore(Relevance.SemaFileProximityScore);
+  E.setSymbolScopeDistance(Derived.ScopeProximityDistance);
+  E.setSemaSaysInScope(Relevance.SemaSaysInScope);
+  E.setScope(Relevance.Scope);
+  E.setContextKind(Relevance.Context);
+  E.setIsInstanceMember(Relevance.IsInstanceMember);
+  E.setHadContextType(Relevance.HadContextType);
+  E.setHadSymbolType(Relevance.HadSymbolType);
+  E.setTypeMatchesPreferred(Relevance.TypeMatchesPreferred);
+  E.setFilterLength(Relevance.FilterLength);
+  return Evaluate(E);
+}
+
 // Produces an integer that sorts in the same order as F.
 // That is: a < b <==> encodeFloat(a) < encodeFloat(b).
 static uint32_t encodeFloat(float F) {
   static_assert(std::numeric_limits<float>::is_iec559, "");
   constexpr uint32_t TopBit = ~(~uint32_t{0} >> 1);
 
   // Get the bits of the float. Endianness is the same as for integers.
   uint32_t U = llvm::FloatToBits(F);

clang-tools-extra/clangd/tool/ClangdMain.cpp

values(clEnumValN(CodeCompleteOptions::AlwaysParse, "always",
                       "Use text-based completion if the parser "
                       "is not ready"),
            clEnumValN(CodeCompleteOptions::NeverParse, "never",
                       "Always used text-based completion")),
     init(CodeCompleteOptions().RunParser),
     Hidden,
 };
 
+opt<CodeCompleteOptions::CodeCompletionRankingModel> RankingModel{
+    "ranking-model",
+    cat(Features),
+    desc("Model to use to rank code-completion items"),
+    values(clEnumValN(CodeCompleteOptions::Heuristics, "heuristics",
+                      "Use hueristics to rank code completion items"),
+           clEnumValN(CodeCompleteOptions::DecisionForest, "decision_forest",
+                      "Use Decision Forest model to rank completion items")),
+    init(CodeCompleteOptions().RankingModel),
+    Hidden,
+};
+
+opt<bool> DecisionForestBase{
+    "decision-forest-base",
+    cat(Features),
+    desc("Base for exponentiating the prediction from DecisionForest."),
+    init(CodeCompleteOptions().DecisionForestBase),
+    Hidden,
+};
+
 // FIXME: also support "plain" style where signatures are always omitted.
 enum CompletionStyleFlag { Detailed, Bundled };
 opt<CompletionStyleFlag> CompletionStyle{
     "completion-style",
     cat(Features),
     desc("Granularity of code completion suggestions"),
     values(clEnumValN(Detailed, "detailed",
                       "One completion item for each semantically distinct "
#endif
   if (!HeaderInsertionDecorators) {
     CCOpts.IncludeIndicator.Insert.clear();
     CCOpts.IncludeIndicator.NoInsert.clear();
   }
   CCOpts.SpeculativeIndexRequest = Opts.StaticIndex;
   CCOpts.EnableFunctionArgSnippets = EnableFunctionArgSnippets;
   CCOpts.AllScopes = AllScopesCompletion;
   CCOpts.RunParser = CodeCompletionParse;
+  CCOpts.RankingModel = RankingModel;
+  CCOpts.DecisionForestBase = DecisionForestBase;
 
   RealThreadsafeFS TFS;
   std::vector<std::unique_ptr<config::Provider>> ProviderStack;
   std::unique_ptr<config::Provider> Config;
   if (EnableConfig) {
     ProviderStack.push_back(
         config::Provider::fromAncestorRelativeYAMLFiles(".clangd", TFS));
     llvm::SmallString<256> UserConfig;

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

 //===-- CodeCompleteTests.cpp -----------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
 #include "ClangdServer.h"
 #include "CodeComplete.h"
 #include "Compiler.h"
-#include "CompletionModel.h"
 #include "Matchers.h"
 #include "Protocol.h"
 #include "Quality.h"
 #include "SourceCode.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
 #include "TestIndex.h"
 #include "TestTU.h"
return codeComplete(FilePath, Test.point(), /*Preamble=*/nullptr, ParseInput,
                       Opts);
 }
 
 Symbol withReferences(int N, Symbol S) {
   S.References = N;
   return S;
 }
 
-TEST(DecisionForestRuntime, SanityTest) {
-  using Example = clangd::Example;
-  using clangd::Evaluate;
-  Example E1;
-  E1.setContextKind(ContextKind::CCC_ArrowMemberAccess);
-  Example E2;
-  E2.setContextKind(ContextKind::CCC_SymbolOrNewName);
-  EXPECT_GT(Evaluate(E1), Evaluate(E2));
+TEST(DecisionForestRankingModel, NameMatchSanityTest) {
+  clangd::CodeCompleteOptions Opts;
+  Opts.RankingModel = CodeCompleteOptions::DecisionForest;
+  auto Results = completions(
+      R"cpp(
+struct MemberAccess {
+  int ABG();
+  int AlphaBetaGamma();
+};
+int func() { MemberAccess().ABG^ }
+)cpp",
+      /*IndexSymbols=*/{}, Opts);
+  EXPECT_THAT(Results.Completions,
+              ElementsAre(Named("ABG"), Named("AlphaBetaGamma")));
+}
+
+TEST(DecisionForestRankingModel, ReferencesAffectRanking) {
+  clangd::CodeCompleteOptions Opts;
+  Opts.RankingModel = CodeCompleteOptions::DecisionForest;
+  constexpr int NumReferences = 100000;
+  EXPECT_THAT(
+      completions("int main() { clang^ }",
+                  {ns("clangA"), withReferences(NumReferences, func("clangD"))},
+                  Opts)
+          .Completions,
+      ElementsAre(Named("clangD"), Named("clangA")));
+  EXPECT_THAT(
+      completions("int main() { clang^ }",
+                  {withReferences(NumReferences, ns("clangA")), func("clangD")},
+                  Opts)
+          .Completions,
+      ElementsAre(Named("clangA"), Named("clangD")));
 }
 
 TEST(CompletionTest, Limit) {
   clangd::CodeCompleteOptions Opts;
   Opts.Limit = 2;
   auto Results = completions(R"cpp(
 struct ClassWithMembers {
   int AAA();