From 3632c0d26dfc3758cf4a5dfa2fcde548d83d316b Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Fri, 3 Apr 2026 12:10:09 +0900
Subject: [PATCH] Add support for # @type var type assertion comments

Support Steep-compatible `# @type var name: Type` comments that override
local variable types at statement boundaries. This enables users to
provide type hints where TypeProf's inference is imprecise, such as
narrowing variables before dynamic dispatch (e.g., send).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 lib/typeprof/core/ast.rb          | 22 ++++++++++++++++++++++
 lib/typeprof/core/ast/misc.rb     | 30 +++++++++++++++++++++++++++++-
 scenario/misc/type_var_comment.rb | 11 +++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 scenario/misc/type_var_comment.rb

diff --git a/lib/typeprof/core/ast.rb b/lib/typeprof/core/ast.rb
index 54f05878..ee91c0bc 100644
--- a/lib/typeprof/core/ast.rb
+++ b/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
diff --git a/lib/typeprof/core/ast/misc.rb b/lib/typeprof/core/ast/misc.rb
index 52db028d..c946740e 100644
--- a/lib/typeprof/core/ast/misc.rb
+++ b/lib/typeprof/core/ast/misc.rb
@@ -12,18 +12,40 @@ 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)
@@ -31,6 +53,12 @@ def install0(genv)
             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
 
diff --git a/scenario/misc/type_var_comment.rb b/scenario/misc/type_var_comment.rb
new file mode 100644
index 00000000..fb18b963
--- /dev/null
+++ b/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