[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>

4 files changed, 67 insertions, 4 deletions

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