Improve prompt generating performance by caching prompt parts(%m, %M)

lib/irb.rb

@@ -87,6 +87,7 @@ class Irb
     # Creates a new irb session
     def initialize(workspace = nil, input_method = nil, from_binding: false)
       @from_binding = from_binding
+      @prompt_part_cache = nil
       @context = Context.new(self, workspace, input_method)
       @context.workspace.load_helper_methods_to_main
       @signal_status = :IN_IRB
@@ -238,13 +239,7 @@ def read_input(prompt)
       end
     end
 
-    def readmultiline
-      prompt = generate_prompt([], false, 0)
-
-      # multiline
-      return read_input(prompt) if @context.io.respond_to?(:check_termination)
-
-      # nomultiline
+    def read_input_nomultiline(prompt)
       code = +''
       line_offset = 0
       loop do
@@ -265,6 +260,20 @@ def readmultiline
       end
     end
 
+    def readmultiline
+      with_prompt_part_cached do
+        prompt = generate_prompt([], false, 0)
+
+        if @context.io.respond_to?(:check_termination)
+          # multiline
+          read_input(prompt)
+        else
+          # nomultiline
+          read_input_nomultiline(prompt)
+        end
+      end
+    end
+
     def each_top_level_statement
       loop do
         code = readmultiline
@@ -567,6 +576,13 @@ def inspect
 
     private
 
+    def with_prompt_part_cached
+      @prompt_part_cache = {}
+      yield
+    ensure
+      @prompt_part_cache = nil
+    end
+
     def generate_prompt(opens, continue, line_offset)
       ltype = @scanner.ltype_from_open_tokens(opens)
       indent = @scanner.calc_indent_level(opens)
@@ -598,25 +614,28 @@ def generate_prompt(opens, continue, line_offset)
     end
 
     def truncate_prompt_main(str) # :nodoc:
-      str = str.tr(CONTROL_CHARACTERS_PATTERN, ' ')
-      if str.size <= PROMPT_MAIN_TRUNCATE_LENGTH
-        str
-      else
-        str[0, PROMPT_MAIN_TRUNCATE_LENGTH - PROMPT_MAIN_TRUNCATE_OMISSION.size] + PROMPT_MAIN_TRUNCATE_OMISSION
+      if str.size > PROMPT_MAIN_TRUNCATE_LENGTH
+        str = str[0, PROMPT_MAIN_TRUNCATE_LENGTH - PROMPT_MAIN_TRUNCATE_OMISSION.size] + PROMPT_MAIN_TRUNCATE_OMISSION
       end
+      str.tr(CONTROL_CHARACTERS_PATTERN, ' ')
     end
 
     def format_prompt(format, ltype, indent, line_no) # :nodoc:
+      part_cache = @prompt_part_cache || {}
       format.gsub(/%([0-9]+)?([a-zA-Z%])/) do
         case $2
         when "N"
           @context.irb_name
         when "m"
-          main_str = "#{@context.safe_method_call_on_main(:to_s)}" rescue "!#{$!.class}"
-          truncate_prompt_main(main_str)
+          part_cache[:m] ||= (
+            main_str = "#{@context.safe_method_call_on_main(:to_s)}" rescue "!#{$!.class}"
+            truncate_prompt_main(main_str)
+          )
         when "M"
-          main_str = "#{@context.safe_method_call_on_main(:inspect)}" rescue "!#{$!.class}"
-          truncate_prompt_main(main_str)
+          part_cache[:M] ||= (
+            main_str = "#{@context.safe_method_call_on_main(:inspect)}" rescue "!#{$!.class}"
+            truncate_prompt_main(main_str)
+          )
         when "l"
           ltype
         when "i"

test/irb/test_context.rb

@@ -644,6 +644,27 @@ def test_prompt_main_inspect_escape
       assert_equal("irb(main\\n main)>", irb.send(:format_prompt, 'irb(%M)>', nil, 1, 1))
     end
 
+    def test_prompt_part_cached
+      main = Object.new
+      def main.to_s; "to_s#{rand}"; end
+      def main.inspect; "inspect#{rand}"; end
+      irb = IRB::Irb.new(IRB::WorkSpace.new(main), TestInputMethod.new)
+      format = '[%m %M %m %M]>'
+      pattern = /\A\[(to_s[\d.]+) (inspect[\d.]+) \1 \2\]>\z/
+
+      prompt1, prompt2 = nil
+      irb.send(:with_prompt_part_cached) do
+        prompt1 = irb.send(:format_prompt, format, nil, 1, 1)
+        prompt2 = irb.send(:format_prompt, format, nil, 1, 1)
+      end
+      assert_equal(prompt1, prompt2)
+      assert_match(pattern, prompt1)
+
+      prompt3 = irb.send(:format_prompt, format, nil, 1, 1)
+      assert_not_equal(prompt1, prompt3)
+      assert_match(pattern, prompt3)
+    end
+
     def test_prompt_main_truncate
       main = Struct.new(:to_s).new("a" * 100)
       def main.inspect; to_s.inspect; end

test/irb/yamatanooroti/test_rendering.rb

@@ -112,6 +112,36 @@ def test_nomultiline
     close
   end
 
+  def test_main_to_s_call_cached
+    start_terminal(25, 80, %W{ruby -I#{@pwd}/lib #{@pwd}/exe/irb}, startup_message: /irb\(main\)/)
+    write(<<~'EOC')
+      @count = 0
+      def self.to_s; @count += 1; "[#{@count}]"; end
+      if false
+        123
+      end
+      if false
+        123
+      end
+    EOC
+    assert_screen(<<~'EOC')
+      irb(main):001> @count = 0
+      => 0
+      irb(main):002> def self.to_s; @count += 1; "[#{@count}]"; end
+      => :to_s
+      irb([1]):003* if false
+      irb([1]):004*   123
+      irb([1]):005> end
+      => nil
+      irb([2]):006* if false
+      irb([2]):007*   123
+      irb([2]):008> end
+      => nil
+      irb([3]):009>
+    EOC
+    close
+  end
+
   def test_multiline_paste
     start_terminal(25, 80, %W{ruby -I#{@pwd}/lib #{@pwd}/exe/irb}, startup_message: /irb\(main\)/)
     write(<<~EOC)