Improved ensure clause handling

lib/typeprof/core/ast/control.rb

@@ -488,17 +488,20 @@ def install0(genv)
         end
 
         if @ensure_clause
-          ensure_old_vtxs = {}
           vars.each do |var|
-            ensure_old_vtxs[var] = @lenv.get_var(var)
+            union_vtx = old_vtxs[var].new_vertex(genv, self)
+            @changes.add_edge(genv, body_vtxs[var], union_vtx)
+            clause_vtxs_list.each do |clause_vtx|
+              @changes.add_edge(genv, clause_vtx[var], union_vtx)
+            end
+            @lenv.set_var(var, union_vtx)
           end
 
           @ensure_clause.install(genv)
 
           clause_vtxs_list << {}
           vars.each do |var|
             clause_vtxs_list.last[var] = @lenv.get_var(var)
-            @lenv.set_var(var, ensure_old_vtxs[var])
           end
         end
 

scenario/control/begin-variable-tracking.rb

@@ -12,19 +12,22 @@ def foo
     check_else(x)
     x = :e
   ensure
+    check_ensure(x)
     x = :f
   end
   check_after(x)
 end
 
 def check_rescue(n) = nil
 def check_else(n) = nil
+def check_ensure(n) = nil
 def check_after(n) = nil
 
 ## assert
 class Object
   def foo: -> nil
   def check_rescue: (:a | :c) -> nil
   def check_else: (:c) -> nil
+  def check_ensure: (:a | :c | :d | :e) -> nil
   def check_after: (:a | :c | :d | :e | :f) -> nil
 end

scenario/control/ensure.rb

@@ -6,14 +6,13 @@ def foo(n)
   n = "str"
   1.0
 ensure
-  ## TODO: bar should accept "Integer | String" ???
   bar(n)
 end
 
 foo(1)
 
 ## assert
 class Object
-  def bar: (String) -> nil
+  def bar: (Integer | String) -> nil
   def foo: (Integer) -> Float
 end