marshal.c: Make partial objects table conditional

Only allocate and use the partial objects table when a load proc is active or loaded objects must be frozen.

Co-authored-by: Codex <noreply@openai.com>

benchmark/marshal_load_partial_objects.yml

@@ -0,0 +1,74 @@
+prelude: |
+  MarshalLoadPartialObjectsObject = Class.new do
+    def initialize(a, b, c, d)
+      @a = a
+      @b = b
+      @c = c
+      @d = d
+    end
+  end
+
+  MarshalLoadPartialObjectsStruct = Struct.new(:a, :b, :c, :d)
+  MarshalLoadPartialObjectsData = Data.define(:a, :b, :c, :d)
+
+  integer_array_10 = 10.times.to_a
+  integer_array_1000 = 1000.times.to_a
+  unique_strings = 1000.times.map { |i| "string #{i}" }
+  shared_string = +"shared string"
+  shared_strings = Array.new(1000, shared_string)
+  integer_hash = 1000.times.to_h { |i| [i, i] }
+  symbol_string_hash = 1000.times.to_h { |i| ["key#{i}".to_sym, "value #{i}"] }
+  objects = 1000.times.map { |i| MarshalLoadPartialObjectsObject.new(i, i + 1, i + 2, i + 3) }
+  structs = 1000.times.map { |i| MarshalLoadPartialObjectsStruct.new(i, i + 1, i + 2, i + 3) }
+  data_objects = 1000.times.map { |i| MarshalLoadPartialObjectsData.new(i, i + 1, i + 2, i + 3) }
+
+  integer_array_10_dump = Marshal.dump(integer_array_10)
+  integer_array_1000_dump = Marshal.dump(integer_array_1000)
+  unique_strings_dump = Marshal.dump(unique_strings)
+  shared_strings_dump = Marshal.dump(shared_strings)
+  integer_hash_dump = Marshal.dump(integer_hash)
+  symbol_string_hash_dump = Marshal.dump(symbol_string_hash)
+  objects_dump = Marshal.dump(objects)
+  structs_dump = Marshal.dump(structs)
+  data_objects_dump = Marshal.dump(data_objects)
+
+  load_proc = ->(object) { object }
+
+benchmark:
+  marshal_load_integer_array_10: 'Marshal.load(integer_array_10_dump)'
+  marshal_load_integer_array_10_with_proc: 'Marshal.load(integer_array_10_dump, load_proc)'
+  marshal_load_integer_array_10_freeze: 'Marshal.load(integer_array_10_dump, freeze: true)'
+
+  marshal_load_integer_array_1000: 'Marshal.load(integer_array_1000_dump)'
+  marshal_load_integer_array_1000_with_proc: 'Marshal.load(integer_array_1000_dump, load_proc)'
+  marshal_load_integer_array_1000_freeze: 'Marshal.load(integer_array_1000_dump, freeze: true)'
+
+  marshal_load_unique_strings: 'Marshal.load(unique_strings_dump)'
+  marshal_load_unique_strings_with_proc: 'Marshal.load(unique_strings_dump, load_proc)'
+  marshal_load_unique_strings_freeze: 'Marshal.load(unique_strings_dump, freeze: true)'
+
+  marshal_load_shared_strings: 'Marshal.load(shared_strings_dump)'
+  marshal_load_shared_strings_with_proc: 'Marshal.load(shared_strings_dump, load_proc)'
+  marshal_load_shared_strings_freeze: 'Marshal.load(shared_strings_dump, freeze: true)'
+
+  marshal_load_integer_hash: 'Marshal.load(integer_hash_dump)'
+  marshal_load_integer_hash_with_proc: 'Marshal.load(integer_hash_dump, load_proc)'
+  marshal_load_integer_hash_freeze: 'Marshal.load(integer_hash_dump, freeze: true)'
+
+  marshal_load_symbol_string_hash: 'Marshal.load(symbol_string_hash_dump)'
+  marshal_load_symbol_string_hash_with_proc: 'Marshal.load(symbol_string_hash_dump, load_proc)'
+  marshal_load_symbol_string_hash_freeze: 'Marshal.load(symbol_string_hash_dump, freeze: true)'
+
+  marshal_load_objects: 'Marshal.load(objects_dump)'
+  marshal_load_objects_with_proc: 'Marshal.load(objects_dump, load_proc)'
+  marshal_load_objects_freeze: 'Marshal.load(objects_dump, freeze: true)'
+
+  marshal_load_structs: 'Marshal.load(structs_dump)'
+  marshal_load_structs_with_proc: 'Marshal.load(structs_dump, load_proc)'
+  marshal_load_structs_freeze: 'Marshal.load(structs_dump, freeze: true)'
+
+  marshal_load_data_objects: 'Marshal.load(data_objects_dump)'
+  marshal_load_data_objects_with_proc: 'Marshal.load(data_objects_dump, load_proc)'
+  marshal_load_data_objects_freeze: 'Marshal.load(data_objects_dump, freeze: true)'
+
+loop_count: 1000

