Remove Arena_pin in favor of adopting the UPB freezing API.

PiperOrigin-RevId: 639960700
protocolbuffers · Jun 4, 2024 · 1194440 · 1194440
1 parent 5d40c4f
commit 1194440
Show file tree

Hide file tree

Showing 13 changed files with 496 additions and 193 deletions.
diff --git a/ruby/ext/google/protobuf_c/map.c b/ruby/ext/google/protobuf_c/map.c
@@ -63,8 +63,8 @@ static VALUE Map_alloc(VALUE klass) {
   return TypedData_Wrap_Struct(klass, &Map_type, self);
 }
 
-VALUE Map_GetRubyWrapper(upb_Map* map, upb_CType key_type, TypeInfo value_type,
-                         VALUE arena) {
+VALUE Map_GetRubyWrapper(const upb_Map* map, upb_CType key_type,
+                         TypeInfo value_type, VALUE arena) {
   PBRUBY_ASSERT(map);
 
   VALUE val = ObjectCache_Get(map);
@@ -83,7 +83,6 @@ VALUE Map_GetRubyWrapper(upb_Map* map, upb_CType key_type, TypeInfo value_type,
     }
     return ObjectCache_TryAdd(map, val);
   }
-
   return val;
 }
 
@@ -105,8 +104,9 @@ static TypeInfo Map_keyinfo(Map* self) {
 }
 
 static upb_Map* Map_GetMutable(VALUE _self) {
-  rb_check_frozen(_self);
-  return (upb_Map*)ruby_to_Map(_self)->map;
+  const upb_Map* map = ruby_to_Map(_self)->map;
+  Protobuf_CheckNotFrozen(_self, upb_Map_IsFrozen(map));
+  return (upb_Map*)map;
 }
 
 VALUE Map_CreateHash(const upb_Map* map, upb_CType key_type,
@@ -439,14 +439,14 @@ static VALUE Map_has_key(VALUE _self, VALUE key) {
  * nil if none was present. Throws an exception if the key is of the wrong type.
  */
 static VALUE Map_delete(VALUE _self, VALUE key) {
+  upb_Map* map = Map_GetMutable(_self);
   Map* self = ruby_to_Map(_self);
-  rb_check_frozen(_self);
 
   upb_MessageValue key_upb =
       Convert_RubyToUpb(key, "", Map_keyinfo(self), NULL);
   upb_MessageValue val_upb;
 
-  if (upb_Map_Delete(Map_GetMutable(_self), key_upb, &val_upb)) {
+  if (upb_Map_Delete(map, key_upb, &val_upb)) {
     return Convert_UpbToRuby(val_upb, self->value_type_info, self->arena);
   } else {
     return Qnil;
@@ -560,29 +560,79 @@ VALUE Map_eq(VALUE _self, VALUE _other) {
 
 /*
  * call-seq:
- *     Message.freeze => self
+ *     Map.frozen? => bool
+ *
+ * Returns true if the map is frozen in either Ruby or the underlying
+ * representation. Freezes the Ruby map object if it is not already frozen in
+ * Ruby but it is frozen in the underlying representation.
+ */
+VALUE Map_frozen(VALUE _self) {
+  Map* self = ruby_to_Map(_self);
+  if (!upb_Map_IsFrozen(self->map)) {
+    PBRUBY_ASSERT(!RB_OBJ_FROZEN(_self));
+    return Qfalse;
+  }
+
+  // Lazily freeze the Ruby wrapper.
+  if (!RB_OBJ_FROZEN(_self)) RB_OBJ_FREEZE(_self);
+  return Qtrue;
+}
+
+/*
+ * call-seq:
+ *     Map.freeze => self
  *
- * Freezes the message object. We have to intercept this so we can pin the
- * Ruby object into memory so we don't forget it's frozen.
+ * Freezes the map object. We have to intercept this so we can freeze the
+ * underlying representation, not just the Ruby wrapper.
  */
 VALUE Map_freeze(VALUE _self) {
   Map* self = ruby_to_Map(_self);
+  if (RB_OBJ_FROZEN(_self)) {
+    PBRUBY_ASSERT(upb_Map_IsFrozen(self->map));
+    return _self;
+  }
+
+  if (!upb_Map_IsFrozen(self->map)) {
+    if (self->value_type_info.type == kUpb_CType_Message) {
+      upb_Map_Freeze(
+          Map_GetMutable(_self),
+          upb_MessageDef_MiniTable(self->value_type_info.def.msgdef));
+    } else {
+      upb_Map_Freeze(Map_GetMutable(_self), NULL);
+    }
+  }
 
-  if (RB_OBJ_FROZEN(_self)) return _self;
-  Arena_Pin(self->arena, _self);
   RB_OBJ_FREEZE(_self);
 
-  if (self->value_type_info.type == kUpb_CType_Message) {
-    size_t iter = kUpb_Map_Begin;
-    upb_MessageValue key, val;
+  return _self;
+}
+
+VALUE Map_EmptyFrozen(const upb_FieldDef* f) {
+  PBRUBY_ASSERT(upb_FieldDef_IsMap(f));
+  VALUE val = ObjectCache_Get(f);
 
-    while (upb_Map_Next(self->map, &key, &val, &iter)) {
-      VALUE val_val =
-          Convert_UpbToRuby(val, self->value_type_info, self->arena);
-      Message_freeze(val_val);
+  if (val == Qnil) {
+    const upb_FieldDef* key_f = map_field_key(f);
+    const upb_FieldDef* val_f = map_field_value(f);
+    upb_CType key_type = upb_FieldDef_CType(key_f);
+    TypeInfo value_type_info = TypeInfo_get(val_f);
+    val = Map_alloc(cMap);
+    Map* self;
+    TypedData_Get_Struct(val, Map, &Map_type, self);
+    self->arena = Arena_new();
+    self->map =
+        upb_Map_New(Arena_get(self->arena), key_type, value_type_info.type);
+    self->key_type = key_type;
+    self->value_type_info = value_type_info;
+    if (self->value_type_info.type == kUpb_CType_Message) {
+      const upb_MessageDef* val_m = value_type_info.def.msgdef;
+      self->value_type_class = Descriptor_DefToClass(val_m);
     }
+    return ObjectCache_TryAdd(f, Map_freeze(val));
   }
-  return _self;
+  PBRUBY_ASSERT(RB_OBJ_FROZEN(val));
+  PBRUBY_ASSERT(upb_Map_IsFrozen(ruby_to_Map(val)->map));
+  return val;
 }
 
 /*
@@ -671,6 +721,7 @@ void Map_register(VALUE module) {
   rb_define_method(klass, "clone", Map_dup, 0);
   rb_define_method(klass, "==", Map_eq, 1);
   rb_define_method(klass, "freeze", Map_freeze, 0);
+  rb_define_method(klass, "frozen?", Map_frozen, 0);
   rb_define_method(klass, "hash", Map_hash, 0);
   rb_define_method(klass, "to_h", Map_to_h, 0);
   rb_define_method(klass, "inspect", Map_inspect, 0);

diff --git a/ruby/ext/google/protobuf_c/map.h b/ruby/ext/google/protobuf_c/map.h
@@ -11,10 +11,14 @@
 #include "protobuf.h"
 #include "ruby-upb.h"
 
+// Returns a frozen sentinel Ruby wrapper object for an empty upb_Map with the
+// key and value types specified by the field. Creates one if it doesn't exist.
+VALUE Map_EmptyFrozen(const upb_FieldDef* f);
+
 // Returns a Ruby wrapper object for the given map, which will be created if
 // one does not exist already.
-VALUE Map_GetRubyWrapper(upb_Map *map, upb_CType key_type, TypeInfo value_type,
-                         VALUE arena);
+VALUE Map_GetRubyWrapper(const upb_Map *map, upb_CType key_type,
+                         TypeInfo value_type, VALUE arena);
 
 // Gets the underlying upb_Map for this Ruby map object, which must have
 // key/value type that match |field|. If this is not a map or the type doesn't

diff --git a/ruby/ext/google/protobuf_c/message.c b/ruby/ext/google/protobuf_c/message.c
@@ -80,11 +80,12 @@ const upb_Message* Message_Get(VALUE msg_rb, const upb_MessageDef** m) {
 }
 
 upb_Message* Message_GetMutable(VALUE msg_rb, const upb_MessageDef** m) {
-  rb_check_frozen(msg_rb);
-  return (upb_Message*)Message_Get(msg_rb, m);
+  const upb_Message* upb_msg = Message_Get(msg_rb, m);
+  Protobuf_CheckNotFrozen(msg_rb, upb_Message_IsFrozen(upb_msg));
+  return (upb_Message*)upb_msg;
 }
 
-void Message_InitPtr(VALUE self_, upb_Message* msg, VALUE arena) {
+void Message_InitPtr(VALUE self_, const upb_Message* msg, VALUE arena) {
   Message* self = ruby_to_Message(self_);
   self->msg = msg;
   RB_OBJ_WRITE(self_, &self->arena, arena);
@@ -105,7 +106,7 @@ void Message_CheckClass(VALUE klass) {
   }
 }
 
-VALUE Message_GetRubyWrapper(upb_Message* msg, const upb_MessageDef* m,
+VALUE Message_GetRubyWrapper(const upb_Message* msg, const upb_MessageDef* m,
                              VALUE arena) {
   if (msg == NULL) return Qnil;
 
@@ -116,7 +117,6 @@ VALUE Message_GetRubyWrapper(upb_Message* msg, const upb_MessageDef* m,
     val = Message_alloc(klass);
     Message_InitPtr(val, msg, arena);
   }
-
   return val;
 }
 
@@ -288,13 +288,42 @@ static void Message_setfield(upb_Message* msg, const upb_FieldDef* f, VALUE val,
   upb_Message_SetFieldByDef(msg, f, msgval, arena);
 }
 
+VALUE Message_getfield_frozen(const upb_Message* msg, const upb_FieldDef* f,
+                              VALUE arena) {
+  upb_MessageValue msgval = upb_Message_GetFieldByDef(msg, f);
+  if (upb_FieldDef_IsMap(f)) {
+    if (msgval.map_val == NULL) {
+      return Map_EmptyFrozen(f);
+    }
+    const upb_FieldDef* key_f = map_field_key(f);
+    const upb_FieldDef* val_f = map_field_value(f);
+    upb_CType key_type = upb_FieldDef_CType(key_f);
+    TypeInfo value_type_info = TypeInfo_get(val_f);
+    return Map_GetRubyWrapper(msgval.map_val, key_type, value_type_info, Qnil);
+  }
+  if (upb_FieldDef_IsRepeated(f)) {
+    if (msgval.array_val == NULL) {
+      return RepeatedField_EmptyFrozen(f);
+    }
+    return RepeatedField_GetRubyWrapper(msgval.array_val, TypeInfo_get(f),
+                                        Qnil);
+  }
+  VALUE ret;
+  if (upb_FieldDef_IsSubMessage(f)) {
+    const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
+    ret = Message_GetRubyWrapper(msgval.msg_val, m, arena);
+  } else {
+    ret = Convert_UpbToRuby(msgval, TypeInfo_get(f), Qnil);
+  }
+  return ret;
+}
+
 VALUE Message_getfield(VALUE _self, const upb_FieldDef* f) {
   Message* self = ruby_to_Message(_self);
-  // This is a special-case: upb_Message_Mutable() for map & array are logically
-  // const (they will not change what is serialized) but physically
-  // non-const, as they do allocate a repeated field or map. The logical
-  // constness means it's ok to do even if the message is frozen.
-  upb_Message* msg = (upb_Message*)self->msg;
+  if (upb_Message_IsFrozen(self->msg)) {
+    return Message_getfield_frozen(self->msg, f, self->arena);
+  }
+  upb_Message* msg = Message_GetMutable(_self, NULL);
   upb_Arena* arena = Arena_get(self->arena);
   if (upb_FieldDef_IsMap(f)) {
     upb_Map* map = upb_Message_Mutable(msg, f, arena).map;
@@ -307,12 +336,12 @@ VALUE Message_getfield(VALUE _self, const upb_FieldDef* f) {
     upb_Array* arr = upb_Message_Mutable(msg, f, arena).array;
     return RepeatedField_GetRubyWrapper(arr, TypeInfo_get(f), self->arena);
   } else if (upb_FieldDef_IsSubMessage(f)) {
-    if (!upb_Message_HasFieldByDef(self->msg, f)) return Qnil;
+    if (!upb_Message_HasFieldByDef(msg, f)) return Qnil;
     upb_Message* submsg = upb_Message_Mutable(msg, f, arena).msg;
     const upb_MessageDef* m = upb_FieldDef_MessageSubDef(f);
     return Message_GetRubyWrapper(submsg, m, self->arena);
   } else {
-    upb_MessageValue msgval = upb_Message_GetFieldByDef(self->msg, f);
+    upb_MessageValue msgval = upb_Message_GetFieldByDef(msg, f);
     return Convert_UpbToRuby(msgval, TypeInfo_get(f), self->arena);
   }
 }
@@ -436,7 +465,6 @@ static VALUE Message_method_missing(int argc, VALUE* argv, VALUE _self) {
       if (argc != 2) {
         rb_raise(rb_eArgError, "Expected 2 arguments, received %d", argc);
       }
-      rb_check_frozen(_self);
       break;
     default:
       if (argc != 1) {
@@ -812,33 +840,42 @@ static VALUE Message_to_h(VALUE _self) {
 
 /*
  * call-seq:
- *     Message.freeze => self
+ *     Message.frozen? => bool
  *
- * Freezes the message object. We have to intercept this so we can pin the
- * Ruby object into memory so we don't forget it's frozen.
+ * Returns true if the message is frozen in either Ruby or the underlying
+ * representation. Freezes the Ruby message object if it is not already frozen
+ * in Ruby but it is frozen in the underlying representation.
  */
-VALUE Message_freeze(VALUE _self) {
+VALUE Message_frozen(VALUE _self) {
   Message* self = ruby_to_Message(_self);
+  if (!upb_Message_IsFrozen(self->msg)) {
+    PBRUBY_ASSERT(!RB_OBJ_FROZEN(_self));
+    return Qfalse;
+  }
 
-  if (RB_OBJ_FROZEN(_self)) return _self;
-  Arena_Pin(self->arena, _self);
-  RB_OBJ_FREEZE(_self);
+  // Lazily freeze the Ruby wrapper.
+  if (!RB_OBJ_FROZEN(_self)) RB_OBJ_FREEZE(_self);
+  return Qtrue;
+}
 
-  int n = upb_MessageDef_FieldCount(self->msgdef);
-  for (int i = 0; i < n; i++) {
-    const upb_FieldDef* f = upb_MessageDef_Field(self->msgdef, i);
-    VALUE field = Message_getfield(_self, f);
-
-    if (field != Qnil) {
-      if (upb_FieldDef_IsMap(f)) {
-        Map_freeze(field);
-      } else if (upb_FieldDef_IsRepeated(f)) {
-        RepeatedField_freeze(field);
-      } else if (upb_FieldDef_IsSubMessage(f)) {
-        Message_freeze(field);
-      }
-    }
+/*
+ * call-seq:
+ *     Message.freeze => self
+ *
+ * Freezes the message object. We have to intercept this so we can freeze the
+ * underlying representation, not just the Ruby wrapper.
+ */
+VALUE Message_freeze(VALUE _self) {
+  Message* self = ruby_to_Message(_self);
+  if (RB_OBJ_FROZEN(_self)) {
+    PBRUBY_ASSERT(upb_Message_IsFrozen(self->msg));
+    return _self;
   }
+  if (!upb_Message_IsFrozen(self->msg)) {
+    upb_Message_Freeze(Message_GetMutable(_self, NULL),
+                       upb_MessageDef_MiniTable(self->msgdef));
+  }
+  RB_OBJ_FREEZE(_self);
   return _self;
 }
 
@@ -1352,6 +1389,7 @@ static void Message_define_class(VALUE klass) {
   rb_define_method(klass, "==", Message_eq, 1);
   rb_define_method(klass, "eql?", Message_eq, 1);
   rb_define_method(klass, "freeze", Message_freeze, 0);
+  rb_define_method(klass, "frozen?", Message_frozen, 0);
   rb_define_method(klass, "hash", Message_hash, 0);
   rb_define_method(klass, "to_h", Message_to_h, 0);
   rb_define_method(klass, "inspect", Message_inspect, 0);

diff --git a/ruby/ext/google/protobuf_c/message.h b/ruby/ext/google/protobuf_c/message.h
@@ -36,7 +36,7 @@ const upb_Message* Message_GetUpbMessage(VALUE value, const upb_MessageDef* m,
 
 // Gets or constructs a Ruby wrapper object for the given message. The wrapper
 // object will reference |arena| and ensure that it outlives this object.
-VALUE Message_GetRubyWrapper(upb_Message* msg, const upb_MessageDef* m,
+VALUE Message_GetRubyWrapper(const upb_Message* msg, const upb_MessageDef* m,
                              VALUE arena);
 
 // Gets the given field from this message.

diff --git a/ruby/ext/google/protobuf_c/protobuf.c b/ruby/ext/google/protobuf_c/protobuf.c
@@ -221,15 +221,6 @@ void Arena_fuse(VALUE _arena, upb_Arena *other) {
 
 VALUE Arena_new() { return Arena_alloc(cArena); }
 
-void Arena_Pin(VALUE _arena, VALUE obj) {
-  Arena *arena;
-  TypedData_Get_Struct(_arena, Arena, &Arena_type, arena);
-  if (arena->pinned_objs == Qnil) {
-    RB_OBJ_WRITE(_arena, &arena->pinned_objs, rb_ary_new());
-  }
-  rb_ary_push(arena->pinned_objs, obj);
-}
-
 void Arena_register(VALUE module) {
   VALUE internal = rb_define_module_under(module, "Internal");
   VALUE klass = rb_define_class_under(internal, "Arena", rb_cObject);
@@ -354,3 +345,14 @@ __attribute__((visibility("default"))) void Init_protobuf_c() {
   rb_define_singleton_method(protobuf, "deep_copy", Google_Protobuf_deep_copy,
                              1);
 }
+
+// -----------------------------------------------------------------------------
+// Utilities
+// -----------------------------------------------------------------------------
+
+// Raises a Ruby error if val is frozen in Ruby or UPB.
+void Protobuf_CheckNotFrozen(VALUE val, bool upb_frozen) {
+  if (RB_UNLIKELY(rb_obj_frozen_p(val)||upb_frozen)) {
+    rb_error_frozen_object(val);
+  }
+}