[ruby/reline] [ruby/irb] Handle multiple newlines in a token correctly

Co-authored-by: manga_osyo <manga.osyo@gmail.com>
Co-authored-by: ima1zumi <mariimaizumi5@gmail.com>

https://github.com/ruby/irb/commit/c59a9be82f

https://github.com/ruby/reline/commit/a7922da16b

2 files changed, 62 insertions, 6 deletions

diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index 71145224ad..b07027898c 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -66,14 +66,19 @@ class RubyLex
         unprocessed_tokens = []
         line_num_offset = 0
         tokens.each do |t|
-          code << t[2]
           partial_tokens << t
           unprocessed_tokens << t
           if t[2].include?("\n")
-            ltype, indent, continue, code_block_open = check_state(code, partial_tokens)
-            result << @prompt.call(ltype, indent, continue || code_block_open, @line_no + line_num_offset)
-            line_num_offset += 1
+            t_str = t[2]
+            t_str.each_line("\n") do |s|
+              code << s << "\n"
+              ltype, indent, continue, code_block_open = check_state(code, partial_tokens)
+              result << @prompt.call(ltype, indent, continue || code_block_open, @line_no + line_num_offset)
+              line_num_offset += 1
+            end
             unprocessed_tokens = []
+          else
+            code << t[2]
           end
         end
         unless unprocessed_tokens.empty?
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
index 25f037ec3b..255a5a1a63 100644
--- a/test/irb/test_ruby_lex.rb
+++ b/test/irb/test_ruby_lex.rb
@@ -7,7 +7,7 @@ module TestIRB
   class TestRubyLex < Test::Unit::TestCase
     Row = Struct.new(:content, :current_line_spaces, :new_line_spaces, :nesting_level)
 
-    class MockIO
+    class MockIO_AutoIndent
       def initialize(params, &assertion)
         @params = params
         @assertion = assertion
@@ -25,7 +25,7 @@ module TestIRB
       byte_pointer = lines.last.length
 
       ruby_lex = RubyLex.new()
-      io = MockIO.new([lines, last_line_index, byte_pointer, add_new_line]) do |auto_indent|
+      io = MockIO_AutoIndent.new([lines, last_line_index, byte_pointer, add_new_line]) do |auto_indent|
         error_message = "Calculated the wrong number of spaces for:\n #{lines.join("\n")}"
         assert_equal(correct_space_count, auto_indent, error_message)
       end
@@ -262,5 +262,56 @@ module TestIRB
         assert_nesting_level(lines, row.nesting_level)
       end
     end
+
+    PromptRow = Struct.new(:prompt, :content)
+
+    class MockIO_DynamicPrompt
+      def initialize(params, &assertion)
+        @params = params
+        @assertion = assertion
+      end
+
+      def dynamic_prompt(&block)
+        result = block.call(@params)
+        @assertion.call(result)
+      end
+    end
+
+    def assert_dynamic_prompt(lines, expected_prompt_list)
+      skip if RUBY_ENGINE == 'truffleruby'
+      ruby_lex = RubyLex.new()
+      io = MockIO_DynamicPrompt.new(lines) do |prompt_list|
+        error_message = "Calculated the wrong number of spaces for:\n #{lines.join("\n")}"
+        assert_equal(expected_prompt_list, prompt_list)
+      end
+      ruby_lex.set_prompt do |ltype, indent, continue, line_no|
+        '%03d:%01d:%1s:%s ' % [line_no, indent, ltype, continue ? '*' : '>']
+      end
+      ruby_lex.set_input(io)
+    end
+
+    def test_dyanmic_prompt
+      input_with_prompt = [
+        PromptRow.new('001:1: :* ', %q(def hoge)),
+        PromptRow.new('002:1: :* ', %q(  3)),
+        PromptRow.new('003:0: :> ', %q(end)),
+      ]
+
+      lines = input_with_prompt.map(&:content)
+      expected_prompt_list = input_with_prompt.map(&:prompt)
+      assert_dynamic_prompt(lines, expected_prompt_list)
+    end
+
+    def test_dyanmic_prompt_with_blank_line
+      input_with_prompt = [
+        PromptRow.new('001:0:]:* ', %q(%w[)),
+        PromptRow.new('002:0:]:* ', %q()),
+        PromptRow.new('003:0: :> ', %q(])),
+      ]
+
+      lines = input_with_prompt.map(&:content)
+      expected_prompt_list = input_with_prompt.map(&:prompt)
+      assert_dynamic_prompt(lines, expected_prompt_list)
+    end
   end
 end