Lazily convert Prism/RBS locations to CodeRange objects

lib/typeprof/core/ast/call.rb

@@ -1,12 +1,12 @@
 module TypeProf::Core
   class AST
     class CallBaseNode < Node
-      def initialize(raw_node, recv, mid, mid_code_range, raw_args, last_arg, raw_block, lenv, forwarding_arguments: false)
+      def initialize(raw_node, recv, mid, mid_code_range_loc, raw_args, last_arg, raw_block, lenv, forwarding_arguments: false)
         super(raw_node, lenv)
 
         @recv = recv
         @mid = mid
-        @mid_code_range = mid_code_range
+        @mid_code_range_loc = mid_code_range_loc
 
         # args
         @positional_args = []
@@ -95,7 +95,11 @@ def initialize(raw_node, recv, mid, mid_code_range, raw_args, last_arg, raw_bloc
         @yield = raw_node.type == :yield_node
       end
 
-      attr_reader :recv, :mid, :mid_code_range, :yield
+      attr_reader :recv, :mid, :yield
+
+      def mid_code_range
+        @mid_code_range ||= @lenv.code_range_from_node(@mid_code_range_loc) if @mid_code_range_loc
+      end
       attr_reader :positional_args, :splat_flags, :keyword_args
       attr_reader :block_tbl, :block_f_args, :block_opt_positional_defaults, :block_body, :block_pass, :anonymous_block_forwarding
       attr_reader :block_multi_targets
@@ -230,7 +234,7 @@ def block_last_stmt_code_range
       end
 
       def retrieve_at(pos, &blk)
-        yield self if @mid_code_range && @mid_code_range.include?(pos)
+        yield self if mid_code_range&.include?(pos)
         each_subnode do |subnode|
           next unless subnode
           subnode.retrieve_at(pos, &blk)
@@ -257,10 +261,9 @@ class CallNode < CallBaseNode
       def initialize(raw_node, lenv)
         recv = raw_node.receiver ? AST.create_node(raw_node.receiver, lenv) : nil
         mid = raw_node.name
-        mid_code_range = lenv.code_range_from_node(raw_node.message_loc) if raw_node.message_loc
         raw_args = raw_node.arguments
         raw_block = raw_node.block
-        super(raw_node, recv, mid, mid_code_range, raw_args, nil, raw_block, lenv)
+        super(raw_node, recv, mid, raw_node.message_loc, raw_args, nil, raw_block, lenv)
       end
 
       def narrowings
@@ -331,9 +334,8 @@ def initialize(raw_node, lenv)
     class OperatorNode < CallBaseNode
       def initialize(raw_node, recv, lenv)
         mid = raw_node.binary_operator
-        mid_code_range = lenv.code_range_from_node(raw_node.binary_operator_loc)
         last_arg = AST.create_node(raw_node.value, lenv)
-        super(raw_node, recv, mid, mid_code_range, nil, last_arg, nil, lenv)
+        super(raw_node, recv, mid, raw_node.binary_operator_loc, nil, last_arg, nil, lenv)
       end
     end
 
@@ -364,18 +366,16 @@ class CallReadNode < CallBaseNode
       def initialize(raw_node, lenv)
         recv = AST.create_node(raw_node.receiver, lenv)
         mid = raw_node.read_name
-        mid_code_range = lenv.code_range_from_node(raw_node.message_loc)
-        super(raw_node, recv, mid, mid_code_range, nil, nil, nil, lenv)
+        super(raw_node, recv, mid, raw_node.message_loc, nil, nil, nil, lenv)
       end
     end
 
     class CallWriteNode < CallBaseNode
       def initialize(raw_node, rhs, lenv)
         recv = AST.create_node(raw_node.receiver, lenv)
         mid = raw_node.is_a?(Prism::CallTargetNode) ? raw_node.name : raw_node.write_name
-        mid_code_range = lenv.code_range_from_node(raw_node.message_loc)
         @rhs = rhs
-        super(raw_node, recv, mid, mid_code_range, nil, rhs, nil, lenv)
+        super(raw_node, recv, mid, raw_node.message_loc, nil, rhs, nil, lenv)
       end
 
       attr_reader :rhs

lib/typeprof/core/ast/const.rb

@@ -8,7 +8,7 @@ def initialize(raw_node, lenv)
           @cbase = nil
           @toplevel = false
           @cname = raw_node.name
-          @cname_code_range = lenv.code_range_from_node(raw_node.location)
+          @cname_code_range_loc = raw_node.location
         when :constant_path_node, :constant_path_target_node
           if raw_node.parent
             @cbase = AST.create_node(raw_node.parent, lenv)
@@ -18,14 +18,18 @@ def initialize(raw_node, lenv)
             @toplevel = true
           end
           @cname = raw_node.name
-          @cname_code_range = lenv.code_range_from_node(raw_node.name_loc)
+          @cname_code_range_loc = raw_node.name_loc
         else
           raise raw_node.type.to_s
         end
         @strict_const_scope = lenv.strict_const_scope
       end
 
-      attr_reader :cname, :cbase, :toplevel, :cname_code_range, :strict_const_scope
+      attr_reader :cname, :cbase, :toplevel, :strict_const_scope
+
+      def cname_code_range
+        @cname_code_range ||= @lenv.code_range_from_node(@cname_code_range_loc) if @cname_code_range_loc
+      end
 
       def attrs = { cname:, toplevel:, strict_const_scope: }
       def subnodes = { cbase: }
@@ -57,24 +61,28 @@ def initialize(raw_node, rhs, lenv)
           # C = expr
           @cpath = nil
           @static_cpath = lenv.cref.cpath + [raw_node.name]
-          @cname_code_range = lenv.code_range_from_node(raw_node.respond_to?(:name_loc) ? raw_node.name_loc : raw_node)
+          @cname_code_range_loc = raw_node.respond_to?(:name_loc) ? raw_node.name_loc : raw_node.location
         when :constant_path_write_node, :constant_path_operator_write_node, :constant_path_or_write_node, :constant_path_and_write_node
           # expr::C = expr
           @cpath = AST.create_node(raw_node.target, lenv)
           @static_cpath = AST.parse_cpath(raw_node.target, lenv.cref)
-          @cname_code_range = lenv.code_range_from_node(raw_node.target)
+          @cname_code_range_loc = raw_node.target.location
         when :constant_path_target_node
           # expr::C, * = ary
           @cpath = ConstantReadNode.new(raw_node, lenv)
           @static_cpath = AST.parse_cpath(raw_node, lenv.cref)
-          @cname_code_range = lenv.code_range_from_node(raw_node)
+          @cname_code_range_loc = raw_node.location
         else
           raise
         end
         @rhs = rhs
       end
 
-      attr_reader :cpath, :rhs, :static_cpath, :cname_code_range
+      attr_reader :cpath, :rhs, :static_cpath
+
+      def cname_code_range
+        @cname_code_range ||= @lenv.code_range_from_node(@cname_code_range_loc) if @cname_code_range_loc
+      end
 
       def subnodes = { cpath:, rhs: }
       def attrs = { static_cpath: }

lib/typeprof/core/ast/method.rb

@@ -135,7 +135,7 @@ def initialize(raw_node, lenv, use_result)
         # TODO: warn "def self.foo" in a metaclass
         singleton = !!raw_node.receiver || lenv.cref.scope_level == :metaclass
         mid = raw_node.name
-        mid_code_range = lenv.code_range_from_node(raw_node.name_loc)
+        mid_code_range_loc = raw_node.name_loc
         @tbl = raw_node.locals
         raw_args = raw_node.parameters
         raw_body = raw_node.body
@@ -144,7 +144,7 @@ def initialize(raw_node, lenv, use_result)
 
         @singleton = singleton
         @mid = mid
-        @mid_code_range = mid_code_range
+        @mid_code_range_loc = mid_code_range_loc
 
         ncref = CRef.new(lenv.cref.cpath, @singleton ? :class : :instance, @mid, lenv.cref)
         nlenv = LocalEnv.new(@lenv.file_context, ncref, {}, [])
@@ -177,7 +177,11 @@ def initialize(raw_node, lenv, use_result)
         @reusable = !use_result
       end
 
-      attr_reader :singleton, :mid, :mid_code_range
+      attr_reader :singleton, :mid
+
+      def mid_code_range
+        @mid_code_range ||= @lenv.code_range_from_node(@mid_code_range_loc) if @mid_code_range_loc
+      end
       attr_reader :tbl
       attr_reader :req_positionals
       attr_reader :opt_positionals
@@ -203,7 +207,6 @@ def subnodes = {
       def attrs = {
         singleton:,
         mid:,
-        mid_code_range:,
         tbl:,
         req_positionals:,
         opt_positionals:,
@@ -217,7 +220,7 @@ def attrs = {
         reusable:,
       }
 
-      def mname_code_range(_name) = @mid_code_range
+      def mname_code_range(_name) = mid_code_range
 
       def define(genv) # NOT define0
         return define_copy(genv) if @prev_node && @reusable

lib/typeprof/core/ast/module.rb

@@ -18,11 +18,15 @@ def initialize(raw_node, lenv, raw_cpath, meta, raw_scope, use_result)
           @body = nil
         end
 
-        @cname_code_range = meta ? nil : lenv.code_range_from_node(raw_node.constant_path)
+        @cname_code_range_loc = meta ? nil : raw_node.constant_path
         @mod_cdef = nil
       end
 
-      attr_reader :tbl, :cpath, :static_cpath, :cname_code_range, :body
+      attr_reader :tbl, :cpath, :static_cpath, :body
+
+      def cname_code_range
+        @cname_code_range ||= @lenv.code_range_from_node(@cname_code_range_loc) if @cname_code_range_loc
+      end
 
       def subnodes = { cpath:, body: }
       def attrs = { static_cpath:, tbl: }

lib/typeprof/core/ast/sig_decl.rb

@@ -190,7 +190,7 @@ class SigDefNode < Node
       def initialize(raw_decl, lenv)
         super(raw_decl, lenv)
         @mid = raw_decl.name
-        @mid_code_range = lenv.code_range_from_node(raw_decl.location[:name])
+        @mid_code_range_loc = raw_decl.location[:name]
         @singleton = raw_decl.singleton?
         @instance = raw_decl.instance?
         @method_types = OverloadSet.new(raw_decl.overloads.map do |overload|
@@ -200,12 +200,16 @@ def initialize(raw_decl, lenv)
         @overloading = raw_decl.overloading
       end
 
-      attr_reader :mid, :singleton, :instance, :method_types, :overloading, :mid_code_range
+      attr_reader :mid, :singleton, :instance, :method_types, :overloading
+
+      def mid_code_range
+        @mid_code_range ||= @lenv.code_range_from_node(@mid_code_range_loc) if @mid_code_range_loc
+      end
 
       def subnodes = { method_types: @method_types.to_a }
-      def attrs = { mid:, mid_code_range:, singleton:, instance:, overloading: }
+      def attrs = { mid:, singleton:, instance:, overloading: }
 
-      def mname_code_range(_name) = @mid_code_range
+      def mname_code_range(_name) = mid_code_range
 
       def install0(genv)
         [[@singleton, true], [@instance, false]].each do |enabled, singleton|