Addressed code-review comments.

6 files changed, 72 insertions, 135 deletions

diff --git a/ruby/ext/google/protobuf_c/defs.c b/ruby/ext/google/protobuf_c/defs.c
index 499e041b..a18aaac4 100644
--- a/ruby/ext/google/protobuf_c/defs.c
+++ b/ruby/ext/google/protobuf_c/defs.c
@@ -226,6 +226,7 @@ DEFINE_CLASS(Descriptor, "Google::Protobuf::Descriptor");
 void Descriptor_mark(void* _self) {
   Descriptor* self = _self;
   rb_gc_mark(self->klass);
+  rb_gc_mark(self->typeclass_references);
 }
 
 void Descriptor_free(void* _self) {
@@ -270,6 +271,7 @@ VALUE Descriptor_alloc(VALUE klass) {
   self->fill_method = NULL;
   self->pb_serialize_handlers = NULL;
   self->json_serialize_handlers = NULL;
+  self->typeclass_references = rb_ary_new();
   return ret;
 }
 
diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index f1b951fc..1cff9049 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -64,7 +64,7 @@ static const void *newsubmsghandlerdata(upb_handlers* h, uint32_t ofs,
 static void *startseq_handler(void* closure, const void* hd) {
   MessageHeader* msg = closure;
   const size_t *ofs = hd;
-  return (void*)DEREF(Message_data(msg), *ofs, VALUE);
+  return (void*)DEREF(msg, *ofs, VALUE);
 }
 
 // Handlers that append primitive values to a repeated field (a regular Ruby
@@ -115,7 +115,7 @@ static void* str_handler(void *closure,
   const size_t *ofs = hd;
   VALUE str = rb_str_new2("");
   rb_enc_associate(str, kRubyStringUtf8Encoding);
-  DEREF(Message_data(msg), *ofs, VALUE) = str;
+  DEREF(msg, *ofs, VALUE) = str;
   return (void*)str;
 }
 
@@ -127,7 +127,7 @@ static void* bytes_handler(void *closure,
   const size_t *ofs = hd;
   VALUE str = rb_str_new2("");
   rb_enc_associate(str, kRubyString8bitEncoding);
-  DEREF(Message_data(msg), *ofs, VALUE) = str;
+  DEREF(msg, *ofs, VALUE) = str;
   return (void*)str;
 }
 
@@ -163,12 +163,12 @@ static void *submsg_handler(void *closure, const void *hd) {
       get_def_obj((void*)submsgdata->md);
   VALUE subklass = Descriptor_msgclass(subdesc);
 
-  if (DEREF(Message_data(msg), submsgdata->ofs, VALUE) == Qnil) {
-    DEREF(Message_data(msg), submsgdata->ofs, VALUE) =
+  if (DEREF(msg, submsgdata->ofs, VALUE) == Qnil) {
+    DEREF(msg, submsgdata->ofs, VALUE) =
         rb_class_new_instance(0, NULL, subklass);
   }
 
-  VALUE submsg_rb = DEREF(Message_data(msg), submsgdata->ofs, VALUE);
+  VALUE submsg_rb = DEREF(msg, submsgdata->ofs, VALUE);
   MessageHeader* submsg;
   TypedData_Get_Struct(submsg_rb, MessageHeader, &Message_type, submsg);
   return submsg;
@@ -199,7 +199,7 @@ typedef struct {
 static void *startmapentry_handler(void *closure, const void *hd) {
   MessageHeader* msg = closure;
   const map_handlerdata_t* mapdata = hd;
-  VALUE map_rb = DEREF(Message_data(msg), mapdata->ofs, VALUE);
+  VALUE map_rb = DEREF(msg, mapdata->ofs, VALUE);
 
   map_parse_frame_t* frame = ALLOC(map_parse_frame_t);
   frame->map = map_rb;
@@ -238,7 +238,8 @@ static bool endmap_handler(void *closure, const void *hd, upb_status* s) {
 // pass the handlerdata down to the sub-message handler setup.
 static map_handlerdata_t* new_map_handlerdata(
     size_t ofs,
-    const upb_msgdef* mapentry_def) {
+    const upb_msgdef* mapentry_def,
+    Descriptor* desc) {
 
   map_handlerdata_t* hd = ALLOC(map_handlerdata_t);
   hd->ofs = ofs;
@@ -252,6 +253,11 @@ static map_handlerdata_t* new_map_handlerdata(
   hd->value_field_type = upb_fielddef_type(value_field);
   hd->value_field_typeclass = field_type_class(value_field);
 
+  // Ensure that value_field_typeclass is properly GC-rooted.
+  if (hd->value_field_typeclass != Qnil) {
+    rb_ary_push(desc->typeclass_references, hd->value_field_typeclass);
+  }
+
   return hd;
 }
 
@@ -314,9 +320,7 @@ static void add_handlers_for_singular_field(upb_handlers *h,
     case UPB_TYPE_INT64:
     case UPB_TYPE_UINT64:
     case UPB_TYPE_DOUBLE:
-      // The shim writes directly at the given offset (instead of using
-      // DEREF()) so we need to add the msg overhead.
-      upb_shim_set(h, f, offset + sizeof(MessageHeader), -1);
+      upb_shim_set(h, f, offset, -1);
       break;
     case UPB_TYPE_STRING:
     case UPB_TYPE_BYTES: {
@@ -343,9 +347,10 @@ static void add_handlers_for_singular_field(upb_handlers *h,
 // Adds handlers to a map field.
 static void add_handlers_for_mapfield(upb_handlers* h,
                                       const upb_fielddef* fielddef,
-                                      size_t offset) {
+                                      size_t offset,
+                                      Descriptor* desc) {
   const upb_msgdef* map_msgdef = upb_fielddef_msgsubdef(fielddef);
-  map_handlerdata_t* hd = new_map_handlerdata(offset, map_msgdef);
+  map_handlerdata_t* hd = new_map_handlerdata(offset, map_msgdef, desc);
   upb_handlers_addcleanup(h, hd, free);
   upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
   upb_handlerattr_sethandlerdata(&attr, hd);
@@ -355,10 +360,11 @@ static void add_handlers_for_mapfield(upb_handlers* h,
 
 // Adds handlers to a map-entry msgdef.
 static void add_handlers_for_mapentry(const upb_msgdef* msgdef,
-                                      upb_handlers* h) {
+                                      upb_handlers* h,
+                                      Descriptor* desc) {
   const upb_fielddef* key_field = map_entry_key(msgdef);
   const upb_fielddef* value_field = map_entry_value(msgdef);
-  map_handlerdata_t* hd = new_map_handlerdata(0, msgdef);
+  map_handlerdata_t* hd = new_map_handlerdata(0, msgdef, desc);
   upb_handlers_addcleanup(h, hd, free);
   upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
   upb_handlerattr_sethandlerdata(&attr, hd);
@@ -366,15 +372,10 @@ static void add_handlers_for_mapentry(const upb_msgdef* msgdef,
 
   add_handlers_for_singular_field(
       h, key_field,
-      // Convert the offset into map_parse_frame_t to an offset understood by the
-      // singular field handlers, so that we don't have to use special
-      // map-key/value-specific handlers. The ordinary singular field handlers expect
-      // a Message* and assume offset is relative to the data section at the end, so
-      // we compensate for that addition.
-      offsetof(map_parse_frame_t, key_storage) - sizeof(MessageHeader));
+      offsetof(map_parse_frame_t, key_storage));
   add_handlers_for_singular_field(
       h, value_field,
-      offsetof(map_parse_frame_t, value_storage) - sizeof(MessageHeader));
+      offsetof(map_parse_frame_t, value_storage));
 }
 
 static void add_handlers_for_message(const void *closure, upb_handlers *h) {
@@ -384,7 +385,7 @@ static void add_handlers_for_message(const void *closure, upb_handlers *h) {
   // If this is a mapentry message type, set up a special set of handlers and
   // bail out of the normal (user-defined) message type handling.
   if (upb_msgdef_mapentry(msgdef)) {
-    add_handlers_for_mapentry(msgdef, h);
+    add_handlers_for_mapentry(msgdef, h, desc);
     return;
   }
 
@@ -402,10 +403,11 @@ static void add_handlers_for_message(const void *closure, upb_handlers *h) {
        !upb_msg_done(&i);
        upb_msg_next(&i)) {
     const upb_fielddef *f = upb_msg_iter_field(&i);
-    size_t offset = desc->layout->offsets[upb_fielddef_index(f)];
+    size_t offset = desc->layout->offsets[upb_fielddef_index(f)] +
+        sizeof(MessageHeader);
 
     if (is_map_field(f)) {
-      add_handlers_for_mapfield(h, f, offset);
+      add_handlers_for_mapfield(h, f, offset, desc);
     } else if (upb_fielddef_isseq(f)) {
       add_handlers_for_repeated_field(h, f, offset);
     } else {
@@ -796,38 +798,38 @@ static void putmsg(VALUE msg_rb, const Descriptor* desc,
 
   MessageHeader* msg;
   TypedData_Get_Struct(msg_rb, MessageHeader, &Message_type, msg);
-  void* msg_data = Message_data(msg);
 
   upb_msg_iter i;
   for (upb_msg_begin(&i, desc->msgdef);
        !upb_msg_done(&i);
        upb_msg_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
-    uint32_t offset = desc->layout->offsets[upb_fielddef_index(f)];
+    uint32_t offset =
+        desc->layout->offsets[upb_fielddef_index(f)] + sizeof(MessageHeader);
 
     if (is_map_field(f)) {
-      VALUE map = DEREF(msg_data, offset, VALUE);
+      VALUE map = DEREF(msg, offset, VALUE);
       if (map != Qnil) {
         putmap(map, f, sink, depth);
       }
     } else if (upb_fielddef_isseq(f)) {
-      VALUE ary = DEREF(msg_data, offset, VALUE);
+      VALUE ary = DEREF(msg, offset, VALUE);
       if (ary != Qnil) {
         putary(ary, f, sink, depth);
       }
     } else if (upb_fielddef_isstring(f)) {
-      VALUE str = DEREF(msg_data, offset, VALUE);
+      VALUE str = DEREF(msg, offset, VALUE);
       if (RSTRING_LEN(str) > 0) {
         putstr(str, f, sink);
       }
     } else if (upb_fielddef_issubmsg(f)) {
-      putsubmsg(DEREF(msg_data, offset, VALUE), f, sink, depth);
+      putsubmsg(DEREF(msg, offset, VALUE), f, sink, depth);
     } else {
       upb_selector_t sel = getsel(f, upb_handlers_getprimitivehandlertype(f));
 
 #define T(upbtypeconst, upbtype, ctype, default_value)                \
   case upbtypeconst: {                                                \
-      ctype value = DEREF(msg_data, offset, ctype);                   \
+      ctype value = DEREF(msg, offset, ctype);                        \
       if (value != default_value) {                                   \
         upb_sink_put##upbtype(sink, sel, value);                      \
       }                                                               \
diff --git a/ruby/ext/google/protobuf_c/map.c b/ruby/ext/google/protobuf_c/map.c
index 56548799..2c2ae240 100644
--- a/ruby/ext/google/protobuf_c/map.c
+++ b/ruby/ext/google/protobuf_c/map.c
@@ -49,8 +49,14 @@
 // field values), but keys are a bit special. Since we use a strtable, we need
 // to store keys as sequences of bytes such that equality of those bytes maps
 // one-to-one to equality of keys. We store strings directly (i.e., they map to
-// their own bytes) and integers as sequences of either 4 or 8 bytes in
-// host-byte-order as either a uint32_t or a uint64_t.
+// their own bytes) and integers as native integers (using the native_slot
+// abstraction).
+
+// Note that there is another tradeoff here in keeping string keys as native
+// strings rather than Ruby strings: traversing the Map requires conversion to
+// Ruby string values on every traversal, potentially creating more garbage. We
+// should consider ways to cache a Ruby version of the key if this becomes an
+// issue later.
 
 // Forms a key to use with the underlying strtable from a Ruby key value. |buf|
 // must point to TABLE_KEY_BUF_LENGTH bytes of temporary space, used to
@@ -73,61 +79,13 @@ static void table_key(Map* self, VALUE key,
 
     case UPB_TYPE_BOOL:
     case UPB_TYPE_INT32:
-    case UPB_TYPE_INT64: {
-      // Signed numeric types: use an int64 in host-native byte order.
-      int64_t key_val = 0;
-
-      // Do a range/value check.
-      switch (self->key_type) {
-        case UPB_TYPE_BOOL:
-          if (key != Qtrue && key != Qfalse) {
-            rb_raise(rb_eTypeError, "Key must be true or false");
-          }
-          key_val = (key == Qtrue) ? 1 : 0;
-          break;
-        case UPB_TYPE_INT32:
-          native_slot_check_int_range_precision(self->key_type, key);
-          key_val = NUM2INT(key);
-          break;
-        case UPB_TYPE_INT64:
-          native_slot_check_int_range_precision(self->key_type, key);
-          key_val = NUM2LL(key);
-          break;
-        default:
-          break;
-      }
-
-      int64_t* int64_key = (int64_t*)buf;
-      *int64_key = key_val;
-      *out_key = buf;
-      *out_length = sizeof(int64_t);
-      break;
-    }
-
+    case UPB_TYPE_INT64:
     case UPB_TYPE_UINT32:
-    case UPB_TYPE_UINT64: {
-      // Unsigned numeric types: use a uint64 in host-native byte order.
-      uint64_t key_val = 0;
-
-      // Do a range/value check.
-      native_slot_check_int_range_precision(self->key_type, key);
-      switch (self->key_type) {
-        case UPB_TYPE_UINT32:
-          key_val = NUM2UINT(key);
-          break;
-        case UPB_TYPE_UINT64:
-          key_val = NUM2ULL(key);
-          break;
-        default:
-          break;
-      }
-
-      uint64_t* uint64_key = (uint64_t*)buf;
-      *uint64_key = key_val;
+    case UPB_TYPE_UINT64:
+      native_slot_set(self->key_type, Qnil, buf, key);
       *out_key = buf;
-      *out_length = sizeof(uint64_t);
+      *out_length = native_slot_size(self->key_type);
       break;
-    }
 
     default:
       // Map constructor should not allow a Map with another key type to be
@@ -150,48 +108,14 @@ static VALUE table_key_to_ruby(Map* self, const char* buf, size_t length) {
 
     case UPB_TYPE_BOOL:
     case UPB_TYPE_INT32:
-    case UPB_TYPE_INT64: {
-      assert(length == sizeof(int64_t));
-      int64_t* int64_key = (int64_t*)buf;
-
-      if (self->key_type == UPB_TYPE_BOOL) {
-        return *int64_key ? Qtrue : Qfalse;
-      } else {
-        return LL2NUM(*int64_key);
-      }
-    }
-
-    case UPB_TYPE_UINT32:
-    case UPB_TYPE_UINT64: {
-      assert(length == sizeof(uint64_t));
-      uint64_t* uint64_key = (uint64_t*)buf;
-      return ULL2NUM(*uint64_key);
-    }
-
-    default:
-      assert(false);
-      return Qnil;
-  }
-}
-
-static upb_ctype_t upb_table_value_type(upb_fieldtype_t value_type) {
-  switch (value_type) {
-    case UPB_TYPE_BOOL:
-    case UPB_TYPE_INT32:
     case UPB_TYPE_INT64:
     case UPB_TYPE_UINT32:
     case UPB_TYPE_UINT64:
-    case UPB_TYPE_ENUM:
-    case UPB_TYPE_FLOAT:
-    case UPB_TYPE_DOUBLE:
-    case UPB_TYPE_STRING:
-    case UPB_TYPE_BYTES:
-    case UPB_TYPE_MESSAGE:
-      return UPB_CTYPE_UINT64;
+      return native_slot_get(self->key_type, Qnil, buf);
 
     default:
       assert(false);
-      return 0;
+      return Qnil;
   }
 }
 
@@ -321,7 +245,9 @@ VALUE Map_init(int argc, VALUE* argv, VALUE _self) {
     init_value_arg = 3;
   }
 
-  if (!upb_strtable_init(&self->table, upb_table_value_type(self->value_type))) {
+  // Table value type is always UINT64: this ensures enough space to store the
+  // native_slot value.
+  if (!upb_strtable_init(&self->table, UPB_CTYPE_UINT64)) {
     rb_raise(rb_eRuntimeError, "Could not allocate table.");
   }
 
@@ -528,8 +454,7 @@ VALUE Map_clear(VALUE _self) {
   // Uninit and reinit the table -- this is faster than iterating and doing a
   // delete-lookup on each key.
   upb_strtable_uninit(&self->table);
-  if (!upb_strtable_init(&self->table,
-                         upb_table_value_type(self->value_type))) {
+  if (!upb_strtable_init(&self->table, UPB_CTYPE_INT64)) {
     rb_raise(rb_eRuntimeError, "Unable to re-initialize table");
   }
   return Qnil;
diff --git a/ruby/ext/google/protobuf_c/protobuf.h b/ruby/ext/google/protobuf_c/protobuf.h
index 6dac6029..88ae62e4 100644
--- a/ruby/ext/google/protobuf_c/protobuf.h
+++ b/ruby/ext/google/protobuf_c/protobuf.h
@@ -110,6 +110,10 @@ struct Descriptor {
   const upb_pbdecodermethod* fill_method;
   const upb_handlers* pb_serialize_handlers;
   const upb_handlers* json_serialize_handlers;
+  // Handlers hold type class references for sub-message fields directly in some
+  // cases. We need to keep these rooted because they might otherwise be
+  // collected.
+  VALUE typeclass_references;
 };
 
 struct FieldDescriptor {
@@ -252,7 +256,7 @@ void native_slot_set(upb_fieldtype_t type,
                      VALUE value);
 VALUE native_slot_get(upb_fieldtype_t type,
                       VALUE type_class,
-                      void* memory);
+                      const void* memory);
 void native_slot_init(upb_fieldtype_t type, void* memory);
 void native_slot_mark(upb_fieldtype_t type, void* memory);
 void native_slot_dup(upb_fieldtype_t type, void* to, void* from);
@@ -389,7 +393,7 @@ struct MessageLayout {
 MessageLayout* create_layout(const upb_msgdef* msgdef);
 void free_layout(MessageLayout* layout);
 VALUE layout_get(MessageLayout* layout,
-                 void* storage,
+                 const void* storage,
                  const upb_fielddef* field);
 void layout_set(MessageLayout* layout,
                 void* storage,
diff --git a/ruby/ext/google/protobuf_c/storage.c b/ruby/ext/google/protobuf_c/storage.c
index 235fbff2..14f49d44 100644
--- a/ruby/ext/google/protobuf_c/storage.c
+++ b/ruby/ext/google/protobuf_c/storage.c
@@ -200,7 +200,9 @@ void native_slot_set(upb_fieldtype_t type, VALUE type_class,
   }
 }
 
-VALUE native_slot_get(upb_fieldtype_t type, VALUE type_class, void* memory) {
+VALUE native_slot_get(upb_fieldtype_t type,
+                      VALUE type_class,
+                      const void* memory) {
   switch (type) {
     case UPB_TYPE_FLOAT:
       return DBL2NUM(DEREF(memory, float));
@@ -211,7 +213,7 @@ VALUE native_slot_get(upb_fieldtype_t type, VALUE type_class, void* memory) {
     case UPB_TYPE_STRING:
     case UPB_TYPE_BYTES:
     case UPB_TYPE_MESSAGE:
-      return *((VALUE *)memory);
+      return DEREF(memory, VALUE);
     case UPB_TYPE_ENUM: {
       int32_t val = DEREF(memory, int32_t);
       VALUE symbol = enum_lookup(type_class, INT2NUM(val));
@@ -332,19 +334,19 @@ bool is_map_field(const upb_fielddef* field) {
       upb_fielddef_type(field) != UPB_TYPE_MESSAGE) {
     return false;
   }
-  const upb_msgdef* subdef = (const upb_msgdef*)upb_fielddef_subdef(field);
+  const upb_msgdef* subdef = upb_fielddef_msgsubdef(field);
   return upb_msgdef_mapentry(subdef);
 }
 
 const upb_fielddef* map_field_key(const upb_fielddef* field) {
   assert(is_map_field(field));
-  const upb_msgdef* subdef = (const upb_msgdef*)upb_fielddef_subdef(field);
+  const upb_msgdef* subdef = upb_fielddef_msgsubdef(field);
   return map_entry_key(subdef);
 }
 
 const upb_fielddef* map_field_value(const upb_fielddef* field) {
   assert(is_map_field(field));
-  const upb_msgdef* subdef = (const upb_msgdef*)upb_fielddef_subdef(field);
+  const upb_msgdef* subdef = upb_fielddef_msgsubdef(field);
   return map_entry_value(subdef);
 }
 
@@ -414,7 +416,7 @@ VALUE field_type_class(const upb_fielddef* field) {
 }
 
 VALUE layout_get(MessageLayout* layout,
-                 void* storage,
+                 const void* storage,
                  const upb_fielddef* field) {
   void* memory = ((uint8_t *)storage) +
       layout->offsets[upb_fielddef_index(field)];
diff --git a/ruby/ext/google/protobuf_c/upb.c b/ruby/ext/google/protobuf_c/upb.c
index 0015aad1..abea7711 100644
--- a/ruby/ext/google/protobuf_c/upb.c
+++ b/ruby/ext/google/protobuf_c/upb.c
@@ -3420,8 +3420,10 @@ char *upb_strdup2(const char *s, size_t len) {
   // have a null-terminating byte since it may be a raw binary buffer.
   size_t n = len + 1;
   char *p = malloc(n);
-  if (p) memcpy(p, s, len);
-  p[len] = 0;
+  if (p) {
+    memcpy(p, s, len);
+    p[len] = 0;
+  }
   return p;
 }