From 2aa0bd1693298dc58379d6e8f7dc5adc4b107505 Mon Sep 17 00:00:00 2001
From: pvcresin <pvcresin0730@gmail.com>
Date: Mon, 28 Sep 2026 18:46:57 +0900
Subject: [PATCH] Fix local variable inference across retry

---
 lib/typeprof/core/ast/control.rb | 45 ++++++++++++++++++++++++++++++++
 scenario/control/retry.rb        | 28 ++++++++++++++++++++
 scenario/known-issues/retry.rb   | 20 --------------
 3 files changed, 73 insertions(+), 20 deletions(-)
 delete mode 100644 scenario/known-issues/retry.rb

diff --git a/lib/typeprof/core/ast/control.rb b/lib/typeprof/core/ast/control.rb
index 6c315699d..5a625f7a9 100644
--- a/lib/typeprof/core/ast/control.rb
+++ b/lib/typeprof/core/ast/control.rb
@@ -620,8 +620,13 @@ def undefine0(genv)
       def install0(genv)
         ret = Vertex.new(self)
 
+        retrying_clauses = @rescue_clauses.select {|clause| has_retry?(clause) }
+
         vars = []
         @body.modified_vars(@lenv.locals.keys, vars) if @body
+        unless retrying_clauses.empty?
+          retrying_clauses.each {|clause| clause.modified_vars(@lenv.locals.keys, vars) }
+        end
         vars.uniq!
 
         old_vtxs = {}
@@ -630,6 +635,17 @@ def install0(genv)
           old_vtxs[var] = vtx
         end
 
+        retry_vtxs = {}
+        unless retrying_clauses.empty?
+          # A retry starts the begin body again with the locals from the rescue clause.
+          vars.each do |var|
+            vtx = Vertex.new(self)
+            @changes.add_edge(genv, old_vtxs[var], vtx)
+            retry_vtxs[var] = vtx
+            @lenv.set_var(var, vtx)
+          end
+        end
+
         @changes.add_edge(genv, @body.install(genv), ret)
 
         body_vtxs = {}
@@ -639,12 +655,14 @@ def install0(genv)
 
         clause_vtxs_list = []
         @rescue_clauses.each do |clause|
+          rescue_input_vtxs = {}
           vars.each do |var|
             old_vtx = old_vtxs[var]
             nvtx = old_vtx.new_vertex(genv, self)
 
             @changes.add_edge(genv, body_vtxs[var], nvtx) unless body_vtxs[var] == old_vtxs[var]
 
+            rescue_input_vtxs[var] = nvtx
             @lenv.set_var(var, nvtx)
           end
 
@@ -654,6 +672,14 @@ def install0(genv)
           vars.each do |var|
             clause_vtxs_list.last[var] = @lenv.get_var(var)
           end
+
+          if retrying_clauses.include?(clause)
+            vars.each do |var|
+              # Avoid a redundant cycle when neither the body nor rescue changes this local.
+              next if clause_vtxs_list.last[var] == rescue_input_vtxs[var] && body_vtxs[var] == retry_vtxs[var]
+              @changes.add_edge(genv, clause_vtxs_list.last[var], retry_vtxs[var])
+            end
+          end
         end
 
         if @else_clause
@@ -703,6 +729,25 @@ def install0(genv)
 
         ret
       end
+
+      private
+
+      def has_retry?(root)
+        nodes = [root]
+        until nodes.empty?
+          node = nodes.pop
+          next if node != root && (node.is_a?(DefNode) || node.is_a?(LambdaNode))
+          if node != root && node.is_a?(BeginNode)
+            # Its body and else remain in this retry scope; its rescue clauses form a new one.
+            nodes << node.body
+            nodes << node.else_clause if node.else_clause
+            next
+          end
+          return true if node.is_a?(RetryNode)
+          node.each_subnode {|subnode| nodes << subnode }
+        end
+        false
+      end
     end
 
     class RetryNode < Node
diff --git a/scenario/control/retry.rb b/scenario/control/retry.rb
index 2bb609aba..6d3d5b2d3 100644
--- a/scenario/control/retry.rb
+++ b/scenario/control/retry.rb
@@ -19,10 +19,38 @@ def bar(n)
   :a
 end
 
+def baz
+  n = 1
+  begin
+    raise if rand < 0.5
+    n
+  rescue
+    n = "str"
+    retry
+  end
+end
+
+def qux
+  n = 1
+  begin
+    raise if rand < 0.5
+    n
+  rescue
+    n = "str"
+    begin
+      retry
+    end
+  end
+end
+
 foo(1)
+baz
+qux
 
 ## assert
 class Object
   def foo: (Integer) -> (:a | :b | :c | Float)
   def bar: (Integer | String) -> :a
+  def baz: -> (Integer | String)
+  def qux: -> (Integer | String)
 end
diff --git a/scenario/known-issues/retry.rb b/scenario/known-issues/retry.rb
deleted file mode 100644
index 034204a73..000000000
--- a/scenario/known-issues/retry.rb
+++ /dev/null
@@ -1,20 +0,0 @@
-## update
-def foo
-  n = 1
-  begin
-    raise if rand < 0.5
-    n
-  rescue
-    n = "str"
-    retry
-  end
-end
-  
-foo
-  
-## assert
-class Object
-  def foo: () -> (Integer | String)
-end
-
-## diagnostics