Support the assignment syntax (=> e) in rescue clauses

lib/typeprof/core/ast.rb

@@ -58,6 +58,7 @@ def self.create_node(raw_node, lenv, use_result = true)
       when :begin_node then BeginNode.new(raw_node, lenv)
       when :retry_node then RetryNode.new(raw_node, lenv)
       when :rescue_modifier_node then RescueModifierNode.new(raw_node, lenv)
+      when :rescue_node then RescueNode.new(raw_node, lenv)
 
       # constants
       when :constant_read_node, :constant_path_node

lib/typeprof/core/ast/control.rb

@@ -350,41 +350,88 @@ def install0(genv)
       end
     end
 
+    class RescueNode < Node
+      def initialize(raw_node, lenv)
+        super(raw_node, lenv)
+
+        @exceptions = raw_node.exceptions.map {|raw_cond| AST.create_node(raw_cond, lenv) }
+        @statements = AST.create_node(raw_node.statements, lenv) if raw_node.statements
+        if raw_node.reference && @statements
+          @reference = AST.create_target_node(raw_node.reference, @statements.lenv)
+        end
+      end
+
+      attr_reader :exceptions, :reference, :statements
+
+      def subnodes = { exceptions:, reference:, statements: }
+
+      def define0(genv)
+        @exceptions.each {|exc| exc.define(genv) }
+        @reference.define(genv) if @reference
+        @statements.define(genv) if @statements
+      end
+
+      def undefine0(genv)
+        @exceptions.each {|exc| exc.undefine(genv) }
+        @reference.undefine(genv) if @reference
+        @statements.undefine(genv) if @statements
+      end
+
+      def install0(genv)
+        cond_vtxs = @exceptions.map do |exc|
+          case exc
+          when AST::SplatNode
+            ary_vtx = exc.expr.install(genv)
+            @changes.add_splat_box(genv, ary_vtx).ret
+          else
+            exc.install(genv)
+          end
+        end
+
+        if @reference
+          @reference.install(genv)
+          cond_vtxs.each do |cond_vtx|
+            instance_ty_box = @changes.add_instance_type_box(genv, cond_vtx)
+            @changes.add_edge(genv, instance_ty_box.ret, @reference.rhs.ret)
+          end
+        end
+
+        if @statements
+          @statements.install(genv)
+        else
+          Source.new(genv.nil_type)
+        end
+      end
+    end
+
     class BeginNode < Node
       def initialize(raw_node, lenv)
         super(raw_node, lenv)
         @body = raw_node.statements ? AST.create_node(raw_node.statements, lenv) : DummyNilNode.new(code_range, lenv)
-        @rescue_conds = []
+
         @rescue_clauses = []
         raw_res = raw_node.rescue_clause
         while raw_res
-          raw_res.exceptions.each do |raw_cond|
-            @rescue_conds << AST.create_node(raw_cond, lenv)
-          end
-          if raw_res.statements
-            @rescue_clauses << AST.create_node(raw_res.statements, lenv)
-          end
+          @rescue_clauses << AST.create_node(raw_res, lenv)
           raw_res = raw_res.subsequent
         end
         @else_clause = raw_node.else_clause&.statements ? AST.create_node(raw_node.else_clause.statements, lenv) : DummyNilNode.new(code_range, lenv)
         @ensure_clause = raw_node.ensure_clause&.statements ? AST.create_node(raw_node.ensure_clause.statements, lenv) : DummyNilNode.new(code_range, lenv)
       end
 
-      attr_reader :body, :rescue_conds, :rescue_clauses, :else_clause, :ensure_clause
+      attr_reader :body, :rescue_clauses, :else_clause, :ensure_clause
 
-      def subnodes = { body:, rescue_conds:, rescue_clauses:, else_clause:, ensure_clause: }
+      def subnodes = { body:, rescue_clauses:, else_clause:, ensure_clause: }
 
       def define0(genv)
         @body.define(genv)
-        @rescue_conds.each {|cond| cond.define(genv) }
         @rescue_clauses.each {|clause| clause.define(genv) }
         @else_clause.define(genv) if @else_clause
         @ensure_clause.define(genv) if @ensure_clause
       end
 
       def undefine0(genv)
         @body.undefine(genv)
-        @rescue_conds.each {|cond| cond.undefine(genv) }
         @rescue_clauses.each {|clause| clause.undefine(genv) }
         @else_clause.undefine(genv) if @else_clause
         @ensure_clause.undefine(genv) if @ensure_clause
@@ -393,8 +440,7 @@ def undefine0(genv)
       def install0(genv)
         ret = Vertex.new(self)
         @changes.add_edge(genv, @body.install(genv), ret)
-        @rescue_conds.each {|cond| cond.install(genv) }
-        @rescue_clauses.each {|clause| @changes.add_edge(genv, clause.install(genv), ret) }
+        @rescue_clauses.each { |clause| @changes.add_edge(genv, clause.install(genv), ret) }
         @changes.add_edge(genv, @else_clause.install(genv), ret) if @else_clause
         @ensure_clause.install(genv) if @ensure_clause
         ret

lib/typeprof/core/graph/box.rb

@@ -1032,4 +1032,24 @@ def run0(genv, changes)
       end
     end
   end
+
+  class InstanceTypeBox < Box
+    def initialize(node, genv, singleton_ty_vtx)
+      super(node)
+      @singleton_ty_vtx = singleton_ty_vtx
+      @ret = Vertex.new(node)
+      genv.add_run(self)
+    end
+
+    attr_reader :ret
+
+    def run0(genv, changes)
+      instance_tys = []
+      @singleton_ty_vtx.each_type do |ty|
+        instance_tys << ty.get_instance_type(genv)
+      end
+      source_vtx = Source.new(*instance_tys)
+      changes.add_edge(genv, source_vtx, @ret)
+    end
+  end
 end

lib/typeprof/core/graph/change_set.rb

@@ -141,6 +141,12 @@ def add_type_read_box(genv, type)
       @new_boxes[key] = TypeReadBox.new(@node, genv, type)
     end
 
+    def add_instance_type_box(genv, singleton_ty_vtx)
+      key = [:instance_type, singleton_ty_vtx]
+      return if @new_boxes[key]
+      @new_boxes[key] = InstanceTypeBox.new(@node, genv, singleton_ty_vtx)
+    end
+
     def add_diagnostic(meth, msg)
       @new_diagnostics << TypeProf::Diagnostic.new(@node, meth, msg)
     end

scenario/control/rescue-assign.rb

@@ -0,0 +1,16 @@
+## update
+def foo(n)
+  raise if n != 0
+  n.to_s
+rescue StandardError => e
+  e.message
+end
+
+foo(1)
+
+## diagnostics
+
+## assert
+class Object
+  def foo: (Integer) -> String?
+end