Powerup show_source by enabling RubyVM.keep_script_lines

lib/irb.rb

@@ -979,12 +979,16 @@ def run(conf = IRB.conf)
       end
 
       begin
-        forced_exit = false
+        if defined?(RubyVM.keep_script_lines)
+          keep_script_lines_backup = RubyVM.keep_script_lines
+          RubyVM.keep_script_lines = true
+        end
 
         forced_exit = catch(:IRB_EXIT) do
           eval_input
         end
       ensure
+        RubyVM.keep_script_lines = keep_script_lines_backup if defined?(RubyVM.keep_script_lines)
         trap("SIGINT", prev_trap)
         conf[:AT_EXIT].each{|hook| hook.call}
 

lib/irb/cmd/edit.rb

@@ -24,11 +24,9 @@ def transform_args(args)
       def execute(*args)
         path = args.first
 
-        if path.nil? && (irb_path = @irb_context.irb_path)
-          path = irb_path
-        end
-
-        if !File.exist?(path)
+        if path.nil?
+          path = @irb_context.irb_path
+        elsif !File.exist?(path)
           source =
             begin
               SourceFinder.new(@irb_context).find_source(path)
@@ -37,14 +35,16 @@ def execute(*args)
               # in this case, we should just ignore the error
             end
 
-          if source
+          if source&.file_exist? && !source.binary_file?
             path = source.file
-          else
-            puts "Can not find file: #{path}"
-            return
           end
         end
 
+        unless File.exist?(path)
+          puts "Can not find file: #{path}"
+          return
+        end
+
         if editor = (ENV['VISUAL'] || ENV['EDITOR'])
           puts "command: '#{editor}'"
           puts "   path: #{path}"

lib/irb/cmd/show_source.rb

@@ -45,15 +45,18 @@ def execute(str = nil)
       private
 
       def show_source(source)
-        file_content = IRB::Color.colorize_code(File.read(source.file))
-        code = file_content.lines[(source.first_line - 1)...source.last_line].join
-        content = <<~CONTENT
+        if source.binary_file?
+          content = "\n#{bold('Defined in binary file')}: #{source.file}\n\n"
+        else
+          code = source.colorized_content || 'Source not available'
+          content = <<~CONTENT
 
-          #{bold("From")}: #{source.file}:#{source.first_line}
+            #{bold("From")}: #{source.file}:#{source.line}
 
-          #{code}
-        CONTENT
+            #{code.chomp}
 
+          CONTENT
+        end
         Pager.page_content(content)
       end
 

lib/irb/context.rb

@@ -557,7 +557,7 @@ def evaluate(line, line_no) # :nodoc:
 
       if IRB.conf[:MEASURE] && !IRB.conf[:MEASURE_CALLBACKS].empty?
         last_proc = proc do
-          result = @workspace.evaluate(line, irb_path, line_no)
+          result = @workspace.evaluate(line, eval_path, line_no)
         end
         IRB.conf[:MEASURE_CALLBACKS].inject(last_proc) do |chain, item|
           _name, callback, arg = item
@@ -568,12 +568,20 @@ def evaluate(line, line_no) # :nodoc:
           end
         end.call
       else
-        result = @workspace.evaluate(line, irb_path, line_no)
+        result = @workspace.evaluate(line, eval_path, line_no)
       end
 
       set_last_value(result)
     end
 
+    private def eval_path
+      # We need to use differente path to distinguish source_location of method defined in the actual file and method defined in irb session.
+      if !defined?(@irb_path_existence) || @irb_path_existence[0] != irb_path
+        @irb_path_existence = [irb_path, File.exist?(irb_path)]
+      end
+      @irb_path_existence[1] ? "#{irb_path}(#{IRB.conf[:IRB_NAME]})" : irb_path
+    end
+
     def inspect_last_value # :nodoc:
       @inspect_method.inspect_value(@last_value)
     end

lib/irb/source_finder.rb

@@ -4,12 +4,58 @@
 
 module IRB
   class SourceFinder
-    Source = Struct.new(
-      :file,       # @param [String] - file name
-      :first_line, # @param [String] - first line
-      :last_line,  # @param [String] - last line
-      keyword_init: true,
-    )
+    class Source
+      attr_reader :file, :line
+      def initialize(file, line, ast_source = nil)
+        @file = file
+        @line = line
+        @ast_source = ast_source
+      end
+
+      def file_exist?
+        File.exist?(@file)
+      end
+
+      def binary_file?
+        # If the line is zero, it means that the target's source is probably in a binary file.
+        @line.zero?
+      end
+
+      def file_content
+        @file_content ||= File.read(@file)
+      end
+
+      def colorized_content
+        if !binary_file? && file_exist?
+          end_line = Source.find_end(file_content, @line)
+          # To correctly colorize, we need to colorize full content and extract the relevant lines.
+          colored = IRB::Color.colorize_code(file_content)
+          colored.lines[@line - 1...end_line].join
+        elsif @ast_source
+          IRB::Color.colorize_code(@ast_source)
+        end
+      end
+
+      def self.find_end(code, first_line)
+        lex = RubyLex.new
+        lines = code.lines[(first_line - 1)..-1]
+        tokens = RubyLex.ripper_lex_without_warning(lines.join)
+        prev_tokens = []
+
+        # chunk with line number
+        tokens.chunk { |tok| tok.pos[0] }.each do |lnum, chunk|
+          code = lines[0..lnum].join
+          prev_tokens.concat chunk
+          continue = lex.should_continue?(prev_tokens)
+          syntax = lex.check_code_syntax(code, local_variables: [])
+          if !continue && syntax == :valid
+            return first_line + lnum
+          end
+        end
+        first_line
+      end
+    end
+
     private_constant :Source
 
     def initialize(irb_context)