marshal.c

@@ -1275,7 +1275,7 @@ mark_load_arg(void *ptr)
         return;
     rb_mark_tbl(p->symbols);
     rb_mark_tbl(p->data);
-    rb_mark_tbl(p->partial_objects);
+    if (p->partial_objects) rb_mark_tbl(p->partial_objects);
     rb_mark_hash(p->compat_tbl);
 }
 
@@ -1658,7 +1658,9 @@ r_entry0(VALUE v, st_index_t num, struct load_arg *arg)
         st_lookup(arg->compat_tbl, v, &real_obj);
     }
     st_insert(arg->data, num, real_obj);
-    st_insert(arg->partial_objects, (st_data_t)real_obj, Qtrue);
+    if (arg->partial_objects) {
+        st_insert(arg->partial_objects, (st_data_t)real_obj, Qtrue);
+    }
     return v;
 }
 
@@ -1693,9 +1695,11 @@ r_leave(VALUE v, struct load_arg *arg, bool partial)
 {
     v = r_fixup_compat(v, arg);
     if (!partial) {
-        st_data_t data;
-        st_data_t key = (st_data_t)v;
-        st_delete(arg->partial_objects, &key, &data);
+        if (arg->partial_objects) {
+            st_data_t data;
+            st_data_t key = (st_data_t)v;
+            st_delete(arg->partial_objects, &key, &data);
+        }
         if (arg->freeze) {
             if (RB_TYPE_P(v, T_MODULE) || RB_TYPE_P(v, T_CLASS)) {
                 // noop
@@ -1890,7 +1894,8 @@ r_object_for(struct load_arg *arg, bool partial, int *ivp, VALUE klass, VALUE ex
             rb_raise(rb_eArgError, "dump format error (unlinked)");
         }
         v = (VALUE)link;
-        if (!st_lookup(arg->partial_objects, (st_data_t)v, &link)) {
+        if (arg->partial_objects &&
+            !st_lookup(arg->partial_objects, (st_data_t)v, &link)) {
             if (arg->freeze && RB_TYPE_P(v, T_STRING)) {
                 v = rb_str_to_interned_str(v);
             }
@@ -2382,8 +2387,10 @@ clear_load_arg(struct load_arg *arg)
     arg->symbols = 0;
     st_free_table(arg->data);
     arg->data = 0;
-    st_free_table(arg->partial_objects);
-    arg->partial_objects = 0;
+    if (arg->partial_objects) {
+        st_free_table(arg->partial_objects);
+        arg->partial_objects = 0;
+    }
     if (arg->compat_tbl) {
         st_free_table(arg->compat_tbl);
         arg->compat_tbl = 0;
@@ -2413,7 +2420,7 @@ rb_marshal_load_with_proc(VALUE port, VALUE proc, bool freeze)
     arg->offset = 0;
     arg->symbols = st_init_numtable();
     arg->data    = rb_init_identtable();
-    arg->partial_objects = rb_init_identtable();
+    arg->partial_objects = (RTEST(proc) || freeze) ? rb_init_identtable() : NULL;
     arg->compat_tbl = 0;
     arg->proc = 0;
     arg->readable = 0;

test/ruby/test_marshal.rb

@@ -769,6 +769,14 @@ class TestMarshal < Test::Unit::TestCase
     assert_equal object, Marshal.load(Marshal.dump(object), :freeze.to_proc)
   end
 
+  def test_marshal_false_proc
+    object = []
+    object << object
+
+    loaded = Marshal.load(Marshal.dump(object), false)
+    assert_same loaded, loaded.first
+  end
+
   def test_marshal_load_extended_class_crash
     assert_separately([], "#{<<-"begin;"}\n#{<<-"end;"}")
     begin;