Store context in RubyLex

Some background for this refactor:

1. Through a RubyLex instance's lifetime, the context passed to its methods
   should be the same.
   Given that `Context` is only initialised in `Irb#initialize`,
   this should be true.

2. When `RubyLex` is initialised, the context object should be accessible.
   This is also true in all 3 of `RubyLex.new`'s invocations.

With the above observations, we should be able to store the context in `RubyLex`
as an instance variable. And doing so will make `RubyLex`'s instance methods
easier to use and maintain.

lib/irb.rb

@@ -468,7 +468,7 @@ def initialize(workspace = nil, input_method = nil)
       @context = Context.new(self, workspace, input_method)
       @context.main.extend ExtendCommandBundle
       @signal_status = :IN_IRB
-      @scanner = RubyLex.new
+      @scanner = RubyLex.new(@context)
     end
 
     # A hook point for `debug` command's TracePoint after :IRB_EXIT as well as its clean-up
@@ -538,7 +538,7 @@ def eval_input
         @context.io.prompt
       end
 
-      @scanner.set_input(@context.io, context: @context) do
+      @scanner.set_input(@context.io) do
         signal_status(:IN_INPUT) do
           if l = @context.io.gets
             print l if @context.verbose?
@@ -556,9 +556,9 @@ def eval_input
         end
       end
 
-      @scanner.set_auto_indent(@context) if @context.auto_indent_mode
+      @scanner.set_auto_indent
 
-      @scanner.each_top_level_statement(@context) do |line, line_no|
+      @scanner.each_top_level_statement do |line, line_no|
         signal_status(:IN_EVAL) do
           begin
             line.untaint if RUBY_VERSION < '2.7'

lib/irb/cmd/show_source.rb

@@ -40,15 +40,15 @@ def find_source(str, irb_context)
             file, line = receiver.method(method).source_location if receiver.respond_to?(method)
           end
           if file && line
-            Source.new(file: file, first_line: line, last_line: find_end(file, line))
+            Source.new(file: file, first_line: line, last_line: find_end(file, line, irb_context))
           end
         end
 
         private
 
-        def find_end(file, first_line)
+        def find_end(file, first_line, irb_context)
           return first_line unless File.exist?(file)
-          lex = RubyLex.new
+          lex = RubyLex.new(irb_context)
           lines = File.read(file).lines[(first_line - 1)..-1]
           tokens = RubyLex.ripper_lex_without_warning(lines.join)
           prev_tokens = []

lib/irb/ruby-lex.rb

@@ -16,7 +16,8 @@ def initialize
     end
   end
 
-  def initialize
+  def initialize(context)
+    @context = context
     @exp_line_no = @line_no = 1
     @indent = 0
     @continue = false
@@ -42,13 +43,13 @@ def self.compile_with_errors_suppressed(code, line_no: 1)
   end
 
   # io functions
-  def set_input(io, context:, &block)
+  def set_input(io, &block)
     @io = io
     if @io.respond_to?(:check_termination)
       @io.check_termination do |code|
         if Reline::IOGate.in_pasting?
-          lex = RubyLex.new
-          rest = lex.check_termination_in_prev_line(code, context: context)
+          lex = RubyLex.new(@context)
+          rest = lex.check_termination_in_prev_line(code)
           if rest
             Reline.delete_text
             rest.bytes.reverse_each do |c|
@@ -61,13 +62,13 @@ def set_input(io, context:, &block)
         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)
+          if @context.symbol_alias?(command) || @context.transform_args?(command)
             next true
           end
 
           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, context: context)
+          tokens = self.class.ripper_lex_without_warning(code, context: @context)
+          ltype, indent, continue, code_block_open = check_state(code, tokens)
           if ltype or indent > 0 or continue or code_block_open
             false
           else
@@ -80,7 +81,7 @@ def set_input(io, context:, &block)
       @io.dynamic_prompt do |lines|
         lines << '' if lines.empty?
         result = []
-        tokens = self.class.ripper_lex_without_warning(lines.map{ |l| l + "\n" }.join, context: context)
+        tokens = self.class.ripper_lex_without_warning(lines.map{ |l| l + "\n" }.join, context: @context)
         code = String.new
         partial_tokens = []
         unprocessed_tokens = []
@@ -93,7 +94,7 @@ def set_input(io, context:, &block)
             t_str.each_line("\n") do |s|
               code << s
               next unless s.include?("\n")
-              ltype, indent, continue, code_block_open = check_state(code, partial_tokens, context: context)
+              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
@@ -104,7 +105,7 @@ def set_input(io, context:, &block)
         end
 
         unless unprocessed_tokens.empty?
-          ltype, indent, continue, code_block_open = check_state(code, unprocessed_tokens, context: context)
+          ltype, indent, continue, code_block_open = check_state(code, unprocessed_tokens)
           result << @prompt.call(ltype, indent, continue || code_block_open, @line_no + line_num_offset)
         end
         result
@@ -187,11 +188,11 @@ def find_prev_spaces(line_index)
     prev_spaces
   end
 
-  def set_auto_indent(context)
-    if @io.respond_to?(:auto_indent) and context.auto_indent_mode
+  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)
+          @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
@@ -200,18 +201,18 @@ def set_auto_indent(context)
           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)
+          @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, context:)
+  def check_state(code, tokens)
     ltype = process_literal_type(tokens)
     indent = process_nesting_level(tokens)
     continue = process_continue(tokens)
