Faster symbol completion with cache and limit

lib/irb/completion.rb

@@ -96,6 +96,31 @@ def command_candidates(target)
       end
     end
 
+    def clear_symbol_cache
+      @sorted_symbol_names = nil
+    end
+
+    def symbol_candidates(prefix, first: 50, last: 50)
+      limit = first + last
+      symbol_names = @sorted_symbol_names ||= Symbol.all_symbols.filter_map do
+        _1.inspect[1..]
+      rescue EncodingError
+        # ignore
+      end.sort
+      start_index = symbol_names.bsearch_index { |name| name >= prefix }
+      return [] unless start_index
+
+      end_index = (start_index...symbol_names.size).bsearch { |i| !symbol_names[i].start_with?(prefix) } || symbol_names.size
+      if end_index - start_index <= limit
+        symbol_names[start_index...end_index]
+      else
+        # To avoid wrong perfect match completion, we should include first and last candidates.
+        # e.g. prefix = 'a', symbol_names = 'aaaa'...'zzzz'
+        # if this method returns first 100 of symbol_names('aaa'..'aadv'), Reline/Readline will wrongly completes the common prefix 'aa'.
+        symbol_names[start_index, first] + symbol_names[end_index - last, last]
+      end
+    end
+
     def retrieve_files_to_require_relative_from_current_dir
       @files_from_current_dir ||= Dir.glob("**/*.{rb,#{RbConfig::CONFIG['DLEXT']}}", base: '.').map { |path|
         path.sub(/\.(rb|#{RbConfig::CONFIG['DLEXT']})\z/, '')
@@ -116,18 +141,27 @@ def completion_candidates(preposing, target, _postposing, bind:)
       # When completing the argument of `help` command, only commands should be candidates
       return command_candidates(target) if preposing.match?(HELP_COMMAND_PREPOSING)
 
-      commands = if preposing.empty?
-        command_candidates(target)
+      type_candidates = type_completion_candidates(preposing, target, bind)
+
+      if preposing.empty?
+        command_candidates(target) | type_candidates
       # It doesn't make sense to propose commands with other preposing
       else
-        []
+        type_candidates
       end
+    end
 
+    def type_completion_candidates(preposing, target, bind)
       result = ReplTypeCompletor.analyze(preposing + target, binding: bind, filename: @context.irb_path)
+      return [] unless result
 
-      return commands unless result
-
-      commands | result.completion_candidates.map { target + _1 }
+      analyze_result = result.instance_variable_get(:@analyze_result)
+      if analyze_result.is_a?(Array) && analyze_result[0] == :symbol && analyze_result[1].is_a?(String)
+        symbol_prefix = analyze_result[1]
+        symbol_candidates(symbol_prefix).map { target + _1[symbol_prefix.size..] }
+      else
+        result.completion_candidates.map { target + _1 }
+      end
     end
 
     def doc_namespace(preposing, matched, _postposing, bind:)
@@ -280,12 +314,9 @@ def retrieve_completion_data(input, bind:, doc_namespace:)
           nil
         else
           sym = $1
-          candidates = Symbol.all_symbols.collect do |s|
-            s.inspect
-          rescue EncodingError
-            # ignore
+          candidates = symbol_candidates(sym[1..]).map do |s|
+            ":#{s}"
           end
-          candidates.grep(/^#{Regexp.quote(sym)}/)
         end
       when /^::([A-Z][^:\.\(\)]*)$/
         # Absolute Constant or class methods

lib/irb/input-method.rb

@@ -216,6 +216,7 @@ def completion_info
     def gets
       Readline.input = @stdin
       Readline.output = @stdout
+      @completor.clear_symbol_cache
       if l = readline(@prompt, false)
         HISTORY.push(l) if !l.empty?
         @line[@line_no += 1] = l + "\n"
@@ -473,6 +474,7 @@ def gets
       Reline.output = @stdout
       Reline.prompt_proc = @prompt_proc
       Reline.auto_indent_proc = @auto_indent_proc if @auto_indent_proc
+      @completor.clear_symbol_cache
       if l = Reline.readmultiline(@prompt, false, &@check_termination_proc)
         Reline::HISTORY.push(l) if !l.empty?
         @line[@line_no += 1] = l + "\n"

test/irb/test_completion.rb

@@ -230,15 +230,27 @@ def test_complete_symbol
         "K".force_encoding(enc).to_sym
       rescue
       end
-      symbols += [:aiueo, :"aiu eo"]
-      candidates = completion_candidates(":a", binding)
-      assert_include(candidates, ":aiueo")
-      assert_not_include(candidates, ":aiu eo")
+      symbols += [:irb_test_symbol_aiueo, :"irb_test_symbol_aiu eo"]
+      candidates = completion_candidates(":irb_test_symbol_a", binding)
+      assert_include(candidates, ":irb_test_symbol_aiueo")
+      assert_not_include(candidates, ":irb_test_symbol_aiu eo")
       assert_empty(completion_candidates(":irb_unknown_symbol_abcdefg", binding))
       # Do not complete empty symbol for performance reason
       assert_empty(completion_candidates(":", binding))
     end
 
+    def test_complete_symbol_limit
+      symbols = 200.times.map { :"irb_test_sym_limit_#{_1}" }.sort
+      assert_equal(11, completion_candidates(":irb_test_sym_limit_9", binding).size)
+      assert_equal([], completion_candidates(":irb_test_sym_nomatch", binding))
+      assert_equal([], completion_candidates(":zzzz", binding))
+
+      limited_candidates = completion_candidates(":irb_test_sym_lim", binding)
+      assert_include(limited_candidates, symbols.first.inspect)
+      assert_include(limited_candidates, symbols.last.inspect)
+      assert_equal(100, limited_candidates.size)
+    end
+
     def test_complete_invalid_three_colons
       assert_empty(completion_candidates(":::A", binding))
       assert_empty(completion_candidates(":::", binding))

test/irb/test_type_completor.rb

@@ -61,6 +61,14 @@ def test_type_completion
       assert_doc_namespace('num.chr.', 'upcase', 'String#upcase', binding: bind)
     end
 
+    def test_complete_symbol_limit
+      symbols = 200.times.map { :"irb_test_sym_limit_#{_1}" }.sort
+      candidates = @completor.completion_candidates('', ':irb_test_sym_lim', '', bind: binding)
+      assert_include(candidates, symbols.first.inspect)
+      assert_include(candidates, symbols.last.inspect)
+      assert_equal(candidates.size, 100)
+    end
+
     def test_inspect
       assert_match(/\AReplTypeCompletor.*\z/, @completor.inspect)
     end