From b386bd0c56909151c9ead118c00fed22bd2b24b7 Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Thu, 22 Aug 2024 16:45:10 +0900
Subject: [PATCH] Support multi_write_node: /(?<a>x)/ =~ "x"; a

---
 lib/typeprof/core/ast.rb      |  1 +
 lib/typeprof/core/ast/misc.rb | 23 +++++++++++++++++++++++
 lib/typeprof/core/env.rb      |  3 ++-
 scenario/misc/match_asgn.rb   | 13 +++++++++++++
 4 files changed, 39 insertions(+), 1 deletion(-)
 create mode 100644 scenario/misc/match_asgn.rb

diff --git a/lib/typeprof/core/ast.rb b/lib/typeprof/core/ast.rb
index 993c38dc..da9c3130 100644
--- a/lib/typeprof/core/ast.rb
+++ b/lib/typeprof/core/ast.rb
@@ -183,6 +183,7 @@ def self.create_node(raw_node, lenv, use_result = true)
         rhs = AndNode.new(raw_node, read, raw_node.value, lenv)
         CallWriteNode.new(raw_node, rhs, lenv)
       when :multi_write_node then MultiWriteNode.new(raw_node, lenv)
+      when :match_write_node then MatchWriteNode.new(raw_node, lenv)
 
       # value
       when :self_node then SelfNode.new(raw_node, lenv)
diff --git a/lib/typeprof/core/ast/misc.rb b/lib/typeprof/core/ast/misc.rb
index b69a2f2b..9095d61f 100644
--- a/lib/typeprof/core/ast/misc.rb
+++ b/lib/typeprof/core/ast/misc.rb
@@ -120,6 +120,29 @@ def retrieve_at(pos, &blk)
       end
     end
 
+    class MatchWriteNode < Node
+      def initialize(raw_node, lenv)
+        super(raw_node, lenv)
+        @call = AST.create_node(raw_node.call, lenv)
+        @targets = raw_node.targets.map do |raw_lhs|
+          AST.create_target_node(raw_lhs, lenv)
+        end
+      end
+
+      attr_reader :call, :targets
+      def subnodes = { call:, targets: }
+
+      def install0(genv)
+        ret = @call.install(genv)
+        @targets.each do |target|
+          target.install(genv)
+          target.rhs.ret || raise(target.rhs.inspect)
+          @changes.add_edge(genv, Source.new(Type::Instance.new(genv, genv.mod_str, [])), target.rhs.ret)
+        end
+        ret
+      end
+    end
+
     class DefinedNode < Node
       def initialize(raw_node, lenv)
         super(raw_node, lenv)
diff --git a/lib/typeprof/core/env.rb b/lib/typeprof/core/env.rb
index 52d9bc4e..1da21dd3 100644
--- a/lib/typeprof/core/env.rb
+++ b/lib/typeprof/core/env.rb
@@ -19,6 +19,7 @@ def initialize
       @mod_ary = resolve_cpath([:Array])
       @mod_hash = resolve_cpath([:Hash])
       @mod_range = resolve_cpath([:Range])
+      @mod_str = resolve_cpath([:String])
 
       @cls_type = Type::Instance.new(self, @mod_class, [])
       @mod_type = Type::Instance.new(self, @mod_module, [])
@@ -40,7 +41,7 @@ def initialize
 
     attr_reader :type_table
 
-    attr_reader :mod_object, :mod_ary, :mod_hash, :mod_range
+    attr_reader :mod_object, :mod_ary, :mod_hash, :mod_range, :mod_str
     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/scenario/misc/match_asgn.rb b/scenario/misc/match_asgn.rb
new file mode 100644
index 00000000..a1f729f5
--- /dev/null
+++ b/scenario/misc/match_asgn.rb
@@ -0,0 +1,13 @@
+## update
+def check
+  if /(?<a>foo)/ =~ "foo"
+    a
+  else
+    1
+  end
+end
+
+## assert
+class Object
+  def check: -> (Integer | String)
+end