[lldb] returning command completions up to a maximum (llvm#135565)

- Adding `max_return_elements` field to `CompletionRequest`.
- adding maximum checks to `SymbolCompleter` and `SourceFileCompleter`.

Fixes llvm#135553

Author: eronnen

lldb/include/lldb/Utility/CompletionRequest.h

@@ -115,6 +115,11 @@ class CompletionRequest {
   CompletionRequest(llvm::StringRef command_line, unsigned raw_cursor_pos,
                     CompletionResult &result);
 
+  /// Sets the maximum number of completions that should be returned.
+  void SetMaxReturnElements(size_t max_return_elements) {
+    m_max_return_elements = max_return_elements;
+  }
+
   /// Returns the raw user input used to create this CompletionRequest cut off
   /// at the cursor position. The cursor will be at the end of the raw line.
   llvm::StringRef GetRawLine() const {
@@ -157,6 +162,23 @@ class CompletionRequest {
 
   size_t GetCursorIndex() const { return m_cursor_index; }
 
+  size_t GetMaxReturnElements() const { return m_max_return_elements; }
+
+  /// Returns true if the maximum number of completions has not been reached
+  /// yet, hence we should keep adding completions.
+  bool ShouldAddCompletions() const {
+    return GetMaxNumberOfCompletionsToAdd() > 0;
+  }
+
+  /// Returns the maximum number of completions that need to be added
+  /// until reaching the maximum
+  size_t GetMaxNumberOfCompletionsToAdd() const {
+    const size_t number_of_results = m_result.GetNumberOfResults();
+    if (number_of_results >= m_max_return_elements)
+      return 0;
+    return m_max_return_elements - number_of_results;
+  }
+
   /// Adds a possible completion string. If the completion was already
   /// suggested before, it will not be added to the list of results. A copy of
   /// the suggested completion is stored, so the given string can be free'd
@@ -231,6 +253,8 @@ class CompletionRequest {
   size_t m_cursor_index;
   /// The cursor position in the argument indexed by m_cursor_index.
   size_t m_cursor_char_position;
+  /// The maximum number of completions that should be returned.
+  size_t m_max_return_elements = std::numeric_limits<size_t>::max();
 
   /// The result this request is supposed to fill out.
   /// We keep this object private to ensure that no backend can in any way

lldb/packages/Python/lldbsuite/test/lldbtest.py

@@ -2257,12 +2257,12 @@ def complete_from_to(self, str_input, patterns):
                 substrs=[p],
             )
 
-    def completions_match(self, command, completions):
+    def completions_match(self, command, completions, max_completions=-1):
         """Checks that the completions for the given command are equal to the
         given list of completions"""
         interp = self.dbg.GetCommandInterpreter()
         match_strings = lldb.SBStringList()
-        interp.HandleCompletion(command, len(command), 0, -1, match_strings)
+        interp.HandleCompletion(command, len(command), 0, max_completions, match_strings)
         # match_strings is a 1-indexed list, so we have to slice...
         self.assertCountEqual(
             completions, list(match_strings)[1:], "List of returned completion is wrong"

lldb/source/API/SBCommandInterpreter.cpp

@@ -263,13 +263,26 @@ int SBCommandInterpreter::HandleCompletionWithDescriptions(
   if (!IsValid())
     return 0;
 
+  if (max_return_elements == 0)
+    return 0;
+
   lldb_private::StringList lldb_matches, lldb_descriptions;
   CompletionResult result;
   CompletionRequest request(current_line, cursor - current_line, result);
+  if (max_return_elements > 0)
+    request.SetMaxReturnElements(max_return_elements);
   m_opaque_ptr->HandleCompletion(request);
   result.GetMatches(lldb_matches);
   result.GetDescriptions(lldb_descriptions);
 
+  // limit the matches to the max_return_elements if necessary
+  if (max_return_elements > 0 &&
+      lldb_matches.GetSize() > static_cast<size_t>(max_return_elements)) {
+    lldb_matches.SetSize(max_return_elements);
+    lldb_descriptions.SetSize(max_return_elements);
+  }
+  int number_of_matches = lldb_matches.GetSize();
+
   // Make the result array indexed from 1 again by adding the 'common prefix'
   // of all completions as element 0. This is done to emulate the old API.
   if (request.GetParsedLine().GetArgumentCount() == 0) {
@@ -303,7 +316,7 @@ int SBCommandInterpreter::HandleCompletionWithDescriptions(
   matches.AppendList(temp_matches_list);
   SBStringList temp_descriptions_list(&lldb_descriptions);
   descriptions.AppendList(temp_descriptions_list);
-  return result.GetNumberOfResults();
+  return number_of_matches;
 }
 
 int SBCommandInterpreter::HandleCompletionWithDescriptions(

lldb/source/Commands/CommandCompletions.cpp

@@ -91,7 +91,7 @@ bool CommandCompletions::InvokeCommonCompletionCallbacks(
        nullptr} // This one has to be last in the list.
   };
 
-  for (int i = 0;; i++) {
+  for (int i = 0; request.ShouldAddCompletions(); i++) {
     if (common_completions[i].type == lldb::eTerminatorCompletion)
       break;
     else if ((common_completions[i].type & completion_mask) ==
@@ -167,7 +167,9 @@ class SourceFileCompleter : public Completer {
         m_matching_files.AppendIfUnique(context.comp_unit->GetPrimaryFile());
       }
     }
-    return Searcher::eCallbackReturnContinue;
+    return m_matching_files.GetSize() >= m_request.GetMaxNumberOfCompletionsToAdd()
+               ? Searcher::eCallbackReturnStop
+               : Searcher::eCallbackReturnContinue;
   }
 
   void DoCompletion(SearchFilter *filter) override {
@@ -230,6 +232,9 @@ class SymbolCompleter : public Completer {
 
       // Now add the functions & symbols to the list - only add if unique:
       for (const SymbolContext &sc : sc_list) {
+        if (m_match_set.size() >= m_request.GetMaxNumberOfCompletionsToAdd())
+          break;
+
         ConstString func_name = sc.GetFunctionName(Mangled::ePreferDemangled);
         // Ensure that the function name matches the regex. This is more than
         // a sanity check. It is possible that the demangled function name
@@ -239,7 +244,9 @@ class SymbolCompleter : public Completer {
           m_match_set.insert(func_name);
       }
     }
-    return Searcher::eCallbackReturnContinue;
+    return m_match_set.size() >= m_request.GetMaxNumberOfCompletionsToAdd()
+               ? Searcher::eCallbackReturnStop
+               : Searcher::eCallbackReturnContinue;
   }
 
   void DoCompletion(SearchFilter *filter) override {
@@ -305,7 +312,8 @@ class ModuleCompleter : public Completer {
           m_request.AddCompletion(cur_file_name);
       }
     }
-    return Searcher::eCallbackReturnContinue;
+    return m_request.ShouldAddCompletions() ? Searcher::eCallbackReturnContinue
+                                            : Searcher::eCallbackReturnStop;
   }
 
   void DoCompletion(SearchFilter *filter) override { filter->Search(*this); }
@@ -429,7 +437,8 @@ static void DiskFilesOrDirectories(const llvm::Twine &partial_name,
   std::error_code EC;
   llvm::vfs::directory_iterator Iter = fs.DirBegin(SearchDir, EC);
   llvm::vfs::directory_iterator End;
-  for (; Iter != End && !EC; Iter.increment(EC)) {
+  for (; Iter != End && !EC && request.ShouldAddCompletions();
+       Iter.increment(EC)) {
     auto &Entry = *Iter;
     llvm::ErrorOr<llvm::vfs::Status> Status = fs.GetStatus(Entry.path());
 

lldb/test/API/commands/expression/completion/TestExprCompletion.py

@@ -297,6 +297,37 @@ def test_expr_completion_with_descriptions(self):
             enforce_order=True,
         )
 
+    def test_expr_completion_max_results(self):
+        self.build()
+        self.main_source = "main.cpp"
+        self.main_source_spec = lldb.SBFileSpec(self.main_source)
+        self.createTestTarget()
+
+        (target, process, thread, bkpt) = lldbutil.run_to_source_breakpoint(
+            self, "// Break here", self.main_source_spec
+        )
+
+        expected_completions = [
+            "some_expr.~Expr()",
+            "some_expr.operator=(",  # Copy operator
+            "some_expr.operator=(",  # Move operator
+            "some_expr.MemberVariableBar",
+            "some_expr.StaticMemberMethodBar()",
+            "some_expr.Self()",
+            "some_expr.FooNoArgsBar()",
+            "some_expr.FooWithArgsBar(",
+            "some_expr.FooNumbersBar1()",
+            "some_expr.FooUnderscoreBar_()",
+            "some_expr.FooWithMultipleArgsBar(",
+        ]
+
+        for i in range(1, len(expected_completions)):
+            self.completions_match(
+                "expr some_expr.",
+                expected_completions[:i],
+                max_completions=i,
+            )
+
     def assume_no_completions(self, str_input, cursor_pos=None):
         interp = self.dbg.GetCommandInterpreter()
         match_strings = lldb.SBStringList()