Add support for # @type var type assertion comments

lib/typeprof/core/ast.rb

@@ -473,6 +473,28 @@ def self.create_rbs_member(raw_decl, lenv)
       end
     end
 
+    def self.parse_type_var_comment(raw_node, lenv)
+      comments = lenv.file_context.comments
+      return nil unless comments
+      node_line = raw_node.location.start_line
+      idx = comments.bsearch_index { |c| c.location.start_line >= node_line }
+      idx = (idx || comments.size) - 1
+      return nil if idx < 0
+      comment = comments[idx]
+      return nil unless comment.location.start_line == node_line - 1
+      text = comment.location.slice
+      if text =~ /\A#\s*@type\s+var\s+(\w+)\s*:\s*(.+)\z/
+        var_name = $1.to_sym
+        type_str = $2
+        rbs_type = RBS::Parser.parse_type(type_str)
+        return nil unless rbs_type
+        rbs_type_node = AST.create_rbs_type(rbs_type, lenv)
+        [var_name, rbs_type_node]
+      end
+    rescue RBS::ParsingError
+      nil
+    end
+
     def self.create_rbs_func_type(raw_decl, raw_type_params, raw_block, lenv)
       SigFuncType.new(raw_decl, raw_type_params, raw_block, lenv)
     end

lib/typeprof/core/ast/misc.rb

@@ -12,25 +12,53 @@ def initialize(raw_node, lenv, use_result)
             DummyNilNode.new(TypeProf::CodeRange.new(last, last), lenv)
           end
         end
+        @type_var_assertions = {}
+        stmts.each_with_index do |n, i|
+          next unless n
+          result = AST.parse_type_var_comment(n, lenv)
+          if result
+            @type_var_assertions[i] = result
+          end
+        end
       end
 
       attr_reader :stmts
 
       def subnodes = { stmts: }
 
+      def define0(genv)
+        @type_var_assertions.each_value do |_var_name, rbs_type_node|
+          rbs_type_node.define(genv)
+        end
+        super(genv)
+      end
+
+      def undefine0(genv)
+        super(genv)
+        @type_var_assertions.each_value do |_var_name, rbs_type_node|
+          rbs_type_node.undefine(genv) if rbs_type_node.static_ret
+        end
+      end
+
       def install0(genv)
         ret = nil
 
         post_stmts = []
 
-        @stmts.each do |stmt|
+        @stmts.each_with_index do |stmt, i|
           next if stmt.nil?
 
           if stmt.is_a?(PostExecutionNode)
             post_stmts << stmt
             next
           end
 
+          if (assertion = @type_var_assertions[i])
+            var_name, rbs_type_node = assertion
+            box = @changes.add_type_read_box(genv, rbs_type_node)
+            @lenv.set_var(var_name, box.ret)
+          end
+
           ret = stmt.install(genv)
         end
 

scenario/misc/type_var_comment.rb

@@ -0,0 +1,11 @@
+## update: test.rb
+def foo(x)
+  a = x
+  # @type var a: Integer
+  a + 1
+end
+
+## assert
+class Object
+  def foo: (untyped) -> Integer
+end