[ruby/irb] Always use local variables in current context to parse code (https://github.com/ruby/irb/pull/397)

* Use local_variables for colorize, code_block_open check, nesting_level and assignment_expression check

* Check if expression is an assignment BEFORE evaluating it. evaluate might define new localvars and change result of assignment_expression?

* Add local_variables dependent code test

* pend local variable dependent test on truffleruby

code_block_open is not working on truffleruby

* Always pass context to RubyLex#lex

* Rename local_variable_assign_code generator method name

* Add assignment expression truncate test

* Add Context#local_variables and make generate_local_variables_assign_code more simple

* Update lib/irb/input-method.rb

Co-authored-by: Stan Lo <stan001212@gmail.com>

* Add a comment why assignment expression check should be done before evaluate

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

Co-authored-by: Stan Lo <stan001212@gmail.com>
Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

9 files changed, 109 insertions, 28 deletions

diff --git a/lib/irb.rb b/lib/irb.rb
index 0c145069c0..749f3ee167 100644
--- a/lib/irb.rb
+++ b/lib/irb.rb
@@ -506,13 +506,15 @@ module IRB
 
       @scanner.set_auto_indent(@context) if @context.auto_indent_mode
 
-      @scanner.each_top_level_statement do |line, line_no|
+      @scanner.each_top_level_statement(@context) do |line, line_no|
         signal_status(:IN_EVAL) do
           begin
             line.untaint if RUBY_VERSION < '2.7'
             if IRB.conf[:MEASURE] && IRB.conf[:MEASURE_CALLBACKS].empty?
               IRB.set_measure_callback
             end
+            # Assignment expression check should be done before @context.evaluate to handle code like `a /2#/ if false; a = 1`
+            is_assignment = assignment_expression?(line)
             if IRB.conf[:MEASURE] && !IRB.conf[:MEASURE_CALLBACKS].empty?
               result = nil
               last_proc = proc{ result = @context.evaluate(line, line_no, exception: exc) }
@@ -529,7 +531,7 @@ module IRB
               @context.evaluate(line, line_no, exception: exc)
             end
             if @context.echo?
-              if assignment_expression?(line)
+              if is_assignment
                 if @context.echo_on_assignment?
                   output_value(@context.echo_on_assignment? == :truncate)
                 end
@@ -827,9 +829,12 @@ module IRB
       # array of parsed expressions. The first element of each expression is the
       # expression's type.
       verbose, $VERBOSE = $VERBOSE, nil
-      result = ASSIGNMENT_NODE_TYPES.include?(Ripper.sexp(line)&.dig(1,-1,0))
+      code = "#{RubyLex.generate_local_variables_assign_code(@context.local_variables) || 'nil;'}\n#{line}"
+      # Get the last node_type of the line. drop(1) is to ignore the local_variables_assign_code part.
+      node_type = Ripper.sexp(code)&.dig(1)&.drop(1)&.dig(-1, 0)
+      ASSIGNMENT_NODE_TYPES.include?(node_type)
+    ensure
       $VERBOSE = verbose
-      result
     end
 
     ATTR_TTY = "\e[%sm"
diff --git a/lib/irb/color.rb b/lib/irb/color.rb
index 7071696cb2..34912420e4 100644
--- a/lib/irb/color.rb
+++ b/lib/irb/color.rb
@@ -123,13 +123,15 @@ module IRB # :nodoc:
       # If `complete` is false (code is incomplete), this does not warn compile_error.
       # This option is needed to avoid warning a user when the compile_error is happening
       # because the input is not wrong but just incomplete.
-      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?)
+      def colorize_code(code, complete: true, ignore_error: false, colorable: colorable?, local_variables: [])
         return code unless colorable
 
         symbol_state = SymbolState.new
         colored = +''
+        lvars_code = RubyLex.generate_local_variables_assign_code(local_variables)
+        code_with_lvars = lvars_code ? "#{lvars_code}\n#{code}" : code
 
-        scan(code, allow_last_error: !complete) do |token, str, expr|
+        scan(code_with_lvars, allow_last_error: !complete) do |token, str, expr|
           # handle uncolorable code
           if token.nil?
             colored << Reline::Unicode.escape_for_print(str)
@@ -152,7 +154,12 @@ module IRB # :nodoc:
             end
           end
         end
