Support splat for masgn

lib/typeprof/core/ast/misc.rb

@@ -60,22 +60,46 @@ def diff(prev_node)
     class MultiWriteNode < Node
       def initialize(raw_node, lenv)
         super(raw_node, lenv)
-        @rhs = AST.create_node(raw_node.value, lenv)
-        @lhss = []
-        raw_node.lefts.each do |raw_lhs|
-          lhs = AST.create_target_node(raw_lhs, lenv)
-          @lhss << lhs
+        @value = AST.create_node(raw_node.value, lenv)
+        @lefts = raw_node.lefts.map do |raw_lhs|
+          AST.create_target_node(raw_lhs, lenv)
         end
+        if raw_node.rest
+          # TODO: need more complex case handling
+          raise unless raw_node.rest.type == :splat_node
+          @rest = AST.create_target_node(raw_node.rest.expression, lenv)
+        end
+        @rights = raw_node.rights.map do |raw_lhs|
+          AST.create_target_node(raw_lhs, lenv)
+        end
+        # TODO: raw_node.rest, raw_node.rights
       end
 
-      attr_reader :rhs, :lhss
+      attr_reader :value, :lefts, :rest, :rights
 
-      def subnodes = { rhs:, lhss: }
+      def subnodes = { value:, lefts:, rest:, rights: }
 
       def install0(genv)
-        @lhss.each {|lhs| lhs.install(genv) }
-        rhs = @rhs.install(genv)
-        box = @changes.add_masgn_box(genv, rhs, @lhss.map {|lhs| lhs.rhs.ret || raise(lhs.rhs.inspect) })
+        value = @value.install(genv)
+
+        @lefts.each {|lhs| lhs.install(genv) }
+        @lefts.each {|lhs| lhs.rhs.ret || raise(lhs.rhs.inspect) }
+        lefts = @lefts.map {|lhs| lhs.rhs.ret }
+
+        if @rest
+          @rest.install(genv)
+          @rest.rhs.ret || raise(@rest.rhs.inspect)
+          rest_elem = Vertex.new(self)
+          @changes.add_edge(genv, Source.new(Type::Instance.new(genv, genv.mod_ary, [rest_elem])), @rest.rhs.ret)
+        end
+
+        if @rights
+          @rights.each {|lhs| lhs.install(genv) }
+          @rights.each {|lhs| lhs.rhs.ret || raise(lhs.rhs.inspect) }
+          rights = @rights.map {|rhs| rhs.ret }
+        end
+
+        box = @changes.add_masgn_box(genv, value, lefts, rest_elem, rights)
         box.ret
       end
 

lib/typeprof/core/env/method.rb

@@ -75,7 +75,8 @@ def initialize(node, f_args, next_boxes)
 
     def accept_args(genv, changes, caller_positionals, caller_ret, ret_check)
       if caller_positionals.size == 1 && @f_args.size >= 2
-        changes.add_masgn_box(genv, caller_positionals[0], @f_args)
+        # TODO: support splat "do |a, *b, c|"
+        changes.add_masgn_box(genv, caller_positionals[0], @f_args, nil, nil)
       else
         caller_positionals.zip(@f_args) do |a_arg, f_arg|
           changes.add_edge(genv, a_arg, f_arg) if f_arg

lib/typeprof/core/graph/box.rb

@@ -951,32 +951,39 @@ def run0(genv, changes)
   end
 
   class MAsgnBox < Box
-    def initialize(node, genv, rhs, lhss)
+    def initialize(node, genv, value, lefts, rest_elem, rights)
       super(node)
-      @rhs = rhs
-      @lhss = lhss
-      @rhs.add_edge(genv, self)
+      @value = value
+      @lefts = lefts
+      @rest_elem = rest_elem
+      @rights = rights
+      @value.add_edge(genv, self)
     end
 
-    attr_reader :node, :rhs, :lhss
+    attr_reader :node, :value, :lefts, :rest_elem, :rights
 
     def destroy(genv)
-      @rhs.remove_edge(genv, self) # TODO: Is this really needed?
+      @value.remove_edge(genv, self) # TODO: Is this really needed?
       super(genv)
     end
 
     def ret = @rhs
 
     def run0(genv, changes)
       edges = []
-      @rhs.each_type do |ty|
+      @value.each_type do |ty|
+        # TODO: call to_ary?
         case ty
         when Type::Array
