Addressed code-review comments.

6 files changed, 59 insertions, 32 deletions

diff --git a/ruby/ext/google/protobuf_c/defs.c b/ruby/ext/google/protobuf_c/defs.c
index 5c94a74a..499e041b 100644
--- a/ruby/ext/google/protobuf_c/defs.c
+++ b/ruby/ext/google/protobuf_c/defs.c
@@ -1077,11 +1077,12 @@ VALUE MessageBuilderContext_repeated(int argc, VALUE* argv, VALUE _self) {
  *     MessageBuilderContext.map(name, key_type, value_type, number,
  *                               value_type_class = nil)
  *
- * Defines a new map field on this message type with the given key and value types, tag
- * number, and type class (for message and enum value types). The key type must
- * be :int32/:uint32/:int64/:uint64, :bool, or :string. The value type type must
- * be a Ruby symbol (as accepted by FieldDescriptor#type=) and the type_class
- * must be a string, if present (as accepted by FieldDescriptor#submsg_name=).
+ * Defines a new map field on this message type with the given key and value
+ * types, tag number, and type class (for message and enum value types). The key
+ * type must be :int32/:uint32/:int64/:uint64, :bool, or :string. The value type
+ * type must be a Ruby symbol (as accepted by FieldDescriptor#type=) and the
+ * type_class must be a string, if present (as accepted by
+ * FieldDescriptor#submsg_name=).
  */
 VALUE MessageBuilderContext_map(int argc, VALUE* argv, VALUE _self) {
   DEFINE_SELF(MessageBuilderContext, self, _self);
diff --git a/ruby/ext/google/protobuf_c/encode_decode.c b/ruby/ext/google/protobuf_c/encode_decode.c
index 6263edcc..f1b951fc 100644
--- a/ruby/ext/google/protobuf_c/encode_decode.c
+++ b/ruby/ext/google/protobuf_c/encode_decode.c
@@ -177,8 +177,8 @@ static void *submsg_handler(void *closure, const void *hd) {
 // Handler data for startmap/endmap handlers.
 typedef struct {
   size_t ofs;
-  const upb_fielddef* key_field;
-  const upb_fielddef* value_field;
+  upb_fieldtype_t key_field_type;
+  upb_fieldtype_t value_field_type;
   VALUE value_field_typeclass;
 } map_handlerdata_t;
 
@@ -194,12 +194,6 @@ typedef struct {
   char value_storage[NATIVE_SLOT_MAX_SIZE];
 } map_parse_frame_t;
 
-// Handler to begin a sequence of map entries: simple no-op that exists only to
-// set context for the map entry handlers.
-static void *startmap_handler(void *closure, const void *hd) {
-  return closure;
-}
-
 // Handler to begin a map entry: allocates a temporary frame. This is the
 // 'startsubmsg' handler on the msgdef that contains the map field.
 static void *startmapentry_handler(void *closure, const void *hd) {
@@ -210,10 +204,8 @@ static void *startmapentry_handler(void *closure, const void *hd) {
   map_parse_frame_t* frame = ALLOC(map_parse_frame_t);
   frame->map = map_rb;
 
-  native_slot_init(upb_fielddef_type(mapdata->key_field),
-                   &frame->key_storage);
-  native_slot_init(upb_fielddef_type(mapdata->value_field),
-                   &frame->value_storage);
+  native_slot_init(mapdata->key_field_type, &frame->key_storage);
+  native_slot_init(mapdata->value_field_type, &frame->value_storage);
 
   return frame;
 }
@@ -225,10 +217,10 @@ static bool endmap_handler(void *closure, const void *hd, upb_status* s) {
   const map_handlerdata_t* mapdata = hd;
 
   VALUE key = native_slot_get(
-      upb_fielddef_type(mapdata->key_field), Qnil,
+      mapdata->key_field_type, Qnil,
       &frame->key_storage);
   VALUE value = native_slot_get(
-      upb_fielddef_type(mapdata->value_field), mapdata->value_field_typeclass,
+      mapdata->value_field_type, mapdata->value_field_typeclass,
       &frame->value_storage);
 
   Map_index_set(frame->map, key, value);
@@ -250,11 +242,15 @@ static map_handlerdata_t* new_map_handlerdata(
 
   map_handlerdata_t* hd = ALLOC(map_handlerdata_t);
   hd->ofs = ofs;
-  hd->key_field = upb_msgdef_itof(mapentry_def, 1);
-  assert(hd->key_field != NULL);
-  hd->value_field = upb_msgdef_itof(mapentry_def, 2);
-  assert(hd->value_field != NULL);
-  hd->value_field_typeclass = field_type_class(hd->value_field);
+  const upb_fielddef* key_field = upb_msgdef_itof(mapentry_def,
+                                                  MAP_KEY_FIELD);
+  assert(key_field != NULL);
+  hd->key_field_type = upb_fielddef_type(key_field);
+  const upb_fielddef* value_field = upb_msgdef_itof(mapentry_def,
+                                                    MAP_VALUE_FIELD);
+  assert(value_field != NULL);
+  hd->value_field_type = upb_fielddef_type(value_field);
+  hd->value_field_typeclass = field_type_class(value_field);
 
   return hd;
 }
@@ -293,6 +289,7 @@ static void add_handlers_for_repeated_field(upb_handlers *h,
                                appendbytes_handler : appendstr_handler,
                                NULL);
       upb_handlers_setstring(h, f, stringdata_handler, NULL);
+      break;
     }
     case UPB_TYPE_MESSAGE: {
       upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
@@ -352,7 +349,6 @@ static void add_handlers_for_mapfield(upb_handlers* h,
   upb_handlers_addcleanup(h, hd, free);
   upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
   upb_handlerattr_sethandlerdata(&attr, hd);
-  upb_handlers_setstartseq(h, fielddef, startmap_handler, &attr);
   upb_handlers_setstartsubmsg(h, fielddef, startmapentry_handler, &attr);
   upb_handlerattr_uninit(&attr);
 }
@@ -360,6 +356,8 @@ 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) {
+  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);
   upb_handlers_addcleanup(h, hd, free);
   upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
@@ -367,7 +365,7 @@ static void add_handlers_for_mapentry(const upb_msgdef* msgdef,
   upb_handlers_setendmsg(h, endmap_handler, &attr);
 
   add_handlers_for_singular_field(
-      h, hd->key_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
@@ -375,7 +373,7 @@ static void add_handlers_for_mapentry(const upb_msgdef* msgdef,
       // we compensate for that addition.
       offsetof(map_parse_frame_t, key_storage) - sizeof(MessageHeader));
   add_handlers_for_singular_field(
-      h, hd->value_field,
+      h, value_field,
       offsetof(map_parse_frame_t, value_storage) - sizeof(MessageHeader));
 }
 
diff --git a/ruby/ext/google/protobuf_c/map.c b/ruby/ext/google/protobuf_c/map.c
index 9d5de9cd..56548799 100644
--- a/ruby/ext/google/protobuf_c/map.c
+++ b/ruby/ext/google/protobuf_c/map.c
@@ -32,6 +32,17 @@
 
 // -----------------------------------------------------------------------------
 // Basic map operations on top of upb's strtable.
+//
+// Note that we roll our own `Map` container here because, as for
+// `RepeatedField`, we want a strongly-typed container. This is so that any user
+// errors due to incorrect map key or value types are raised as close as
+// possible to the error site, rather than at some deferred point (e.g.,
+// serialization).
+//
+// We build our `Map` on top of upb_strtable so that we're able to take
+// advantage of the native_slot storage abstraction, as RepeatedField does.
+// (This is not quite a perfect mapping -- see the key conversions below -- but
+// gives us full support and error-checking for all value types for free.)
 // -----------------------------------------------------------------------------
 
 // Map values are stored using the native_slot abstraction (as with repeated
diff --git a/ruby/ext/google/protobuf_c/message.c b/ruby/ext/google/protobuf_c/message.c
index de38dd7b..ee8881d4 100644
--- a/ruby/ext/google/protobuf_c/message.c
+++ b/ruby/ext/google/protobuf_c/message.c
@@ -142,7 +142,7 @@ int Message_initialize_kwarg(VALUE key, VALUE val, VALUE _self) {
   if (is_map_field(f)) {
     if (TYPE(val) != T_HASH) {
       rb_raise(rb_eArgError,
-               "Expected hashmap as initializer value for map field.");
+               "Expected Hash object as initializer value for map field.");
     }
     VALUE map = layout_get(self->descriptor->layout, Message_data(self), f);
     Map_merge_into_self(map, val);
diff --git a/ruby/ext/google/protobuf_c/protobuf.h b/ruby/ext/google/protobuf_c/protobuf.h
index 91a97a68..6dac6029 100644
--- a/ruby/ext/google/protobuf_c/protobuf.h
+++ b/ruby/ext/google/protobuf_c/protobuf.h
@@ -268,10 +268,19 @@ extern rb_encoding* kRubyString8bitEncoding;
 
 VALUE field_type_class(const upb_fielddef* field);
 
+#define MAP_KEY_FIELD 1
+#define MAP_VALUE_FIELD 2
+
+// These operate on a map field (i.e., a repeated field of submessages whose
+// submessage type is a map-entry msgdef).
 bool is_map_field(const upb_fielddef* field);
 const upb_fielddef* map_field_key(const upb_fielddef* field);
 const upb_fielddef* map_field_value(const upb_fielddef* field);
 
+// These operate on a map-entry msgdef.
+const upb_fielddef* map_entry_key(const upb_msgdef* msgdef);
+const upb_fielddef* map_entry_value(const upb_msgdef* msgdef);
+
 // -----------------------------------------------------------------------------
 // Repeated field container type.
 // -----------------------------------------------------------------------------
diff --git a/ruby/ext/google/protobuf_c/storage.c b/ruby/ext/google/protobuf_c/storage.c
index f20ddec2..235fbff2 100644
--- a/ruby/ext/google/protobuf_c/storage.c
+++ b/ruby/ext/google/protobuf_c/storage.c
@@ -339,15 +339,23 @@ bool is_map_field(const upb_fielddef* field) {
 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_fielddef* key_field = upb_msgdef_itof(subdef, 1);
-  assert(key_field != NULL);
-  return key_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_fielddef* value_field = upb_msgdef_itof(subdef, 2);
+  return map_entry_value(subdef);
+}
+
+const upb_fielddef* map_entry_key(const upb_msgdef* msgdef) {
+  const upb_fielddef* key_field = upb_msgdef_itof(msgdef, MAP_KEY_FIELD);
+  assert(key_field != NULL);
+  return key_field;
+}
+
+const upb_fielddef* map_entry_value(const upb_msgdef* msgdef) {
+  const upb_fielddef* value_field = upb_msgdef_itof(msgdef, MAP_VALUE_FIELD);
   assert(value_field != NULL);
   return value_field;
 }