From e5058b58c22e19d559b1122d94af5af3931aa416 Mon Sep 17 00:00:00 2001
From: Aaron Patterson <tenderlove@ruby-lang.org>
Date: Wed, 12 Oct 2022 14:37:02 -0700
Subject: [PATCH] Only expose Ruby Shape API if VM_CHECK_MODE is enabled

---
 lib/mjit/compiler.rb     |  2 +-
 mjit_c.rb                | 25 +++++++++++++++++++++++--
 shape.c                  | 22 +++++++++++++---------
 test/ruby/test_shapes.rb |  2 +-
 tool/mjit/bindgen.rb     |  2 ++
 5 files changed, 40 insertions(+), 13 deletions(-)

diff --git a/lib/mjit/compiler.rb b/lib/mjit/compiler.rb
index 0d2910bf4e..55fcee6b87 100644
--- a/lib/mjit/compiler.rb
+++ b/lib/mjit/compiler.rb
@@ -356,7 +356,7 @@ module RubyVM::MJIT
       source_shape_id = if dest_shape_id == C.INVALID_SHAPE_ID
                           dest_shape_id
                         else
-                          RubyVM::Shape.find_by_id(dest_shape_id).parent_id
+                          C.rb_shape_get_shape_by_id(dest_shape_id).parent_id
                         end
 
       src = +''
diff --git a/mjit_c.rb b/mjit_c.rb
index 1858f86e4d..4a68ec12ae 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -6,11 +6,11 @@ module RubyVM::MJIT
 
   class << C
     def SHAPE_BITS
-      RubyVM::Shape::SHAPE_BITS
+      Primitive.cexpr! 'UINT2NUM(SHAPE_BITS)'
     end
 
     def SHAPE_FLAG_SHIFT
-      RubyVM::Shape::SHAPE_FLAG_SHIFT
+      Primitive.cexpr! 'UINT2NUM(SHAPE_FLAG_SHIFT)'
     end
 
     def ROBJECT_EMBED_LEN_MAX
@@ -29,6 +29,12 @@ module RubyVM::MJIT
       Primitive.has_cache_for_send(cc.to_i, insn)
     end
 
+    def rb_shape_get_shape_by_id(shape_id)
+      _shape_id = shape_id.to_i
+      shape_addr = Primitive.cexpr! 'PTR2NUM((VALUE)rb_shape_get_shape_by_id((shape_id_t)NUM2UINT(_shape_id)))'
+      rb_shape_t.new(shape_addr)
+    end
+
     def rb_iseq_check(iseq)
       _iseq_addr = iseq.to_i
       iseq_addr = Primitive.cexpr! 'PTR2NUM((VALUE)rb_iseq_check((rb_iseq_t *)NUM2PTR(_iseq_addr)))'
@@ -595,6 +601,21 @@ module RubyVM::MJIT
     @rb_serial_t ||= CType::Immediate.parse("unsigned long long")
   end
 
+  def C.rb_shape
+    @rb_shape ||= CType::Struct.new(
+      "rb_shape", Primitive.cexpr!("SIZEOF(struct rb_shape)"),
+      edges: [CType::Pointer.new { self.rb_id_table }, Primitive.cexpr!("OFFSETOF((*((struct rb_shape *)NULL)), edges)")],
+      edge_name: [self.ID, Primitive.cexpr!("OFFSETOF((*((struct rb_shape *)NULL)), edge_name)")],
+      iv_count: [self.attr_index_t, Primitive.cexpr!("OFFSETOF((*((struct rb_shape *)NULL)), iv_count)")],
+      type: [CType::Immediate.parse("uint8_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_shape *)NULL)), type)")],
+      parent_id: [self.shape_id_t, Primitive.cexpr!("OFFSETOF((*((struct rb_shape *)NULL)), parent_id)")],
+    )
+  end
+
+  def C.rb_shape_t
+    @rb_shape_t ||= self.rb_shape
+  end
+
   def C.VALUE
     @VALUE ||= CType::Immediate.find(Primitive.cexpr!("SIZEOF(VALUE)"), Primitive.cexpr!("SIGNED_TYPE_P(VALUE)"))
   end
diff --git a/shape.c b/shape.c
index f21b9a4ece..1b0f1a5dc9 100644
--- a/shape.c
+++ b/shape.c
@@ -306,6 +306,13 @@ rb_shape_set_shape(VALUE obj, rb_shape_t* shape)
     rb_shape_set_shape_id(obj, rb_shape_id(shape));
 }
 
+VALUE
+rb_shape_flags_mask(void)
+{
+    return SHAPE_FLAG_MASK;
+}
+
+#if VM_CHECK_MODE > 0
 VALUE rb_cShape;
 
 /*
@@ -440,19 +447,19 @@ rb_shape_parent(VALUE self)
     }
 }
 
-VALUE
+static VALUE
 rb_shape_debug_shape(VALUE self, VALUE obj)
 {
     return rb_shape_t_to_rb_cShape(rb_shape_get_shape(obj));
 }
 
-VALUE
+static VALUE
 rb_shape_root_shape(VALUE self)
 {
     return rb_shape_t_to_rb_cShape(rb_shape_get_root_shape());
 }
 
-VALUE
+static VALUE
 rb_shape_frozen_root_shape(VALUE self)
 {
     return rb_shape_t_to_rb_cShape(rb_shape_get_frozen_root_shape());
@@ -505,12 +512,6 @@ next_shape_id(VALUE self)
     return INT2NUM(GET_VM()->next_shape_id);
 }
 
-VALUE
-rb_shape_flags_mask(void)
-{
-    return SHAPE_FLAG_MASK;
-}
-
 static VALUE
 rb_shape_find_by_id(VALUE mod, VALUE id)
 {
@@ -520,10 +521,12 @@ rb_shape_find_by_id(VALUE mod, VALUE id)
     }
     return rb_shape_t_to_rb_cShape(rb_shape_get_shape_by_id(shape_id));
 }
+#endif
 
 void
 Init_shape(void)
 {
+#if VM_CHECK_MODE > 0
     rb_cShape = rb_define_class_under(rb_cRubyVM, "Shape", rb_cObject);
     rb_undef_alloc_func(rb_cShape);
 
@@ -548,4 +551,5 @@ Init_shape(void)
     rb_define_singleton_method(rb_cShape, "of", rb_shape_debug_shape, 1);
     rb_define_singleton_method(rb_cShape, "root_shape", rb_shape_root_shape, 0);
     rb_define_singleton_method(rb_cShape, "frozen_root_shape", rb_shape_frozen_root_shape, 0);
+#endif
 }
diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
index 7142c30cd5..0da296189d 100644
--- a/test/ruby/test_shapes.rb
+++ b/test/ruby/test_shapes.rb
@@ -179,4 +179,4 @@ class TestShapes < Test::Unit::TestCase
       RubyVM::Shape.find_by_id(-1)
     end
   end
-end
+end if defined?(RubyVM::Shape)
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 8c21d42449..77b81814e3 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -380,6 +380,8 @@ generator = BindingGenerator.new(
     rb_mjit_compile_info
     rb_mjit_unit
     rb_serial_t
+    rb_shape
+    rb_shape_t
   ],
   dynamic_types: %w[
     VALUE