Mark weak maps as write barrier protected

For both we mark the lambda finalizer.

ObjectSpace::WeakMap doesn't mark any other reference, so we can just add the flag.

ObjectSpace::WeakKeyMap only ever add new refs in `wkmap_aset`, so we can just trigger the write barrier there.

1 files changed, 6 insertions, 4 deletions

diff --git a/weakmap.c b/weakmap.c
index 6915b3c71d..e6dbb325d7 100644
--- a/weakmap.c
+++ b/weakmap.c
@@ -74,7 +74,7 @@ static const rb_data_type_t weakmap_type = {
         wmap_memsize,
         wmap_compact,
     },
-    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
 };
 
 static VALUE wmap_finalize(RB_BLOCK_CALL_FUNC_ARGLIST(objid, self));
@@ -86,7 +86,7 @@ wmap_allocate(VALUE klass)
     VALUE obj = TypedData_Make_Struct(klass, struct weakmap, &weakmap_type, w);
     w->obj2wmap = rb_init_identtable();
     w->wmap2obj = rb_init_identtable();
-    w->final = rb_func_lambda_new(wmap_finalize, obj, 1, 1);
+    RB_OBJ_WRITE(obj, &w->final, rb_func_lambda_new(wmap_finalize, obj, 1, 1));
     return obj;
 }
 
@@ -522,7 +522,7 @@ static const rb_data_type_t weakkeymap_type = {
         wkmap_memsize,
         wkmap_compact,
     },
-    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED
 };
 
 static VALUE
@@ -559,7 +559,7 @@ wkmap_allocate(VALUE klass)
     VALUE obj = TypedData_Make_Struct(klass, struct weakkeymap, &weakkeymap_type, w);
     w->map = st_init_table(&weakkeymap_hash);
     w->obj2hash = rb_init_identtable();
-    w->final = rb_func_lambda_new(wkmap_finalize, obj, 1, 1);
+    RB_OBJ_WRITE(obj, &w->final, rb_func_lambda_new(wkmap_finalize, obj, 1, 1));
     return obj;
 }
 
@@ -653,6 +653,8 @@ wkmap_aset(VALUE self, VALUE key, VALUE value)
         rb_define_finalizer_no_check(key, w->final);
     }
 
+    RB_OBJ_WRITTEN(self, Qundef, value);
+
     return value;
 }