Support Marshal.{dump,load} for core Set

This was missed when adding core Set, because it's handled
implicitly for T_OBJECT.

Keep marshal compatibility between core Set and stdlib Set,
so you can unmarshal core Set with stdlib Set and vice versa.

Co-authored-by: Nobuyoshi Nakada <[email protected]>

Author: jeremyevans

hash.c

@@ -2253,7 +2253,7 @@ rb_hash_default(int argc, VALUE *argv, VALUE hash)
  *  See {Hash Default}[rdoc-ref:Hash@Hash+Default].
  */
 
-static VALUE
+VALUE
 rb_hash_set_default(VALUE hash, VALUE ifnone)
 {
     rb_hash_modify_check(hash);

internal/hash.h

@@ -72,6 +72,7 @@ struct RHash {
 /* hash.c */
 void rb_hash_st_table_set(VALUE hash, st_table *st);
 VALUE rb_hash_default_value(VALUE hash, VALUE key);
+VALUE rb_hash_set_default(VALUE hash, VALUE ifnone);
 VALUE rb_hash_set_default_proc(VALUE hash, VALUE proc);
 long rb_dbl_long_hash(double d);
 st_table *rb_init_identtable(void);

set.c

@@ -99,6 +99,7 @@ VALUE rb_cSet;
 static ID id_each_entry;
 static ID id_any_p;
 static ID id_new;
+static ID id_i_hash;
 static ID id_set_iter_lev;
 
 #define RSET_INITIALIZED FL_USER1
@@ -1850,6 +1851,66 @@ set_i_hash(VALUE set)
     return ST2FIX(hval);
 }
 
+/* :nodoc: */
+static int
+set_to_hash_i(st_data_t key, st_data_t arg)
+{
+    rb_hash_aset((VALUE)arg, (VALUE)key, Qtrue);
+    return ST_CONTINUE;
+}
+
+static VALUE
+set_i_to_h(VALUE set)
+{
+    st_index_t size = RSET_SIZE(set);
+    VALUE hash;
+    if (RSET_COMPARE_BY_IDENTITY(set)) {
+        hash = rb_ident_hash_new_with_size(size);
+    }
+    else {
+        hash = rb_hash_new_with_size(size);
+    }
+    rb_hash_set_default(hash, Qfalse);
+
+    if (size == 0) return hash;
+
+    set_iter(set, set_to_hash_i, (st_data_t)hash);
+    return hash;
+}
+
+static VALUE
+compat_dumper(VALUE set)
+{
+    VALUE dumper = rb_class_new_instance(0, 0, rb_cObject);
+    rb_ivar_set(dumper, id_i_hash, set_i_to_h(set));
+    return dumper;
+}
+
+static int
+set_i_from_hash_i(st_data_t key, st_data_t val, st_data_t set)
+{
+    if ((VALUE)val != Qtrue) {
+        rb_raise(rb_eRuntimeError, "expect true as Set value: %"PRIsVALUE, rb_obj_class((VALUE)val));
+    }
+    set_i_add((VALUE)set, (VALUE)key);
+    return ST_CONTINUE;
+}
+
+static VALUE
+set_i_from_hash(VALUE set, VALUE hash)
+{
+    Check_Type(hash, T_HASH);
+    if (rb_hash_compare_by_id_p(hash)) set_i_compare_by_identity(set);
+    rb_hash_stlike_foreach(hash, set_i_from_hash_i, (st_data_t)set);
+    return set;
+}
+
+static VALUE
+compat_loader(VALUE self, VALUE a)
+{
+    return set_i_from_hash(self, rb_ivar_get(a, id_i_hash));
+}
+
 /*
  *  Document-class: Set
  *
@@ -2068,6 +2129,7 @@ Init_Set(void)
     id_each_entry = rb_intern_const("each_entry");
     id_any_p = rb_intern_const("any?");
     id_new = rb_intern_const("new");
+    id_i_hash = rb_intern_const("@hash");
     id_set_iter_lev = rb_make_internal_id();
 
     rb_define_alloc_func(rb_cSet, set_s_alloc);
@@ -2132,7 +2194,12 @@ Init_Set(void)
     rb_define_method(rb_cSet, "superset?", set_i_superset, 1);
     rb_define_alias(rb_cSet, ">=", "superset?");
     rb_define_method(rb_cSet, "to_a", set_i_to_a, 0);
+    rb_define_method(rb_cSet, "to_h", set_i_to_h, 0);
     rb_define_method(rb_cSet, "to_set", set_i_to_set, -1);
 
+    /* :nodoc: */
+    VALUE compat = rb_define_class_under(rb_cSet, "compatible", rb_cObject);
+    rb_marshal_define_compat(rb_cSet, compat, compat_dumper, compat_loader);
+
     rb_provide("set.rb");
 }

test/ruby/test_set.rb

@@ -6,6 +6,32 @@ class TC_Set < Test::Unit::TestCase
   class Set2 < Set
   end
 
+  def test_marshal
+    set = Set[1, 2, 3]
+    mset = Marshal.load(Marshal.dump(set))
+    assert_equal(set, mset)
+    assert_equal(set.compare_by_identity?, mset.compare_by_identity?)
+
+    set.compare_by_identity
+    mset = Marshal.load(Marshal.dump(set))
+    assert_equal(set, mset)
+    assert_equal(set.compare_by_identity?, mset.compare_by_identity?)
+
+    set.instance_variable_set(:@a, 1)
+    mset = Marshal.load(Marshal.dump(set))
+    assert_equal(set, mset)
+    assert_equal(set.compare_by_identity?, mset.compare_by_identity?)
+    assert_equal(1, mset.instance_variable_get(:@a))
+
+    old_stdlib_set_data = "\x04\bo:\bSet\x06:\n@hash}\bi\x06Ti\aTi\bTF".b
+    set = Marshal.load(old_stdlib_set_data)
+    assert_equal(Set[1, 2, 3], set)
+
+    old_stdlib_set_cbi_data = "\x04\bo:\bSet\x06:\n@hashC:\tHash}\ai\x06Ti\aTF".b
+    set = Marshal.load(old_stdlib_set_cbi_data)
+    assert_equal(Set[1, 2].compare_by_identity, set)
+  end
+
   def test_aref
     assert_nothing_raised {
       Set[]