[ruby/irb] Refactor RubyLex's input/io methods

(https://github.com/ruby/irb/pull/583)

1. Make `RubyLex#set_input` simply assign the input block. This matches
   the behavior of `RubyLex#set_prompt`.
2. Merge `RubyLex#set_input`'s IO configuration logic with `#set_auto_indent`
   into `#configure_io`.

3 files changed, 27 insertions, 31 deletions

diff --git a/lib/irb.rb b/lib/irb.rb
index 13fbbbad29..26432a0fac 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -537,7 +537,7 @@ module IRB
         @context.io.prompt
       end
 
-      @scanner.set_input(@context.io) do
+      @scanner.set_input do
         signal_status(:IN_INPUT) do
           if l = @context.io.gets
             print l if @context.verbose?
@@ -555,7 +555,7 @@ module IRB
         end
       end
 
-      @scanner.set_auto_indent
+      @scanner.configure_io(@context.io)
 
       @scanner.each_top_level_statement do |line, line_no|
         signal_status(:IN_EVAL) do
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index a156f3707a..ac734072df 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -43,7 +43,11 @@ class RubyLex
   end
 
   # io functions
-  def set_input(io, &block)
+  def set_input(&block)
+    @input = block
+  end
+
+  def configure_io(io)
     @io = io
     if @io.respond_to?(:check_termination)
       @io.check_termination do |code|
@@ -112,10 +116,22 @@ class RubyLex
       end
     end
 
-    if block_given?
-      @input = block
-    else
-      @input = Proc.new{@io.gets}
+    if @io.respond_to?(:auto_indent) and @context.auto_indent_mode
+      @io.auto_indent do |lines, line_index, byte_pointer, is_newline|
+        if is_newline
+          @tokens = self.class.ripper_lex_without_warning(lines[0..line_index].join("\n"), context: @context)
+          prev_spaces = find_prev_spaces(line_index)
+          depth_difference = check_newline_depth_difference
+          depth_difference = 0 if depth_difference < 0
+          prev_spaces + depth_difference * 2
+        else
+          code = line_index.zero? ? '' : lines[0..(line_index - 1)].map{ |l| l + "\n" }.join
+          last_line = lines[line_index]&.byteslice(0, byte_pointer)
+          code += last_line if last_line
+          @tokens = self.class.ripper_lex_without_warning(code, context: @context)
+          check_corresponding_token_depth(lines, line_index)
+        end
+      end
     end
   end
 
@@ -184,26 +200,6 @@ class RubyLex
     prev_spaces
   end
 
-  def set_auto_indent
-    if @io.respond_to?(:auto_indent) and @context.auto_indent_mode
-      @io.auto_indent do |lines, line_index, byte_pointer, is_newline|
-        if is_newline
-          @tokens = self.class.ripper_lex_without_warning(lines[0..line_index].join("\n"), context: @context)
-          prev_spaces = find_prev_spaces(line_index)
-          depth_difference = check_newline_depth_difference
-          depth_difference = 0 if depth_difference < 0
-          prev_spaces + depth_difference * 2
-        else
-          code = line_index.zero? ? '' : lines[0..(line_index - 1)].map{ |l| l + "\n" }.join
-          last_line = lines[line_index]&.byteslice(0, byte_pointer)
-          code += last_line if last_line
-          @tokens = self.class.ripper_lex_without_warning(code, context: @context)
-          check_corresponding_token_depth(lines, line_index)
-        end
-      end
-    end
-  end
-
   def check_state(code, tokens)
     ltype = process_literal_type(tokens)
     indent = process_nesting_level(tokens)
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
index 5188009eee..5413fe51a8 100644
--- a/test/irb/test_ruby_lex.rb
+++ b/test/irb/test_ruby_lex.rb
@@ -41,8 +41,8 @@ module TestIRB
       ruby_lex = RubyLex.new(context)
       mock_io = MockIO_AutoIndent.new(lines, last_line_index, byte_pointer, add_new_line)
 
-      ruby_lex.set_input(mock_io)
-      ruby_lex.set_auto_indent
+      ruby_lex.set_input { @io.gets }
+      ruby_lex.configure_io(mock_io)
       mock_io.calculated_indent
     end
 
@@ -99,7 +99,7 @@ module TestIRB
       ruby_lex = RubyLex.new(context)
 
       io = proc{ lines.join("\n") }
-      ruby_lex.set_input(io) do
+      ruby_lex.set_input do
         lines.join("\n")
       end
       ruby_lex.lex
@@ -658,7 +658,7 @@ module TestIRB
       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)
+      ruby_lex.set_input { @io.gets }
     end
 
     def test_dyanmic_prompt