Addressed code-review comments.

4 files changed, 65 insertions, 58 deletions

diff --git a/ruby/ext/google/protobuf_c/defs.c b/ruby/ext/google/protobuf_c/defs.c
index 371939ae..446a5586 100644
--- a/ruby/ext/google/protobuf_c/defs.c
+++ b/ruby/ext/google/protobuf_c/defs.c
@@ -287,7 +287,7 @@ void Descriptor_register(VALUE module) {
   rb_define_method(klass, "lookup", Descriptor_lookup, 1);
   rb_define_method(klass, "add_field", Descriptor_add_field, 1);
   rb_define_method(klass, "add_oneof", Descriptor_add_oneof, 1);
-  rb_define_method(klass, "oneofs", Descriptor_oneofs, 0);
+  rb_define_method(klass, "each_oneof", Descriptor_each_oneof, 0);
   rb_define_method(klass, "lookup_oneof", Descriptor_lookup_oneof, 1);
   rb_define_method(klass, "msgclass", Descriptor_msgclass, 0);
   rb_define_method(klass, "name", Descriptor_name, 0);
@@ -409,23 +409,23 @@ VALUE Descriptor_add_oneof(VALUE _self, VALUE obj) {
 
 /*
  * call-seq:
- *     Descriptor.oneofs => list of OneofDescriptors
+ *     Descriptor.each_oneof(&block) => nil
  *
- * Returns a list of OneofDescriptors that are part of this message type.
+ * Invokes the given block for each oneof in this message type, passing the
+ * corresponding OneofDescriptor.
  */
-VALUE Descriptor_oneofs(VALUE _self) {
+VALUE Descriptor_each_oneof(VALUE _self) {
   DEFINE_SELF(Descriptor, self, _self);
 
-  VALUE ret = rb_ary_new();
   upb_msg_oneof_iter it;
   for (upb_msg_oneof_begin(&it, self->msgdef);
        !upb_msg_oneof_done(&it);
        upb_msg_oneof_next(&it)) {
     const upb_oneofdef* oneof = upb_msg_iter_oneof(&it);
     VALUE obj = get_def_obj(oneof);
-    rb_ary_push(ret, obj);
+    rb_yield(obj);
   }
-  return ret;
+  return Qnil;
 }
 
 /*
diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index 26a038b6..5e53b694 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -60,10 +60,10 @@ static const void *newsubmsghandlerdata(upb_handlers* h, uint32_t ofs,
 }
 
 typedef struct {
-  size_t ofs;      // union data slot
-  size_t case_ofs; // oneof_case field
-  uint32_t tag;    // tag number to place in data slot
-  const upb_msgdef *md; // msgdef, for oneof submessage handler
+  size_t ofs;              // union data slot
+  size_t case_ofs;         // oneof_case field
+  uint32_t oneof_case_num; // oneof-case number to place in oneof_case field
+  const upb_msgdef *md;    // msgdef, for oneof submessage handler
 } oneof_handlerdata_t;
 
 static const void *newoneofhandlerdata(upb_handlers *h,
@@ -73,7 +73,12 @@ static const void *newoneofhandlerdata(upb_handlers *h,
   oneof_handlerdata_t *hd = ALLOC(oneof_handlerdata_t);
   hd->ofs = ofs;
   hd->case_ofs = case_ofs;
-  hd->tag = upb_fielddef_number(f);
+  // We reuse the field tag number as a oneof union discriminant tag. Note that
+  // we don't expose these numbers to the user, so the only requirement is that
+  // we have some unique ID for each union case/possibility. The field tag
+  // numbers are already present and are easy to use so there's no reason to
+  // create a separate ID space.
+  hd->oneof_case_num = upb_fielddef_number(f);
   if (upb_fielddef_type(f) == UPB_TYPE_MESSAGE) {
     hd->md = upb_fielddef_msgsubdef(f);
   } else {
@@ -294,7 +299,8 @@ static map_handlerdata_t* new_map_handlerdata(
   static bool oneof##type##_handler(void *closure, const void *hd,  \
                                      ctype val) {                   \
     const oneof_handlerdata_t *oneofdata = hd;                      \
-    DEREF(closure, oneofdata->case_ofs, uint32_t) = oneofdata->tag; \
+    DEREF(closure, oneofdata->case_ofs, uint32_t) =                 \
+        oneofdata->oneof_case_num;                                  \
     DEREF(closure, oneofdata->ofs, ctype) = val;                    \
     return true;                                                    \
   }
@@ -317,7 +323,8 @@ static void *oneofstr_handler(void *closure,
   const oneof_handlerdata_t *oneofdata = hd;
   VALUE str = rb_str_new2("");
   rb_enc_associate(str, kRubyStringUtf8Encoding);
-  DEREF(msg, oneofdata->case_ofs, uint32_t) = oneofdata->tag;
+  DEREF(msg, oneofdata->case_ofs, uint32_t) =
+      oneofdata->oneof_case_num;
   DEREF(msg, oneofdata->ofs, VALUE) = str;
   return (void*)str;
 }
@@ -329,7 +336,8 @@ static void *oneofbytes_handler(void *closure,
   const oneof_handlerdata_t *oneofdata = hd;
   VALUE str = rb_str_new2("");
   rb_enc_associate(str, kRubyString8bitEncoding);
-  DEREF(msg, oneofdata->case_ofs, uint32_t) = oneofdata->tag;
+  DEREF(msg, oneofdata->case_ofs, uint32_t) =
+      oneofdata->oneof_case_num;
   DEREF(msg, oneofdata->ofs, VALUE) = str;
   return (void*)str;
 }
@@ -340,13 +348,14 @@ static void *oneofsubmsg_handler(void *closure,
   MessageHeader* msg = closure;
   const oneof_handlerdata_t *oneofdata = hd;
   uint32_t oldcase = DEREF(msg, oneofdata->case_ofs, uint32_t);
-  DEREF(msg, oneofdata->case_ofs, uint32_t) = oneofdata->tag;
+  DEREF(msg, oneofdata->case_ofs, uint32_t) =
+      oneofdata->oneof_case_num;
 
   VALUE subdesc =
       get_def_obj((void*)oneofdata->md);
   VALUE subklass = Descriptor_msgclass(subdesc);
 
-  if (oldcase != oneofdata->tag ||
+  if (oldcase != oneofdata->oneof_case_num ||
       DEREF(msg, oneofdata->ofs, VALUE) == Qnil) {
     DEREF(msg, oneofdata->ofs, VALUE) =
         rb_class_new_instance(0, NULL, subklass);
diff --git a/ruby/ext/google/protobuf_c/protobuf.h b/ruby/ext/google/protobuf_c/protobuf.h
index 84f318b3..5b94ae26 100644
--- a/ruby/ext/google/protobuf_c/protobuf.h
+++ b/ruby/ext/google/protobuf_c/protobuf.h
@@ -190,7 +190,7 @@ VALUE Descriptor_each(VALUE _self);
 VALUE Descriptor_lookup(VALUE _self, VALUE name);
 VALUE Descriptor_add_field(VALUE _self, VALUE obj);
 VALUE Descriptor_add_oneof(VALUE _self, VALUE obj);
-VALUE Descriptor_oneofs(VALUE _self);
+VALUE Descriptor_each_oneof(VALUE _self);
 VALUE Descriptor_lookup_oneof(VALUE _self, VALUE name);
 VALUE Descriptor_msgclass(VALUE _self);
 extern const rb_data_type_t _Descriptor_type;
diff --git a/ruby/ext/google/protobuf_c/storage.c b/ruby/ext/google/protobuf_c/storage.c
index d526679a..adb74394 100644
--- a/ruby/ext/google/protobuf_c/storage.c
+++ b/ruby/ext/google/protobuf_c/storage.c
@@ -473,13 +473,26 @@ VALUE field_type_class(const upb_fielddef* field) {
   return type_class;
 }
 
+static void* slot_memory(MessageLayout* layout,
+                         const void* storage,
+                         const upb_fielddef* field) {
+  return ((uint8_t *)storage) +
+      layout->fields[upb_fielddef_index(field)].offset;
+}
+
+static uint32_t* slot_oneof_case(MessageLayout* layout,
+                                 const void* storage,
+                                 const upb_fielddef* field) {
+  return (uint32_t *)(((uint8_t *)storage) +
+      layout->fields[upb_fielddef_index(field)].case_offset);
+}
+
+
 VALUE layout_get(MessageLayout* layout,
                  const void* storage,
                  const upb_fielddef* field) {
-  void* memory = ((uint8_t *)storage) +
-      layout->fields[upb_fielddef_index(field)].offset;
-  uint32_t* oneof_case = (uint32_t *)(((uint8_t *)storage) +
-      layout->fields[upb_fielddef_index(field)].case_offset);
+  void* memory = slot_memory(layout, storage, field);
+  uint32_t* oneof_case = slot_oneof_case(layout, storage, field);
 
   if (upb_fielddef_containingoneof(field)) {
     if (*oneof_case != upb_fielddef_number(field)) {
@@ -552,10 +565,8 @@ void layout_set(MessageLayout* layout,
                 void* storage,
                 const upb_fielddef* field,
                 VALUE val) {
-  void* memory = ((uint8_t *)storage) +
-      layout->fields[upb_fielddef_index(field)].offset;
-  uint32_t* oneof_case = (uint32_t *)(((uint8_t *)storage) +
-      layout->fields[upb_fielddef_index(field)].case_offset);
+  void* memory = slot_memory(layout, storage, field);
+  uint32_t* oneof_case = slot_oneof_case(layout, storage, field);
 
   if (upb_fielddef_containingoneof(field)) {
     if (val == Qnil) {
@@ -592,10 +603,8 @@ void layout_init(MessageLayout* layout,
        !upb_msg_field_done(&it);
        upb_msg_field_next(&it)) {
     const upb_fielddef* field = upb_msg_iter_field(&it);
-    void* memory = ((uint8_t *)storage) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* oneof_case = (uint32_t *)(((uint8_t *)storage) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
+    void* memory = slot_memory(layout, storage, field);
+    uint32_t* oneof_case = slot_oneof_case(layout, storage, field);
 
     if (upb_fielddef_containingoneof(field)) {
       memset(memory, 0, NATIVE_SLOT_MAX_SIZE);
@@ -652,10 +661,8 @@ void layout_mark(MessageLayout* layout, void* storage) {
        !upb_msg_field_done(&it);
        upb_msg_field_next(&it)) {
     const upb_fielddef* field = upb_msg_iter_field(&it);
-    void* memory = ((uint8_t *)storage) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* oneof_case = (uint32_t *)(((uint8_t *)storage) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
+    void* memory = slot_memory(layout, storage, field);
+    uint32_t* oneof_case = slot_oneof_case(layout, storage, field);
 
     if (upb_fielddef_containingoneof(field)) {
       if (*oneof_case == upb_fielddef_number(field)) {
@@ -675,14 +682,11 @@ void layout_dup(MessageLayout* layout, void* to, void* from) {
        !upb_msg_field_done(&it);
        upb_msg_field_next(&it)) {
     const upb_fielddef* field = upb_msg_iter_field(&it);
-    void* to_memory = ((uint8_t *)to) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* to_oneof_case = (uint32_t *)(((uint8_t *)to) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
-    void* from_memory = ((uint8_t *)from) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* from_oneof_case = (uint32_t *)(((uint8_t *)from) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
+
+    void* to_memory = slot_memory(layout, to, field);
+    uint32_t* to_oneof_case = slot_oneof_case(layout, to, field);
+    void* from_memory = slot_memory(layout, from, field);
+    uint32_t* from_oneof_case = slot_oneof_case(layout, from, field);
 
     if (upb_fielddef_containingoneof(field)) {
       if (*from_oneof_case == upb_fielddef_number(field)) {
@@ -705,14 +709,11 @@ void layout_deep_copy(MessageLayout* layout, void* to, void* from) {
        !upb_msg_field_done(&it);
        upb_msg_field_next(&it)) {
     const upb_fielddef* field = upb_msg_iter_field(&it);
-    void* to_memory = ((uint8_t *)to) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* to_oneof_case = (uint32_t *)(((uint8_t *)to) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
-    void* from_memory = ((uint8_t *)from) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* from_oneof_case = (uint32_t *)(((uint8_t *)from) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
+
+    void* to_memory = slot_memory(layout, to, field);
+    uint32_t* to_oneof_case = slot_oneof_case(layout, to, field);
+    void* from_memory = slot_memory(layout, from, field);
+    uint32_t* from_oneof_case = slot_oneof_case(layout, from, field);
 
     if (upb_fielddef_containingoneof(field)) {
       if (*from_oneof_case == upb_fielddef_number(field)) {
@@ -737,14 +738,11 @@ VALUE layout_eq(MessageLayout* layout, void* msg1, void* msg2) {
        !upb_msg_field_done(&it);
        upb_msg_field_next(&it)) {
     const upb_fielddef* field = upb_msg_iter_field(&it);
-    void* msg1_memory = ((uint8_t *)msg1) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* msg1_oneof_case = (uint32_t *)(((uint8_t *)msg1) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
-    void* msg2_memory = ((uint8_t *)msg2) +
-        layout->fields[upb_fielddef_index(field)].offset;
-    uint32_t* msg2_oneof_case = (uint32_t *)(((uint8_t *)msg2) +
-        layout->fields[upb_fielddef_index(field)].case_offset);
+
+    void* msg1_memory = slot_memory(layout, msg1, field);
+    uint32_t* msg1_oneof_case = slot_oneof_case(layout, msg1, field);
+    void* msg2_memory = slot_memory(layout, msg2, field);
+    uint32_t* msg2_oneof_case = slot_oneof_case(layout, msg2, field);
 
     if (upb_fielddef_containingoneof(field)) {
       if (*msg1_oneof_case != *msg2_oneof_case ||