-    lvars_code = self.class.generate_local_variables_assign_code(context.local_variables)
+    lvars_code = self.class.generate_local_variables_assign_code(@context.local_variables)
     code = "#{lvars_code}\n#{code}" if lvars_code
     code_block_open = check_code_block(code, tokens)
     [ltype, indent, continue, code_block_open]
@@ -232,13 +233,13 @@ def initialize_input
     @code_block_open = false
   end
 
-  def each_top_level_statement(context)
+  def each_top_level_statement
     initialize_input
     catch(:TERM_INPUT) do
       loop do
         begin
           prompt
-          unless l = lex(context)
+          unless l = lex
             throw :TERM_INPUT if @line == ''
           else
             @line_no += l.count("\n")
@@ -268,15 +269,15 @@ def each_top_level_statement(context)
     end
   end
 
-  def lex(context)
+  def lex
     line = @input.call
     if @io.respond_to?(:check_termination)
       return line # multiline
     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, context: context)
+    @tokens = self.class.ripper_lex_without_warning(code, context: @context)
+    @ltype, @indent, @continue, @code_block_open = check_state(code, @tokens)
     line
   end
 
@@ -777,8 +778,8 @@ def process_literal_type(tokens)
     end
   end
 
-  def check_termination_in_prev_line(code, context:)
-    tokens = self.class.ripper_lex_without_warning(code, context: context)
+  def check_termination_in_prev_line(code)
+    tokens = self.class.ripper_lex_without_warning(code, context: @context)
     past_first_newline = false
     index = tokens.rindex do |t|
       # traverse first token before last line

test/irb/test_ruby_lex.rb

@@ -36,13 +36,13 @@ def assert_indenting(lines, correct_space_count, add_new_line)
 
       context = build_context
       context.auto_indent_mode = true
-      ruby_lex = RubyLex.new()
+      ruby_lex = RubyLex.new(context)
       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
-      ruby_lex.set_input(io, context: context)
-      ruby_lex.set_auto_indent(context)
+      ruby_lex.set_input(io)
+      ruby_lex.set_auto_indent
     end
 
     def assert_nesting_level(lines, expected, local_variables: [])
@@ -58,14 +58,14 @@ def assert_code_block_open(lines, expected, local_variables: [])
     end
 
     def ruby_lex_for_lines(lines, local_variables: [])
-      ruby_lex = RubyLex.new()
-
       context = build_context(local_variables)
+      ruby_lex = RubyLex.new(context)
+
       io = proc{ lines.join("\n") }
-      ruby_lex.set_input(io, context: context) do
+      ruby_lex.set_input(io) do
         lines.join("\n")
       end
-      ruby_lex.lex(context)
+      ruby_lex.lex
       ruby_lex
     end
 
@@ -633,7 +633,8 @@ def dynamic_prompt(&block)
 
     def assert_dynamic_prompt(lines, expected_prompt_list)
       pend if RUBY_ENGINE == 'truffleruby'
-      ruby_lex = RubyLex.new()
+      context = build_context
+      ruby_lex = RubyLex.new(context)
       io = MockIO_DynamicPrompt.new(lines) do |prompt_list|
         error_message = <<~EOM
           Expected dynamic prompt:
@@ -647,8 +648,7 @@ def assert_dynamic_prompt(lines, expected_prompt_list)
       ruby_lex.set_prompt do |ltype, indent, continue, line_no|
         '%03d:%01d:%1s:%s ' % [line_no, indent, ltype, continue ? '*' : '>']
       end
-      context = build_context
-      ruby_lex.set_input(io, context: context)
+      ruby_lex.set_input(io)
     end
 
     def test_dyanmic_prompt
@@ -751,9 +751,10 @@ def test_unterminated_code
     end
 
     def test_unterminated_heredoc_string_literal
+      context = build_context
       ['<<A;<<B', "<<A;<<B\n", "%W[\#{<<A;<<B", "%W[\#{<<A;<<B\n"].each do |code|
         tokens = RubyLex.ripper_lex_without_warning(code)
-        string_literal = RubyLex.new.check_string_literal(tokens)
+        string_literal = RubyLex.new(context).check_string_literal(tokens)
         assert_equal('<<A', string_literal&.tok)
       end
     end
@@ -779,16 +780,17 @@ def test_corresponding_token_depth_with_heredoc_and_embdoc
           p(
           )
       EOC
+      context = build_context
       [reference_code, code_with_heredoc, code_with_embdoc].each do |code|
-        lex = RubyLex.new
+        lex = RubyLex.new(context)
         lines = code.lines
         lex.instance_variable_set('@tokens', RubyLex.ripper_lex_without_warning(code))
         assert_equal 2, lex.check_corresponding_token_depth(lines, lines.size)
       end
     end
 
     def test_find_prev_spaces_with_multiline_literal
-      lex = RubyLex.new
+      lex = RubyLex.new(build_context)
       reference_code = <<~EOC.chomp
         if true
           1
@@ -813,8 +815,9 @@ def test_find_prev_spaces_with_multiline_literal
           world
         end
       EOC
+      context = build_context
       [reference_code, code_with_percent_string, code_with_quoted_string].each do |code|
-        lex = RubyLex.new
+        lex = RubyLex.new(context)
         lex.instance_variable_set('@tokens', RubyLex.ripper_lex_without_warning(code))
         prev_spaces = (1..code.lines.size).map { |index| lex.find_prev_spaces index }
         assert_equal [0, 2, 2, 2, 2, 0], prev_spaces