-          @lhss.each_with_index do |lhs, i|
-            edges << [ty.get_elem(genv, i), lhs]
-          end
+          edges.concat(ty.splat_assign(genv, @lefts, @rest_elem, @rights))
         else
-          edges << [Source.new(ty), @lhss[0]]
+          if @lefts.size >= 1
+            edges << [Source.new(ty), @lefts[0]]
+          elsif @rights.size >= 1
+            edges << [Source.new(ty), @rights[0]]
+          else
+            edges << [Source.new(ty), @rest_elem]
+          end
         end
       end
       edges.each do |src, dst|

lib/typeprof/core/graph/change_set.rb

@@ -50,7 +50,7 @@ def new_vertex(genv, sig_type_node)
     end
 
     def add_edge(genv, src, dst)
-      raise unless src.is_a?(BasicVertex)
+      raise src.class.to_s unless src.is_a?(BasicVertex)
       src.add_edge(genv, dst) if !@edges.include?([src, dst]) && !@new_edges.include?([src, dst])
       @new_edges << [src, dst]
     end
@@ -81,10 +81,10 @@ def add_hash_splat_box(genv, arg, unified_key, unified_val)
       @new_boxes[key] = HashSplatBox.new(@node, genv, arg, unified_key, unified_val)
     end
 
-    def add_masgn_box(genv, rhs, lhss)
-      key = [:masgn, rhs, lhss]
+    def add_masgn_box(genv, value, lefts, rest_elem, rights)
+      key = [:masgn, value, lefts, rest_elem, rights]
       return if @new_boxes[key]
-      @new_boxes[key] = MAsgnBox.new(@node, genv, rhs, lhss)
+      @new_boxes[key] = MAsgnBox.new(@node, genv, value, lefts, rest_elem, rights)
     end
 
     def add_method_def_box(genv, cpath, singleton, mid, f_args, ret_boxes)

lib/typeprof/core/type.rb

@@ -75,6 +75,7 @@ def initialize(genv, mod, args)
         raise mod.class.to_s unless mod.is_a?(ModuleEntity)
         @mod = mod
         @args = args
+        raise unless @args.is_a?(::Array)
       end
 
       attr_reader :mod, :args
@@ -165,6 +166,35 @@ def get_elem(genv, idx = nil)
         end
       end
 
+      def splat_assign(genv, lefts, rest_elem, rights)
+        edges = []
+        state = :left
+        j = nil
+        @elems.each_with_index do |elem, i|
+          case state
+          when :left
+            if i < lefts.size
+              edges << [elem, lefts[i]]
+            else
+              break unless rest_elem
+              state = :rest
+              redo
+            end
+          when :rest
+            if @elems.size - i > rights.size
+              edges << [elem, rest_elem]
+            else
+              state = :right
+              j = i
+              redo
+            end
+          when :right
+            edges << [elem, rights[i - j]]
+          end
+        end
+        edges
+      end
+
       def base_type(genv)
         @base_type
       end

scenario/block/rbs_block5.rb

@@ -16,7 +16,7 @@ def accept_x_y_z
 ## assert
 class Object
   def accept_x: -> [Integer, String]?
-  def accept_x_y_z: -> [Integer, String, nil]?
+  def accept_x_y_z: -> [Integer, String, untyped]?
 end
 
 ## update: test.rbs

scenario/variable/masgn-non-array.rb

@@ -0,0 +1,21 @@
+## update: test.rb
+def check
+  *ary = 42
+  ary
+end
+
+## assert
+class Object
+  def check: -> Array[Integer]
+end
+
+## update: test.rb
+def check
+  a, *ary, z = 42
+  [a, ary, z]
+end
+
+## assert
+class Object
+  def check: -> [Integer, Array[untyped], untyped]
+end

scenario/variable/masgn3.rb

@@ -0,0 +1,21 @@
+## update: test.rb
+def check
+  *ary = [:a, :b, :c, :d]
+  ary
+end
+
+## assert
+class Object
+  def check: -> Array[:a | :b | :c | :d]
+end
+
+## update: test.rb
+def check
+  a, *ary, z = [:a, :b, :c, :d]
+  [a, ary, z]
+end
+
+## assert
+class Object
+  def check: -> [:a, Array[:b | :c], :d]
+end