From 6a4e096643059c438db9f2d12cfb6502f155ab76 Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Wed, 7 Aug 2024 17:25:27 +0900
Subject: [PATCH] Tentative support of hash splat `{ **expr }`

---
 lib/typeprof/core/ast/base.rb         |  1 +
 lib/typeprof/core/ast/value.rb        | 22 ++++++++++++++++------
 lib/typeprof/core/env.rb              |  2 +-
 lib/typeprof/core/graph/box.rb        | 26 ++++++++++++++++++++++++++
 lib/typeprof/core/graph/change_set.rb |  6 ++++++
 scenario/hash/hash-splat.rb           | 14 ++++++++++++++
 6 files changed, 64 insertions(+), 7 deletions(-)
 create mode 100644 scenario/hash/hash-splat.rb

diff --git a/lib/typeprof/core/ast/base.rb b/lib/typeprof/core/ast/base.rb
index e8a03f9b8..d132ba7e0 100644
--- a/lib/typeprof/core/ast/base.rb
+++ b/lib/typeprof/core/ast/base.rb
@@ -128,6 +128,7 @@ def diff(prev_node)
               subnode = [subnode] if subnode.is_a?(AST::Node)
               prev_subnode = [prev_subnode] if prev_subnode.is_a?(AST::Node)
               subnode.zip(prev_subnode) do |subnode0, prev_subnode0|
+                next if subnode0 == nil && prev_subnode0 == nil
                 subnode0.diff(prev_subnode0)
                 return unless subnode0.prev_node
               end
diff --git a/lib/typeprof/core/ast/value.rb b/lib/typeprof/core/ast/value.rb
index 9d3cccd67..79642df7f 100644
--- a/lib/typeprof/core/ast/value.rb
+++ b/lib/typeprof/core/ast/value.rb
@@ -246,6 +246,7 @@ def initialize(raw_node, lenv, keywords)
         @keys = []
         @vals = []
         @keywords = keywords
+        @splat = false
 
         raw_node.elements.each do |raw_elem|
           # TODO: Support :assoc_splat_node
@@ -253,16 +254,20 @@ def initialize(raw_node, lenv, keywords)
           when :assoc_node
             @keys << AST.create_node(raw_elem.key, lenv)
             @vals << AST.create_node(raw_elem.value, lenv)
+          when :assoc_splat_node
+            @keys << nil
+            @vals << AST.create_node(raw_elem.value, lenv)
+            @splat = true
           else
-            raise "unknown hash elem"
+            raise "unknown hash elem: #{ raw_elem.type }"
           end
         end
       end
 
-      attr_reader :keys, :vals, :keywords
+      attr_reader :keys, :vals, :splat, :keywords
 
       def subnodes = { keys:, vals: }
-      def attrs = { keywords: }
+      def attrs = { splat:, keywords: }
 
       def install0(genv)
         unified_key = Vertex.new(self)
@@ -276,11 +281,16 @@ def install0(genv)
             @changes.add_edge(genv, v, unified_val)
             literal_pairs[key.lit] = v if key.is_a?(SymbolNode)
           else
-            _h = val.install(genv)
-            # TODO: if h is a hash, we need to connect its elements to the new hash
+            h = val.install(genv)
+            # TODO: do we want to call to_hash on h?
+            @changes.add_hash_splat_box(genv, h, unified_key, unified_val)
           end
         end
-        Source.new(Type::Hash.new(genv, literal_pairs, genv.gen_hash_type(unified_key, unified_val)))
+        if @splat
+          Source.new(genv.gen_hash_type(unified_key, unified_val))
+        else
+          Source.new(Type::Hash.new(genv, literal_pairs, genv.gen_hash_type(unified_key, unified_val)))
+        end
       end
     end
 
diff --git a/lib/typeprof/core/env.rb b/lib/typeprof/core/env.rb
index 4931c7c7f..2a70e5461 100644
--- a/lib/typeprof/core/env.rb
+++ b/lib/typeprof/core/env.rb
@@ -40,7 +40,7 @@ def initialize
 
     attr_reader :type_table
 
-    attr_reader :mod_object, :mod_ary, :mod_range
+    attr_reader :mod_object, :mod_ary, :mod_hash, :mod_range
     attr_reader :cls_type, :mod_type
     attr_reader :obj_type, :nil_type, :true_type, :false_type, :str_type, :int_type, :float_type
     attr_reader :proc_type, :symbol_type, :set_type, :regexp_type, :complex_type
diff --git a/lib/typeprof/core/graph/box.rb b/lib/typeprof/core/graph/box.rb
index f9a0b79ec..378b616de 100644
--- a/lib/typeprof/core/graph/box.rb
+++ b/lib/typeprof/core/graph/box.rb
@@ -306,6 +306,32 @@ def run0(genv, changes)
     end
   end
 
+  class HashSplatBox < Box
+    def initialize(node, genv, hsh, unified_key, unified_val)
+      super(node)
+      @hsh = hsh
+      @unified_key = unified_key
+      @unified_val = unified_val
+      @hsh.add_edge(genv, self)
+    end
+
+    def ret = @hsh # dummy
+
+    attr_reader :hsh, :unified_key, :unified_val
+
+    def run0(genv, changes)
+      @hsh.each_type do |ty|
+        ty = ty.base_type(genv)
+        if ty.mod == genv.mod_hash
+          changes.add_edge(genv, ty.args[0], @unified_key)
+          changes.add_edge(genv, ty.args[1], @unified_val)
+        else
+          "???"
+        end
+      end
+    end
+  end
+
   class MethodDefBox < Box
     def initialize(node, genv, cpath, singleton, mid, f_args, ret_boxes)
       super(node)
diff --git a/lib/typeprof/core/graph/change_set.rb b/lib/typeprof/core/graph/change_set.rb
index 85e8079a4..8b2a3586f 100644
--- a/lib/typeprof/core/graph/change_set.rb
+++ b/lib/typeprof/core/graph/change_set.rb
@@ -75,6 +75,12 @@ def add_splat_box(genv, arg)
       @new_boxes[key] = SplatBox.new(@node, genv, arg)
     end
 
+    def add_hash_splat_box(genv, arg, unified_key, unified_val)
+      key = [:hash_splat, arg, unified_key, unified_val]
+      return if @new_boxes[key]
+      @new_boxes[key] = HashSplatBox.new(@node, genv, arg, unified_key, unified_val)
+    end
+
     def add_masgn_box(genv, rhs, lhss)
       key = [:masgn, rhs, lhss]
       return if @new_boxes[key]
diff --git a/scenario/hash/hash-splat.rb b/scenario/hash/hash-splat.rb
new file mode 100644
index 000000000..d0312113b
--- /dev/null
+++ b/scenario/hash/hash-splat.rb
@@ -0,0 +1,14 @@
+## update
+def foo
+  { **bar, b: 1 }
+end
+
+def bar
+  { a: 1 }
+end
+
+## assert
+class Object
+  def foo: -> Hash[:a | :b, Integer]
+  def bar: -> Hash[:a, Integer]
+end