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

In prompt calculation, `main.to_s` was called on every keystroke and every line in multiline input.
Cache prompt parts(%m, %M) so that `main.to_s` is only called once per read-eval cycle.

Author: tompng

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
@@ -239,6 +240,7 @@ def read_input(prompt)
     end
 
     def readmultiline
+      @prompt_part_cache = {}
       prompt = generate_prompt([], false, 0)
 
       # multiline
@@ -263,6 +265,8 @@ def readmultiline
         continue = @scanner.should_continue?(tokens)
         prompt = generate_prompt(opens, continue, line_offset)
       end
+    ensure
+      @prompt_part_cache = nil
     end
 
     def each_top_level_statement
@@ -598,25 +602,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,26 @@ 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/
+
+      irb.instance_variable_set(:@prompt_part_cache, {})
+      prompt1 = irb.send(:format_prompt, format, nil, 1, 1)
+      prompt2 = irb.send(:format_prompt, format, nil, 1, 1)
+      assert_equal(prompt1, prompt2)
+      assert_match(pattern, prompt1)
+
+      irb.instance_variable_set(:@prompt_part_cache, nil)
+      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_prompt_main_part_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)