typeprof 0.15.3 → 0.20.0

data/lib/typeprof/iseq.rb

@@ -2,6 +2,17 @@ module TypeProf
   class ISeq
     # https://github.com/ruby/ruby/pull/4468
     CASE_WHEN_CHECKMATCH = RubyVM::InstructionSequence.compile("case 1; when Integer; end").to_a.last.any? {|insn,| insn == :checkmatch }
+    # https://github.com/ruby/ruby/blob/v3_0_2/vm_core.h#L1206
+    VM_ENV_DATA_SIZE = 3
+    # Check if Ruby 3.1 or later
+    RICH_AST = begin RubyVM::AbstractSyntaxTree.parse("1", keep_script_lines: true).node_id; true; rescue; false; end
+
+    FileInfo = Struct.new(
+      :node_id2node,
+      :definition_table,
+      :caller_table,
+      :created_iseqs,
+    )
 
     class << self
       def compile(file)
@@ -20,17 +31,62 @@ module TypeProf
         opt[:operands_unification] = false
         opt[:coverage_enabled] = false
 
+        parse_opts = {}
+        parse_opts[:keep_script_lines] = true if RICH_AST
+
         if str
+          node = RubyVM::AbstractSyntaxTree.parse(str, **parse_opts)
           iseq = RubyVM::InstructionSequence.compile(str, path, **opt)
         else
+          node = RubyVM::AbstractSyntaxTree.parse_file(path, **parse_opts)
           iseq = RubyVM::InstructionSequence.compile_file(path, **opt)
         end
 
-        return new(iseq.to_a)
+        node_id2node = {}
+        build_ast_node_id_table(node, node_id2node) if RICH_AST
+
+        file_info = FileInfo.new(node_id2node, CodeRangeTable.new, CodeRangeTable.new, [])
+        iseq_rb = new(iseq.to_a, file_info)
+        iseq_rb.collect_local_variable_info(file_info) if RICH_AST
+        file_info.created_iseqs.each do |iseq|
+          iseq.unify_instructions
+        end
+
+        return iseq_rb, file_info.definition_table, file_info.caller_table
+      end
+
+      private def build_ast_node_id_table(node, tbl = {})
+        tbl[node.node_id] = node
+        node.children.each do |child|
+          build_ast_node_id_table(child, tbl) if child.is_a?(RubyVM::AbstractSyntaxTree::Node)
+        end
+        tbl
+      end
+
+      def code_range_from_node(node)
+        CodeRange.new(
+          CodeLocation.new(node.first_lineno, node.first_column),
+          CodeLocation.new(node.last_lineno, node.last_column),
+        )
+      end
+
+      def find_node_by_id(node, id)
+        node = RubyVM::AbstractSyntaxTree.parse(node) if node.is_a?(String)
+
+        return node if id == node.node_id
+
+        node.children.each do |child|
+          if child.is_a?(RubyVM::AbstractSyntaxTree::Node)
+            ret = find_node_by_id(child, id)
+            return ret if ret
+          end
+        end
+
+        nil
       end
     end
 
-    Insn = Struct.new(:insn, :operands, :lineno)
+    Insn = Struct.new(:insn, :operands, :lineno, :code_range, :definitions)
     class Insn
       def check?(insn_cmp, operands_cmp = nil)
         return insn == insn_cmp && (!operands_cmp || operands == operands_cmp)
@@ -39,14 +95,19 @@ module TypeProf
 
     ISEQ_FRESH_ID = [0]
 
-    def initialize(iseq)
+    def initialize(iseq, file_info)
+      file_info.created_iseqs << self
+
       @id = (ISEQ_FRESH_ID[0] += 1)
 
-      _magic, _major_version, _minor_version, _format_type, _misc,
+      _magic, _major_version, _minor_version, _format_type, misc,
         @name, @path, @absolute_path, @start_lineno, @type,
         @locals, @fargs_format, catch_table, insns = *iseq
 
-      convert_insns(insns)
+      fl, fc, ll, lc = misc[:code_location]
+      @iseq_code_range = CodeRange.new(CodeLocation.new(fl, fc), CodeLocation.new(ll, lc))
+
+      convert_insns(insns, misc[:node_ids] || [], file_info)
 
       add_body_start_marker(insns)
 
@@ -54,13 +115,13 @@ module TypeProf
 
       labels = create_label_table(insns)
 
