From dc6be92cd0bec4227bead56845890fdd44eaa8bb Mon Sep 17 00:00:00 2001
From: Yusuke Endoh <mame@ruby-lang.org>
Date: Fri, 13 Dec 2024 16:01:02 +0900
Subject: [PATCH] Add a special handling for `obj.class`

`self.class` will return the class of its context
---
 lib/typeprof/core/builtin.rb  | 20 ++++++++++++++++----
 lib/typeprof/core/env.rb      |  2 +-
 scenario/misc/class_method.rb | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 5 deletions(-)
 create mode 100644 scenario/misc/class_method.rb

diff --git a/lib/typeprof/core/builtin.rb b/lib/typeprof/core/builtin.rb
index 848deb8f..f0aa80cf 100644
--- a/lib/typeprof/core/builtin.rb
+++ b/lib/typeprof/core/builtin.rb
@@ -5,10 +5,21 @@ def initialize(genv)
     end
 
     def class_new(changes, node, ty, a_args, ret)
-      ty = ty.get_instance_type(@genv)
-      recv = Source.new(ty)
-      changes.add_method_call_box(@genv, recv, :initialize, a_args, false)
-      changes.add_edge(@genv, Source.new(ty), ret)
+      if ty.is_a?(Type::Singleton)
+        ty = ty.get_instance_type(@genv)
+        recv = Source.new(ty)
+        changes.add_method_call_box(@genv, recv, :initialize, a_args, false)
+        changes.add_edge(@genv, Source.new(ty), ret)
+      end
+      true
+    end
+
+    def object_class(changes, node, ty, a_args, ret)
+      ty = ty.base_type(@genv)
+      mod = ty.is_a?(Type::Instance) ? ty.mod : @genv.mod_class
+      ty = Type::Singleton.new(@genv, mod)
+      vtx = Source.new(ty)
+      changes.add_edge(@genv, vtx, ret)
       true
     end
 
@@ -119,6 +130,7 @@ def hash_aset(changes, node, ty, a_args, ret)
     def deploy
       {
         class_new: [[:Class], false, :new],
+        object_class: [[:Object], false, :class],
         proc_call: [[:Proc], false, :call],
         array_aref: [[:Array], false, :[]],
         array_aset: [[:Array], false, :[]=],
diff --git a/lib/typeprof/core/env.rb b/lib/typeprof/core/env.rb
index fed95dad..d49b3a76 100644
--- a/lib/typeprof/core/env.rb
+++ b/lib/typeprof/core/env.rb
@@ -42,7 +42,7 @@ def initialize
 
     attr_reader :type_table
 
-    attr_reader :mod_object, :mod_ary, :mod_hash, :mod_range, :mod_str
+    attr_reader :mod_class, :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
     attr_reader :int_type, :float_type, :rational_type, :complex_type
diff --git a/scenario/misc/class_method.rb b/scenario/misc/class_method.rb
new file mode 100644
index 00000000..1478f345
--- /dev/null
+++ b/scenario/misc/class_method.rb
@@ -0,0 +1,34 @@
+## update: test.rb
+def self_class
+  self.class
+end
+
+def int_class
+  1.class
+end
+
+def array_class
+  [1].class
+end
+
+def hash_class
+  { 1 => "str" }.class
+end
+
+def class_class
+  Object.class
+end
+
+def unknown_class(x)
+  x.class
+end
+
+## assert
+class Object
+  def self_class: -> singleton(Object)
+  def int_class: -> singleton(Integer)
+  def array_class: -> singleton(Array)
+  def hash_class: -> singleton(Hash)
+  def class_class: -> singleton(Class)
+  def unknown_class: (untyped) -> untyped
+end