@@ -27,40 +73,28 @@ def find_source(signature, super_level = 0)
         owner = eval(Regexp.last_match[:owner], context_binding)
         method = Regexp.last_match[:method]
         return unless owner.respond_to?(:instance_method)
-        file, line = method_target(owner, super_level, method, "owner")
+        method = method_target(owner, super_level, method, "owner")
+        file, line = method&.source_location
       when /\A((?<receiver>.+)(\.|::))?(?<method>[^ :.]+)\z/ # method, receiver.method, receiver::method
         receiver = eval(Regexp.last_match[:receiver] || 'self', context_binding)
         method = Regexp.last_match[:method]
         return unless receiver.respond_to?(method, true)
-        file, line = method_target(receiver, super_level, method, "receiver")
+        method = method_target(receiver, super_level, method, "receiver")
+        file, line = method&.source_location
       end
-      # If the line is zero, it means that the target's source is probably in a binary file, which we should ignore.
-      if file && line && !line.zero? && File.exist?(file)
-        Source.new(file: file, first_line: line, last_line: find_end(file, line))
+      return unless file && line
+
+      if File.exist?(file)
+        Source.new(file, line)
+      elsif method
+        # Method defined with eval, probably in IRB session
+        source = RubyVM::AbstractSyntaxTree.of(method)&.source rescue nil
+        Source.new(file, line, source)
       end
     end
 
     private
 
-    def find_end(file, first_line)
-      lex = RubyLex.new
-      lines = File.read(file).lines[(first_line - 1)..-1]
-      tokens = RubyLex.ripper_lex_without_warning(lines.join)
-      prev_tokens = []
-
-      # chunk with line number
-      tokens.chunk { |tok| tok.pos[0] }.each do |lnum, chunk|
-        code = lines[0..lnum].join
-        prev_tokens.concat chunk
-        continue = lex.should_continue?(prev_tokens)
-        syntax = lex.check_code_syntax(code, local_variables: [])
-        if !continue && syntax == :valid
-          return first_line + lnum
-        end
-      end
-      first_line
-    end
-
     def method_target(owner_receiver, super_level, method, type)
       case type
       when "owner"
@@ -71,7 +105,7 @@ def method_target(owner_receiver, super_level, method, type)
       super_level.times do |s|
         target_method = target_method.super_method if target_method
       end
-      target_method.nil? ? nil : target_method.source_location
+      target_method
     rescue NameError
       nil
     end

test/irb/cmd/test_show_source.rb

@@ -301,7 +301,37 @@ class Bar
       assert_match(%r[#{@ruby_file.to_path}:5\s+class Bar\r\n  end], out)
     end
 
-    def test_show_source_ignores_binary_source_file
+    def test_show_source_keep_script_lines
+      pend unless defined?(RubyVM.keep_script_lines)
+
+      write_ruby <<~RUBY
+        binding.irb
+      RUBY
+
+      out = run_ruby_file do
+        type "def foo; end"
+        type "show_source foo"
+        type "exit"
+      end
+
+      assert_match(%r[#{@ruby_file.to_path}\(irb\):1\s+def foo; end], out)
+    end
+
+    def test_show_source_unavailable_source
+      write_ruby <<~RUBY
+        binding.irb
+      RUBY
+
+      out = run_ruby_file do
+        type "RubyVM.keep_script_lines = false if defined?(RubyVM.keep_script_lines)"
+        type "def foo; end"
+        type "show_source foo"
+        type "exit"
+      end
+      assert_match(%r[#{@ruby_file.to_path}\(irb\):2\s+Source not available], out)
+    end
+
+    def test_show_source_shows_binary_source
       write_ruby <<~RUBY
         # io-console is an indirect dependency of irb
         require "io/console"
@@ -317,7 +347,7 @@ def test_show_source_ignores_binary_source_file
 
       # A safeguard to make sure the test subject is actually defined
       refute_match(/NameError/, out)
-      assert_match(%r[Error: Couldn't locate a definition for IO::ConsoleMode], out)
+      assert_match(%r[Defined in binary file:.+io/console], out)
     end
   end
 end

test/irb/test_cmd.rb

@@ -848,6 +848,16 @@ def test_edit_without_arg
       assert_match("command: ': code'", out)
     end
 
+    def test_edit_without_arg_and_non_existing_irb_path
+      out, err = execute_lines(
+        "edit",
+        irb_path: '/path/to/file.rb(irb)'
+      )
+
+      assert_empty err
+      assert_match(/Can not find file: \/path\/to\/file\.rb\(irb\)/, out)
+    end
+
     def test_edit_with_path
       out, err = execute_lines(
         "edit #{__FILE__}"

test/irb/test_context.rb

@@ -666,6 +666,15 @@ def test_lineno
         ], out)
     end
 
+    def test_eval_path
+      @context.irb_path = __FILE__
+      assert_equal("#{__FILE__}(irb)", @context.send(:eval_path))
+      @context.irb_path = 'file/does/not/exist'
+      assert_equal('file/does/not/exist', @context.send(:eval_path))
+      @context.irb_path = "#{__FILE__}(irb)"
+      assert_equal("#{__FILE__}(irb)", @context.send(:eval_path))
+    end
+
     def test_build_completor
       verbose, $VERBOSE = $VERBOSE, nil
       original_completor = IRB.conf[:COMPLETOR]