-      @insns = setup_insns(insns, labels)
+      @insns = setup_insns(insns, labels, file_info)
 
       @fargs_format[:opt] = @fargs_format[:opt].map {|l| labels[l] } if @fargs_format[:opt]
 
       @catch_table = []
       catch_table.map do |type, iseq, first, last, cont, stack_depth|
-        iseq = iseq ? ISeq.new(iseq) : nil
+        iseq = iseq ? ISeq.new(iseq, file_info) : nil
         target = labels[cont]
         entry = [type, iseq, target, stack_depth]
         labels[first].upto(labels[last]) do |i|
@@ -69,17 +130,78 @@ module TypeProf
         end
       end
 
+      def_node_id = misc[:def_node_id]
+      if def_node_id && file_info.node_id2node[def_node_id] && (@type == :method || @type == :block)
+        def_node = file_info.node_id2node[def_node_id]
+        method_name_token_range = extract_method_name_token_range(def_node)
+        if method_name_token_range
+          @callers = Utils::MutableSet.new
+          file_info.caller_table[method_name_token_range] = @callers
+        end
+      end
+
       rename_insn_types
+    end
 
-      unify_instructions
+    def extract_method_name_token_range(node)
+      case @type
+      when :method
+        regex = if node.type == :DEFS
+           /^def\s+(?:\w+)\s*\.\s*(\w+)/
+        else
+          /^def\s+(\w+)/
+        end
+        return nil unless node.source =~ regex
+        zero_loc = CodeLocation.new(1, 0)
+        name_start = $~.begin(1)
+        name_length = $~.end(1) - name_start
+        name_head_loc = zero_loc.advance_cursor(name_start, node.source)
+        name_tail_loc = name_head_loc.advance_cursor(name_length, node.source)
+        return CodeRange.new(
+          CodeLocation.new(
+            node.first_lineno + (name_head_loc.lineno - 1),
+            name_head_loc.lineno == 1 ? node.first_column + name_head_loc.column : name_head_loc.column
+          ),
+          CodeLocation.new(
+            node.first_lineno + (name_tail_loc.lineno - 1),
+            name_tail_loc.lineno == 1 ? node.first_column + name_tail_loc.column : name_tail_loc.column
+          ),
+        )
+      when :block
+        return ISeq.code_range_from_node(node)
+      end
     end
 
     def source_location(pc)
       "#{ @path }:#{ @insns[pc].lineno }"
     end
 
+    def detailed_source_location(pc)
+      code_range = @insns[pc].code_range
+      if code_range
+        [@path, code_range]
+      else
+        [@path]
+      end
+    end
+
+    def add_called_iseq(pc, callee_iseq)
+      if callee_iseq && @insns[pc].definitions
+        @insns[pc].definitions << [callee_iseq.path, callee_iseq.iseq_code_range]
+      end
+      if callee_iseq.callers
+        callee_iseq.callers << [@path, @insns[pc].code_range]
+      end
+    end
+
+    def add_def_loc(pc, detailed_loc)
+      if detailed_loc && @insns[pc].definitions
+        @insns[pc].definitions << detailed_loc
+      end
+    end
+
     attr_reader :name, :path, :absolute_path, :start_lineno, :type, :locals, :fargs_format, :catch_table, :insns
-    attr_reader :id
+    attr_reader :id, :iseq_code_range, :callers
 
     def pretty_print(q)
       q.text "ISeq["
@@ -118,7 +240,7 @@ module TypeProf
     end
 
     # Remove lineno entry and convert instructions to Insn instances
-    def convert_insns(insns)
+    def convert_insns(insns, node_ids, file_info)
       ninsns = []
       lineno = 0
       insns.each do |e|
@@ -129,7 +251,25 @@ module TypeProf
           ninsns << e
         when Array
           insn, *operands = e
-          ninsns << Insn.new(insn, operands, lineno)
+          node_id = node_ids.shift
+          node = file_info.node_id2node[node_id]
+          if node
+            code_range = ISeq.code_range_from_node(node)
+            case insn
+            when :send, :invokesuper
+              opt, blk_iseq = operands
+              opt[:node_id] = node_id
+              if blk_iseq
+                misc = blk_iseq[4] # iseq's "misc" field
+                misc[:def_node_id] = node_id
+              end
+            when :definemethod, :definesmethod
+              iseq = operands[1]
+              misc = iseq[4] # iseq's "misc" field
+              misc[:def_node_id] = node_id
+            end
+          end
+          ninsns << Insn.new(insn, operands, lineno, code_range, nil)
         else
           raise "unknown iseq entry: #{ e }"
         end