-        colored
+
+        if lvars_code
+          colored.sub(/\A.+\n/, '')
+        else
+          colored
+        end
       end
 
       private
diff --git a/lib/irb/context.rb b/lib/irb/context.rb
index 5e07f5dfb0..d238da9350 100644
--- a/lib/irb/context.rb
+++ b/lib/irb/context.rb
@@ -518,5 +518,9 @@ module IRB
     end
     alias __to_s__ to_s
     alias to_s inspect
+
+    def local_variables # :nodoc:
+      workspace.binding.local_variables
+    end
   end
 end
diff --git a/lib/irb/input-method.rb b/lib/irb/input-method.rb
index aa5cb5adb9..b0110dd09b 100644
--- a/lib/irb/input-method.rb
+++ b/lib/irb/input-method.rb
@@ -286,7 +286,8 @@ module IRB
         if IRB.conf[:USE_COLORIZE]
           proc do |output, complete: |
             next unless IRB::Color.colorable?
-            IRB::Color.colorize_code(output, complete: complete)
+            lvars = IRB.CurrentContext&.local_variables || []
+            IRB::Color.colorize_code(output, complete: complete, local_variables: lvars)
           end
         else
           proc do |output|
diff --git a/lib/irb/ruby-lex.rb b/lib/irb/ruby-lex.rb
index cb6d669a72..54ea2a9e7b 100644
--- a/lib/irb/ruby-lex.rb
+++ b/lib/irb/ruby-lex.rb
@@ -136,16 +136,18 @@ class RubyLex
     :on_param_error
   ]
 
+  def self.generate_local_variables_assign_code(local_variables)
+    "#{local_variables.join('=')}=nil;" unless local_variables.empty?
+  end
+
   def self.ripper_lex_without_warning(code, context: nil)
     verbose, $VERBOSE = $VERBOSE, nil
-    if context
-      lvars = context.workspace&.binding&.local_variables
-      if lvars && !lvars.empty?
-        code = "#{lvars.join('=')}=nil\n#{code}"
-        line_no = 0
-      else
-        line_no = 1
-      end
+    lvars_code = generate_local_variables_assign_code(context&.local_variables || [])
+    if lvars_code
+      code = "#{lvars_code}\n#{code}"
+      line_no = 0
+    else
+      line_no = 1
     end
 
     compile_with_errors_suppressed(code, line_no: line_no) do |inner_code, line_no|
@@ -214,6 +216,8 @@ class RubyLex
     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 || [])
+    code = "#{lvars_code}\n#{code}" if lvars_code
     code_block_open = check_code_block(code, tokens)
     [ltype, indent, continue, code_block_open]
   end
@@ -233,13 +237,13 @@ class RubyLex
     @code_block_open = false
   end
 
-  def each_top_level_statement
+  def each_top_level_statement(context)
     initialize_input
     catch(:TERM_INPUT) do
       loop do
         begin
           prompt
-          unless l = lex
+          unless l = lex(context)
             throw :TERM_INPUT if @line == ''
           else
             @line_no += l.count("\n")
@@ -269,18 +273,15 @@ class RubyLex
     end
   end
 
-  def lex
+  def lex(context)
     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)
-    @continue = process_continue
-    @code_block_open = check_code_block(code)
-    @indent = process_nesting_level
-    @ltype = process_literal_type
+    @tokens = self.class.ripper_lex_without_warning(code, context: context)
+    @ltype, @indent, @continue, @code_block_open = check_state(code, @tokens, context: context)
     line
   end
 
diff --git a/test/irb/test_color.rb b/test/irb/test_color.rb
index 34af4e2bea..02c79e3443 100644
--- a/test/irb/test_color.rb
+++ b/test/irb/test_color.rb
@@ -156,6 +156,17 @@ module TestIRB
       end
     end
 
+    def test_colorize_code_with_local_variables
+      code = "a /(b +1)/i"
+      result_without_lvars = "a #{RED}#{BOLD}/#{CLEAR}#{RED}(b +1)#{CLEAR}#{RED}#{BOLD}/i#{CLEAR}"
+      result_with_lvar = "a /(b #{BLUE}#{BOLD}+1#{CLEAR})/i"
+      result_with_lvars = "a /(b +#{BLUE}#{BOLD}1#{CLEAR})/i"
+
+      assert_equal_with_term(result_without_lvars, code)
+      assert_equal_with_term(result_with_lvar, code, local_variables: ['a'])
+      assert_equal_with_term(result_with_lvars, code, local_variables: ['a', 'b'])
+    end
+
     def test_colorize_code_complete_true
       unless complete_option_supported?
         pend '`complete: true` is the same as `complete: false` in Ruby 2.6-'
