Simplify each_top_level_statement

lib/irb/ruby-lex.rb

@@ -18,10 +18,7 @@ def initialize
 
   def initialize(context)
     @context = context
-    @exp_line_no = @line_no = 1
-    @indent = 0
-    @continue = false
-    @line = ""
+    @line_no = 1
     @prompt = nil
   end
 
@@ -42,6 +39,11 @@ def self.compile_with_errors_suppressed(code, line_no: 1)
     result
   end
 
+  def single_line_command?(code)
+    command = code.split(/\s/, 2).first
+    @context.symbol_alias?(command) || @context.transform_args?(command)
+  end
+
   # io functions
   def set_input(&block)
     @input = block
@@ -65,14 +67,9 @@ def configure_io(io)
           end
         else
           # Accept any single-line input for symbol aliases or commands that transform args
-          command = code.split(/\s/, 2).first
-          if @context.symbol_alias?(command) || @context.transform_args?(command)
-            next true
-          end
+          next true if single_line_command?(code)
 
-          code.gsub!(/\s*\z/, '').concat("\n")
-          tokens = self.class.ripper_lex_without_warning(code, context: @context)
-          ltype, indent, continue, code_block_open = check_state(code, tokens)
+          ltype, indent, continue, code_block_open = check_code_state(code)
           if ltype or indent > 0 or continue or code_block_open
             false
           else
@@ -210,67 +207,56 @@ def check_state(code, tokens)
     [ltype, indent, continue, code_block_open]
   end
 
-  def prompt
-    if @prompt
-      @prompt.call(@ltype, @indent, @continue, @line_no)
-    end
+  def check_code_state(code)
+    check_target_code = code.gsub(/\s*\z/, '').concat("\n")
+    tokens = self.class.ripper_lex_without_warning(check_target_code, context: @context)
+    check_state(check_target_code, tokens)
   end
 
-  def initialize_input
-    @ltype = nil
-    @indent = 0
-    @continue = false
-    @line = ""
-    @exp_line_no = @line_no
-    @code_block_open = false
+  def save_prompt_to_context_io(ltype, indent, continue, line_num_offset)
+    # Implicitly saves prompt string to `@context.io.prompt`. This will be used in the next `@input.call`.
+    @prompt.call(ltype, indent, continue, @line_no + line_num_offset)
   end
 
-  def each_top_level_statement
-    initialize_input
-    catch(:TERM_INPUT) do
-      loop do
-        begin
-          prompt
-          unless l = lex
-            throw :TERM_INPUT if @line == ''
-          else
-            @line_no += l.count("\n")
-            if l == "\n"
-              @exp_line_no += 1
-              next
-            end
-            @line.concat l
-            if @code_block_open or @ltype or @continue or @indent > 0
-              next
-            end
-          end
-          if @line != "\n"
-            @line.force_encoding(@io.encoding)
-            yield @line, @exp_line_no
-          end
-          raise TerminateLineInput if @io.eof?
-          @line = ''
-          @exp_line_no = @line_no
-
-          @indent = 0
-        rescue TerminateLineInput
-          initialize_input
-          prompt
-        end
+  def readmultiline
+    save_prompt_to_context_io(nil, 0, false, 0)
+
+    # multiline
+    return @input.call if @io.respond_to?(:check_termination)
+
+    # nomultiline
+    code = ''
+    line_offset = 0
+    loop do
+      line = @input.call
+      unless line
+        return code.empty? ? nil : code
       end
+
+      code << line
+      # Accept any single-line input for symbol aliases or commands that transform args
+      return code if single_line_command?(code)
+
+      ltype, indent, continue, code_block_open = check_code_state(code)
+      return code unless ltype or indent > 0 or continue or code_block_open
+
+      line_offset += 1
+      save_prompt_to_context_io(ltype, indent, continue, line_offset)
     end
   end
 
-  def lex
-    line = @input.call
-    if @io.respond_to?(:check_termination)
-      return line # multiline
+  def each_top_level_statement
+    loop do
+      code = readmultiline
+      break unless code
+
+      if code != "\n"
+        code.force_encoding(@io.encoding)
+        yield code, @line_no
+      end
+      @line_no += code.count("\n")
+    rescue TerminateLineInput
     end
-    code = @line + (line.nil? ? '' : line)
-    code.gsub!(/\s*\z/, '').concat("\n")
-    @tokens = self.class.ripper_lex_without_warning(code, context: @context)
-    @ltype, @indent, @continue, @code_block_open = check_state(code, @tokens)
-    line
   end
 
   def process_continue(tokens)

test/irb/test_ruby_lex.rb

@@ -83,27 +83,22 @@ def assert_row_indenting(lines, row)
     end
 
     def assert_nesting_level(lines, expected, local_variables: [])
-      ruby_lex = ruby_lex_for_lines(lines, local_variables: local_variables)
+      indent, _code_block_open = check_state(lines, local_variables: local_variables)
       error_message = "Calculated the wrong number of nesting level for:\n #{lines.join("\n")}"
-      assert_equal(expected, ruby_lex.instance_variable_get(:@indent), error_message)
+      assert_equal(expected, indent, error_message)
     end
 
     def assert_code_block_open(lines, expected, local_variables: [])
-      ruby_lex = ruby_lex_for_lines(lines, local_variables: local_variables)
+      _indent, code_block_open = check_state(lines, local_variables: local_variables)
       error_message = "Wrong result of code_block_open for:\n #{lines.join("\n")}"
-      assert_equal(expected, ruby_lex.instance_variable_get(:@code_block_open), error_message)
+      assert_equal(expected, code_block_open, error_message)
     end
 
-    def ruby_lex_for_lines(lines, local_variables: [])
+    def check_state(lines, local_variables: [])
       context = build_context(local_variables)
       ruby_lex = RubyLex.new(context)
-
-      io = proc{ lines.join("\n") }
-      ruby_lex.set_input do
-        lines.join("\n")
-      end
-      ruby_lex.lex
-      ruby_lex
+      _ltype, indent, _continue, code_block_open = ruby_lex.check_code_state(lines.join("\n"))
+      [indent, code_block_open]
     end
 
     def test_auto_indent