@@ -156,7 +296,7 @@ module TypeProf
           i = insns.index(label) + 1
         end
 
-        insns.insert(i, Insn.new(:_iseq_body_start, [], @start_lineno))
+        insns.insert(i, Insn.new(:_iseq_body_start, [], @start_lineno, nil, nil))
       end
     end
 
@@ -198,7 +338,7 @@ module TypeProf
       labels
     end
 
-    def setup_insns(insns, labels)
+    def setup_insns(insns, labels, file_info)
       ninsns = []
       insns.each do |e|
         case e
@@ -208,7 +348,7 @@ module TypeProf
           operands = (INSN_TABLE[e.insn] || []).zip(e.operands).map do |type, operand|
             case type
             when "ISEQ"
-              operand && ISeq.new(operand)
+              operand && ISeq.new(operand, file_info)
             when "lindex_t", "rb_num_t", "VALUE", "ID", "GENTRY", "CALL_DATA"
               operand
             when "OFFSET"
@@ -221,7 +361,12 @@ module TypeProf
             end
           end
 
-          ninsns << Insn.new(e.insn, operands, e.lineno)
+          if e.code_range && should_collect_defs(e.insn)
+            definition = Utils::MutableSet.new
+            file_info.definition_table[e.code_range] = definition
+          end
+
+          ninsns << Insn.new(e.insn, operands, e.lineno, e.code_range, definition)
         else
           raise "unknown iseq entry: #{ e }"
         end
@@ -229,6 +374,63 @@ module TypeProf
       ninsns
     end
 
+    def should_collect_defs(insn_kind)
+      case insn_kind
+      when :send, :getinstancevariable, :getconstant
+        return true
+      else
+        return false
+      end
+    end
+
+    # Collect local variable use and definition info recursively
+    def collect_local_variable_info(file_info, absolute_level = 0, parent_variable_tables = {})
+      # e.g.
+      # variable_tables[abs_level][idx] = [[path, code_range]]
+      current_variables = []
+      variable_tables = parent_variable_tables.merge({
+        absolute_level => current_variables
+      })
+
+      dummy_def_range = CodeRange.new(
+        CodeLocation.new(@start_lineno, 0),
+        CodeLocation.new(@start_lineno, 1),
+      )
+      # Fill tail elements with parameters
+      (@fargs_format[:lead_num] || 0).times do |offset|
+        current_variables[VM_ENV_DATA_SIZE + @locals.length - offset - 1] ||= Utils::MutableSet.new
+        current_variables[VM_ENV_DATA_SIZE + @locals.length - offset - 1] << [@path, dummy_def_range]
+      end
+
+      @insns.each do |insn|
+        next unless insn.insn == :getlocal || insn.insn == :setlocal
+
+        idx = insn.operands[0]
+        # note: level is relative value to the current level
+        level = insn.operands[1]
+        target_abs_level = absolute_level - level
+        variable_tables[target_abs_level] ||= {}
+        variable_tables[target_abs_level][idx] ||= Utils::MutableSet.new
+
+        case insn.insn
+        when :setlocal
+          variable_tables[target_abs_level][idx] << [path, insn.code_range]
+        when :getlocal
+          file_info.definition_table[insn.code_range] = variable_tables[target_abs_level][idx]
+        end
+      end
+
+      @insns.each do |insn|
+        insn.operands.each do |operand|
+          next unless operand.is_a?(ISeq)
+          operand.collect_local_variable_info(
+            file_info, absolute_level + 1,
+            variable_tables
+          )
+        end
+      end
+    end
+
     def rename_insn_types
       @insns.each do |insn|
         case insn.insn
@@ -331,7 +533,7 @@ module TypeProf
               break unless @insns[j + 1].check?(:putobject, [true])
               break unless @insns[j + 2].check?(:getconstant) # TODO: support A::B::C
               break unless @insns[j + 3].check?(:topn, [1])
-              break unless @insns[j + 4].check?(:send, [{:mid=>:===, :flag=>20, :orig_argc=>1}, nil])
+              break unless @insns[j + 4].check?(:send) && @insns[j + 4].operands[0].slice(:mid, :flag, :orig_argc) == {:mid=>:===, :flag=>20, :orig_argc=>1}
               break unless @insns[j + 5].check?(:branch)
               target_pc = @insns[j + 5].operands[1]
               break unless @insns[target_pc].check?(:pop, [])