diff --git a/test/irb/test_context.rb b/test/irb/test_context.rb
index b3fc49e7e3..998cdd8591 100644
--- a/test/irb/test_context.rb
+++ b/test/irb/test_context.rb
@@ -225,6 +225,16 @@ module TestIRB
       end
     end
 
+    def test_assignment_expression_with_local_variable
+      input = TestInputMethod.new
+      irb = IRB::Irb.new(IRB::WorkSpace.new(Object.new), input)
+      code = "a /1;x=1#/"
+      refute(irb.assignment_expression?(code), "#{code}: should not be an assignment expression")
+      irb.context.workspace.binding.eval('a = 1')
+      assert(irb.assignment_expression?(code), "#{code}: should be an assignment expression")
+      refute(irb.assignment_expression?(""), "empty code should not be an assignment expression")
+    end
+
     def test_echo_on_assignment
       input = TestInputMethod.new([
         "a = 1\n",
diff --git a/test/irb/test_ruby_lex.rb b/test/irb/test_ruby_lex.rb
index beda53fc89..1388d08962 100644
--- a/test/irb/test_ruby_lex.rb
+++ b/test/irb/test_ruby_lex.rb
@@ -34,13 +34,27 @@ module TestIRB
       ruby_lex.set_auto_indent(context)
     end
 
-    def assert_nesting_level(lines, expected)
+    def assert_nesting_level(lines, expected, local_variables: [])
+      ruby_lex = ruby_lex_for_lines(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)
+    end
+
+    def assert_code_block_open(lines, expected, local_variables: [])
+      ruby_lex = ruby_lex_for_lines(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)
+    end
+
+    def ruby_lex_for_lines(lines, local_variables: [])
       ruby_lex = RubyLex.new()
       io = proc{ lines.join("\n") }
       ruby_lex.set_input(io, io)
-      ruby_lex.lex
-      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)
+      unless local_variables.empty?
+        context = OpenStruct.new(local_variables: local_variables)
+      end
+      ruby_lex.lex(context)
+      ruby_lex
     end
 
     def test_auto_indent
@@ -514,6 +528,15 @@ module TestIRB
       end
     end
 
+    def test_local_variables_dependent_code
+      pend if RUBY_ENGINE == 'truffleruby'
+      lines = ["a /1#/ do", "2"]
+      assert_nesting_level(lines, 1)
+      assert_code_block_open(lines, true)
+      assert_nesting_level(lines, 0, local_variables: ['a'])
+      assert_code_block_open(lines, false, local_variables: ['a'])
+    end
+
     def test_heredoc_with_indent
       input_with_correct_indents = [
         Row.new(%q(<<~Q), 0, 0, 0),
diff --git a/test/irb/yamatanooroti/test_rendering.rb b/test/irb/yamatanooroti/test_rendering.rb
index 7ed98b11c1..f9a130b7d4 100644
--- a/test/irb/yamatanooroti/test_rendering.rb
+++ b/test/irb/yamatanooroti/test_rendering.rb
@@ -216,6 +216,25 @@ begin
       EOC
     end
 
+    def test_assignment_expression_truncate
+      write_irbrc <<~'LINES'
+        puts 'start IRB'
+      LINES
+      start_terminal(40, 80, %W{ruby -I#{@pwd}/lib #{@pwd}/exe/irb}, startup_message: 'start IRB')
+      # Assignment expression code that turns into non-assignment expression after evaluation
+      code = "a /'/i if false; a=1; x=1000.times.to_a#'.size"
+      write(code + "\n")
+      close
+      assert_screen(<<~EOC)
+        start IRB
+        irb(main):001:0> #{code}
+        =>
+        [0,
+        ...
+        irb(main):002:0>
+      EOC
+    end
+
     private def write_irbrc(content)
       File.open(@irbrc_file, 'w') do |f|
         f.write content