Support for maps in the MRI C Ruby extension.

This adds the Map container and support for parsing and serializing maps
in the protobuf wire format (as defined by the C++ implementation, with
MapEntry submessages in a repeated field). JSON map
serialization/parsing are not yet supported as these will require some
changes to upb as well.

ruby/ext/google/protobuf_c/defs.c

@@ -923,6 +923,7 @@ DEFINE_CLASS(MessageBuilderContext,
 void MessageBuilderContext_mark(void* _self) {
   MessageBuilderContext* self = _self;
   rb_gc_mark(self->descriptor);
+  rb_gc_mark(self->builder);
 }
 
 void MessageBuilderContext_free(void* _self) {
@@ -935,6 +936,7 @@ VALUE MessageBuilderContext_alloc(VALUE klass) {
   VALUE ret = TypedData_Wrap_Struct(
       klass, &_MessageBuilderContext_type, self);
   self->descriptor = Qnil;
+  self->builder = Qnil;
   return ret;
 }
 
@@ -943,24 +945,29 @@ void MessageBuilderContext_register(VALUE module) {
       module, "MessageBuilderContext", rb_cObject);
   rb_define_alloc_func(klass, MessageBuilderContext_alloc);
   rb_define_method(klass, "initialize",
-                   MessageBuilderContext_initialize, 1);
+                   MessageBuilderContext_initialize, 2);
   rb_define_method(klass, "optional", MessageBuilderContext_optional, -1);
   rb_define_method(klass, "required", MessageBuilderContext_required, -1);
   rb_define_method(klass, "repeated", MessageBuilderContext_repeated, -1);
+  rb_define_method(klass, "map", MessageBuilderContext_map, -1);
   cMessageBuilderContext = klass;
   rb_gc_register_address(&cMessageBuilderContext);
 }
 
 /*
  * call-seq:
- *     MessageBuilderContext.new(desc) => context
+ *     MessageBuilderContext.new(desc, builder) => context
  *
- * Create a new builder context around the given message descriptor. This class
- * is intended to serve as a DSL context to be used with #instance_eval.
+ * Create a new message builder context around the given message descriptor and
+ * builder context. This class is intended to serve as a DSL context to be used
+ * with #instance_eval.
  */
-VALUE MessageBuilderContext_initialize(VALUE _self, VALUE msgdef) {
+VALUE MessageBuilderContext_initialize(VALUE _self,
+                                       VALUE msgdef,
+                                       VALUE builder) {
   DEFINE_SELF(MessageBuilderContext, self, _self);
   self->descriptor = msgdef;
+  self->builder = builder;
   return Qnil;
 }
 
@@ -1065,6 +1072,96 @@ VALUE MessageBuilderContext_repeated(int argc, VALUE* argv, VALUE _self) {
                           name, type, number, type_class);
 }
 
+/*
+ * call-seq:
+ *     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=).
+ */
+VALUE MessageBuilderContext_map(int argc, VALUE* argv, VALUE _self) {
+  DEFINE_SELF(MessageBuilderContext, self, _self);
+
+  if (argc < 4) {
+    rb_raise(rb_eArgError, "Expected at least 4 arguments.");
+  }
+  VALUE name = argv[0];
+  VALUE key_type = argv[1];
+  VALUE value_type = argv[2];
+  VALUE number = argv[3];
+  VALUE type_class = (argc > 4) ? argv[4] : Qnil;
+
+  // Validate the key type. We can't accept enums, messages, or floats/doubles
+  // as map keys. (We exclude these explicitly, and the field-descriptor setter
+  // below then ensures that the type is one of the remaining valid options.)
+  if (SYM2ID(key_type) == rb_intern("float") ||
+      SYM2ID(key_type) == rb_intern("double") ||
+      SYM2ID(key_type) == rb_intern("enum") ||
+      SYM2ID(key_type) == rb_intern("message")) {
+    rb_raise(rb_eArgError,
+             "Cannot add a map field with a float, double, enum, or message "
+             "type.");
+  }
+
+  // Create a new message descriptor for the map entry message, and create a
+  // repeated submessage field here with that type.
+  VALUE mapentry_desc = rb_class_new_instance(0, NULL, cDescriptor);
+  VALUE mapentry_desc_name = rb_funcall(self->descriptor, rb_intern("name"), 0);
+  mapentry_desc_name = rb_str_cat2(mapentry_desc_name, "_MapEntry_");
+  mapentry_desc_name = rb_str_cat2(mapentry_desc_name,
+                                   rb_id2name(SYM2ID(name)));
+  Descriptor_name_set(mapentry_desc, mapentry_desc_name);
+
+  // The 'mapentry' attribute has no Ruby setter because we do not want the user
+  // attempting to DIY the setup below; we want to ensure that the fields are
+  // correct. So we reach into the msgdef here to set the bit manually.
+  Descriptor* mapentry_desc_self = ruby_to_Descriptor(mapentry_desc);
+  upb_msgdef_setmapentry((upb_msgdef*)mapentry_desc_self->msgdef, true);
+
+  // optional <type> key = 1;
+  VALUE key_field = rb_class_new_instance(0, NULL, cFieldDescriptor);
+  FieldDescriptor_name_set(key_field, rb_str_new2("key"));
+  FieldDescriptor_label_set(key_field, ID2SYM(rb_intern("optional")));
+  FieldDescriptor_number_set(key_field, INT2NUM(1));
+  FieldDescriptor_type_set(key_field, key_type);
+  Descriptor_add_field(mapentry_desc, key_field);
+
+  // optional <type> value = 2;
+  VALUE value_field = rb_class_new_instance(0, NULL, cFieldDescriptor);
+  FieldDescriptor_name_set(value_field, rb_str_new2("value"));
+  FieldDescriptor_label_set(value_field, ID2SYM(rb_intern("optional")));
+  FieldDescriptor_number_set(value_field, INT2NUM(2));
+  FieldDescriptor_type_set(value_field, value_type);
+  if (type_class != Qnil) {
+    VALUE submsg_name = rb_str_new2("."); // prepend '.' to make name absolute.
+    submsg_name = rb_str_append(submsg_name, type_class);
+    FieldDescriptor_submsg_name_set(value_field, submsg_name);
+  }
+  Descriptor_add_field(mapentry_desc, value_field);
+
+  // Add the map-entry message type to the current builder, and use the type to
+  // create the map field itself.
+  Builder* builder_self = ruby_to_Builder(self->builder);
+  rb_ary_push(builder_self->pending_list, mapentry_desc);
+
+  VALUE map_field = rb_class_new_instance(0, NULL, cFieldDescriptor);
+  VALUE name_str = rb_str_new2(rb_id2name(SYM2ID(name)));
+  FieldDescriptor_name_set(map_field, name_str);
+  FieldDescriptor_number_set(map_field, number);
+  FieldDescriptor_label_set(map_field, ID2SYM(rb_intern("repeated")));
+  FieldDescriptor_type_set(map_field, ID2SYM(rb_intern("message")));
+  VALUE submsg_name = rb_str_new2("."); // prepend '.' to make name absolute.
+  submsg_name = rb_str_append(submsg_name, mapentry_desc_name);
+  FieldDescriptor_submsg_name_set(map_field, submsg_name);
+  Descriptor_add_field(self->descriptor, map_field);
+
+  return Qnil;
+}
+
 // -----------------------------------------------------------------------------
 // EnumBuilderContext.
 // -----------------------------------------------------------------------------
@@ -1190,7 +1287,8 @@ void Builder_register(VALUE module) {
 VALUE Builder_add_message(VALUE _self, VALUE name) {
   DEFINE_SELF(Builder, self, _self);
   VALUE msgdef = rb_class_new_instance(0, NULL, cDescriptor);
-  VALUE ctx = rb_class_new_instance(1, &msgdef, cMessageBuilderContext);
+  VALUE args[2] = { msgdef, _self };
+  VALUE ctx = rb_class_new_instance(2, args, cMessageBuilderContext);
   VALUE block = rb_block_proc();
   rb_funcall(msgdef, rb_intern("name="), 1, name);
   rb_funcall_with_block(ctx, rb_intern("instance_eval"), 0, NULL, block);

ruby/ext/google/protobuf_c/extconf.rb

@@ -5,6 +5,6 @@ require 'mkmf'
 $CFLAGS += " -O3 -std=c99 -Wno-unused-function -DNDEBUG "
 
 $objs = ["protobuf.o", "defs.o", "storage.o", "message.o",
-         "repeated_field.o", "encode_decode.o", "upb.o"]
+         "repeated_field.o", "map.o", "encode_decode.o", "upb.o"]
 
 create_makefile("google/protobuf_c")

ruby/ext/google/protobuf_c/message.c

@@ -139,7 +139,14 @@ int Message_initialize_kwarg(VALUE key, VALUE val, VALUE _self) {
              "Unknown field name in initialization map entry.");
   }
 
-  if (upb_fielddef_label(f) == UPB_LABEL_REPEATED) {
+  if (is_map_field(f)) {
+    if (TYPE(val) != T_HASH) {
+      rb_raise(rb_eArgError,
+               "Expected hashmap as initializer value for map field.");
+    }
+    VALUE map = layout_get(self->descriptor->layout, Message_data(self), f);
+    Map_merge_into_self(map, val);
+  } else if (upb_fielddef_label(f) == UPB_LABEL_REPEATED) {
     if (TYPE(val) != T_ARRAY) {
       rb_raise(rb_eArgError,
                "Expected array as initializer value for repeated field.");
@@ -450,13 +457,15 @@ VALUE build_module_from_enumdesc(EnumDescriptor* enumdesc) {
  * call-seq:
  *     Google::Protobuf.deep_copy(obj) => copy_of_obj
  *
- * Performs a deep copy of either a RepeatedField instance or a message object,
- * recursively copying its members.
+ * Performs a deep copy of a RepeatedField instance, a Map instance, or a
+ * message object, recursively copying its members.
  */
 VALUE Google_Protobuf_deep_copy(VALUE self, VALUE obj) {
   VALUE klass = CLASS_OF(obj);
   if (klass == cRepeatedField) {
     return RepeatedField_deep_copy(obj);
+  } else if (klass == cMap) {
+    return Map_deep_copy(obj);
   } else {
     return Message_deep_copy(obj);
   }

ruby/ext/google/protobuf_c/protobuf.c

@@ -82,6 +82,7 @@ void Init_protobuf_c() {
   EnumBuilderContext_register(internal);
   Builder_register(internal);
   RepeatedField_register(protobuf);
+  Map_register(protobuf);
 
   rb_define_singleton_method(protobuf, "encode", Google_Protobuf_encode, 1);
   rb_define_singleton_method(protobuf, "decode", Google_Protobuf_decode, 2);

ruby/ext/google/protobuf_c/protobuf.h

@@ -123,6 +123,7 @@ struct EnumDescriptor {
 
 struct MessageBuilderContext {
   VALUE descriptor;
+  VALUE builder;
 };
 
 struct EnumBuilderContext {
@@ -213,10 +214,13 @@ void MessageBuilderContext_free(void* _self);
 VALUE MessageBuilderContext_alloc(VALUE klass);
 void MessageBuilderContext_register(VALUE module);
 MessageBuilderContext* ruby_to_MessageBuilderContext(VALUE value);
-VALUE MessageBuilderContext_initialize(VALUE _self, VALUE descriptor);
+VALUE MessageBuilderContext_initialize(VALUE _self,
+                                       VALUE descriptor,
+                                       VALUE builder);
 VALUE MessageBuilderContext_optional(int argc, VALUE* argv, VALUE _self);
 VALUE MessageBuilderContext_required(int argc, VALUE* argv, VALUE _self);
 VALUE MessageBuilderContext_repeated(int argc, VALUE* argv, VALUE _self);
+VALUE MessageBuilderContext_map(int argc, VALUE* argv, VALUE _self);
 
 void EnumBuilderContext_mark(void* _self);
 void EnumBuilderContext_free(void* _self);
@@ -239,6 +243,8 @@ VALUE Builder_finalize_to_pool(VALUE _self, VALUE pool_rb);
 // Native slot storage abstraction.
 // -----------------------------------------------------------------------------
 
+#define NATIVE_SLOT_MAX_SIZE sizeof(void*)
+
 size_t native_slot_size(upb_fieldtype_t type);
 void native_slot_set(upb_fieldtype_t type,
                      VALUE type_class,
@@ -254,11 +260,18 @@ void native_slot_deep_copy(upb_fieldtype_t type, void* to, void* from);
 bool native_slot_eq(upb_fieldtype_t type, void* mem1, void* mem2);
 
 void native_slot_validate_string_encoding(upb_fieldtype_t type, VALUE value);
+void native_slot_check_int_range_precision(upb_fieldtype_t type, VALUE value);
 
 extern rb_encoding* kRubyStringUtf8Encoding;
 extern rb_encoding* kRubyStringASCIIEncoding;
 extern rb_encoding* kRubyString8bitEncoding;
 
+VALUE field_type_class(const upb_fielddef* field);
+
+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);
+
 // -----------------------------------------------------------------------------
 // Repeated field container type.
 // -----------------------------------------------------------------------------
@@ -282,7 +295,6 @@ extern VALUE cRepeatedField;
 
 RepeatedField* ruby_to_RepeatedField(VALUE value);
 
-void RepeatedField_register(VALUE module);
 VALUE RepeatedField_each(VALUE _self);
 VALUE RepeatedField_index(VALUE _self, VALUE _index);
 void* RepeatedField_index_native(VALUE _self, int index);
@@ -302,6 +314,59 @@ VALUE RepeatedField_hash(VALUE _self);
 VALUE RepeatedField_inspect(VALUE _self);
 VALUE RepeatedField_plus(VALUE _self, VALUE list);
 
+// Defined in repeated_field.c; also used by Map.
+void validate_type_class(upb_fieldtype_t type, VALUE klass);
+
+// -----------------------------------------------------------------------------
+// Map container type.
+// -----------------------------------------------------------------------------
+
+typedef struct {
+  upb_fieldtype_t key_type;
+  upb_fieldtype_t value_type;
+  VALUE value_type_class;
+  upb_strtable table;
+} Map;
+
+void Map_mark(void* self);
+void Map_free(void* self);
+VALUE Map_alloc(VALUE klass);
+VALUE Map_init(int argc, VALUE* argv, VALUE self);
+void Map_register(VALUE module);
+
+extern const rb_data_type_t Map_type;
+extern VALUE cMap;
+
+Map* ruby_to_Map(VALUE value);
+
+VALUE Map_each(VALUE _self);
+VALUE Map_keys(VALUE _self);
+VALUE Map_values(VALUE _self);
+VALUE Map_index(VALUE _self, VALUE key);
+VALUE Map_index_set(VALUE _self, VALUE key, VALUE value);
+VALUE Map_has_key(VALUE _self, VALUE key);
+VALUE Map_delete(VALUE _self, VALUE key);
+VALUE Map_clear(VALUE _self);
+VALUE Map_length(VALUE _self);
+VALUE Map_dup(VALUE _self);
+VALUE Map_deep_copy(VALUE _self);
+VALUE Map_eq(VALUE _self, VALUE _other);
+VALUE Map_hash(VALUE _self);
+VALUE Map_inspect(VALUE _self);
+VALUE Map_merge(VALUE _self, VALUE hashmap);
+VALUE Map_merge_into_self(VALUE _self, VALUE hashmap);
+
+typedef struct {
+  Map* self;
+  upb_strtable_iter it;
+} Map_iter;
+
+void Map_begin(VALUE _self, Map_iter* iter);
+void Map_next(Map_iter* iter);
+bool Map_done(Map_iter* iter);
+VALUE Map_iter_key(Map_iter* iter);
+VALUE Map_iter_value(Map_iter* iter);
+
 // -----------------------------------------------------------------------------
 // Message layout / storage.
 // -----------------------------------------------------------------------------

ruby/ext/google/protobuf_c/repeated_field.c

@@ -324,6 +324,10 @@ VALUE RepeatedField_deep_copy(VALUE _self) {
  * element types are equal, their lengths are equal, and each element is equal.
  * Elements are compared as per normal Ruby semantics, by calling their :==
  * methods (or performing a more efficient comparison for primitive types).
+ *
+ * Repeated fields with dissimilar element types are never equal, even if value
+ * comparison (for example, between integers and floats) would have otherwise
+ * indicated that every element has equal value.
  */
 VALUE RepeatedField_eq(VALUE _self, VALUE _other) {
   if (_self == _other) {
@@ -458,7 +462,7 @@ VALUE RepeatedField_plus(VALUE _self, VALUE list) {
   return dupped;
 }
 
-static void validate_type_class(upb_fieldtype_t type, VALUE klass) {
+void validate_type_class(upb_fieldtype_t type, VALUE klass) {
   if (rb_iv_get(klass, kDescriptorInstanceVar) == Qnil) {
     rb_raise(rb_eArgError,
              "Type class has no descriptor. Please pass a "

ruby/ext/google/protobuf_c/upb.c

@@ -1269,6 +1269,7 @@ upb_msgdef *upb_msgdef_new(const void *owner) {
   if (!upb_def_init(UPB_UPCAST(m), UPB_DEF_MSG, &vtbl, owner)) goto err2;
   if (!upb_inttable_init(&m->itof, UPB_CTYPE_PTR)) goto err2;
   if (!upb_strtable_init(&m->ntof, UPB_CTYPE_PTR)) goto err1;
+  m->map_entry = false;
   return m;
 
 err1:
@@ -1283,6 +1284,7 @@ upb_msgdef *upb_msgdef_dup(const upb_msgdef *m, const void *owner) {
   if (!newm) return NULL;
   bool ok = upb_def_setfullname(UPB_UPCAST(newm),
                                 upb_def_fullname(UPB_UPCAST(m)), NULL);
+  newm->map_entry = m->map_entry;
   UPB_ASSERT_VAR(ok, ok);
   upb_msg_iter i;
   for(upb_msg_begin(&i, m); !upb_msg_done(&i); upb_msg_next(&i)) {
@@ -1386,6 +1388,15 @@ int upb_msgdef_numfields(const upb_msgdef *m) {
   return upb_strtable_count(&m->ntof);
 }
 
+void upb_msgdef_setmapentry(upb_msgdef *m, bool map_entry) {
+  assert(!upb_msgdef_isfrozen(m));
+  m->map_entry = map_entry;
+}
+
+bool upb_msgdef_mapentry(const upb_msgdef *m) {
+  return m->map_entry;
+}
+
 void upb_msg_begin(upb_msg_iter *iter, const upb_msgdef *m) {
   upb_inttable_begin(iter, &m->itof);
 }
@@ -3401,31 +3412,28 @@ int log2ceil(uint64_t v) {
 }
 
 char *upb_strdup(const char *s) {
-  size_t n = strlen(s) + 1;
+  return upb_strdup2(s, strlen(s));
+}
+
+char *upb_strdup2(const char *s, size_t len) {
+  // Always null-terminate, even if binary data; but don't rely on the input to
+  // 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, n);
+  if (p) memcpy(p, s, len);
+  p[len] = 0;
   return p;
 }
 
 // A type to represent the lookup key of either a strtable or an inttable.
-// This is like upb_tabkey, but can carry a size also to allow lookups of
-// non-NULL-terminated strings (we don't store string lengths in the table).
 typedef struct {
   upb_tabkey key;
-  uint32_t len;  // For string keys only.
 } lookupkey_t;
 
-static lookupkey_t strkey(const char *str) {
-  lookupkey_t k;
-  k.key.str = (char*)str;
-  k.len = strlen(str);
-  return k;
-}
-
 static lookupkey_t strkey2(const char *str, size_t len) {
   lookupkey_t k;
-  k.key.str = (char*)str;
-  k.len = len;
+  k.key.s.str = (char*)str;
+  k.key.s.length = len;
   return k;
 }
 
@@ -3607,11 +3615,12 @@ static size_t begin(const upb_table *t) {
 // A simple "subclass" of upb_table that only adds a hash function for strings.
 
 static uint32_t strhash(upb_tabkey key) {
-  return MurmurHash2(key.str, strlen(key.str), 0);
+  return MurmurHash2(key.s.str, key.s.length, 0);
 }
 
 static bool streql(upb_tabkey k1, lookupkey_t k2) {
-  return strncmp(k1.str, k2.key.str, k2.len) == 0 && k1.str[k2.len] == '\0';
+  return k1.s.length == k2.key.s.length &&
+         memcmp(k1.s.str, k2.key.s.str, k1.s.length) == 0;
 }
 
 bool upb_strtable_init(upb_strtable *t, upb_ctype_t ctype) {
@@ -3620,7 +3629,7 @@ bool upb_strtable_init(upb_strtable *t, upb_ctype_t ctype) {
 
 void upb_strtable_uninit(upb_strtable *t) {
   for (size_t i = 0; i < upb_table_size(&t->t); i++)
-    free((void*)t->t.entries[i].key.str);
+    free((void*)t->t.entries[i].key.s.str);
   uninit(&t->t);
 }
 
@@ -3631,26 +3640,30 @@ bool upb_strtable_resize(upb_strtable *t, size_t size_lg2) {
   upb_strtable_iter i;
   upb_strtable_begin(&i, t);
   for ( ; !upb_strtable_done(&i); upb_strtable_next(&i)) {
-    upb_strtable_insert(
-        &new_table, upb_strtable_iter_key(&i), upb_strtable_iter_value(&i));
+    upb_strtable_insert2(
+        &new_table,
+        upb_strtable_iter_key(&i),
+        upb_strtable_iter_keylength(&i),
+        upb_strtable_iter_value(&i));
   }
   upb_strtable_uninit(t);
   *t = new_table;
   return true;
 }
 
-bool upb_strtable_insert(upb_strtable *t, const char *k, upb_value v) {
+bool upb_strtable_insert2(upb_strtable *t, const char *k, size_t len,
+                          upb_value v) {
   if (isfull(&t->t)) {
     // Need to resize.  New table of double the size, add old elements to it.
     if (!upb_strtable_resize(t, t->t.size_lg2 + 1)) {
       return false;
     }
   }
-  if ((k = upb_strdup(k)) == NULL) return false;
+  if ((k = upb_strdup2(k, len)) == NULL) return false;
 
-  lookupkey_t key = strkey(k);
-  uint32_t hash = MurmurHash2(key.key.str, key.len, 0);
-  insert(&t->t, strkey(k), v, hash, &strhash, &streql);
+  lookupkey_t key = strkey2(k, len);
+  uint32_t hash = MurmurHash2(key.key.s.str, key.key.s.length, 0);
+  insert(&t->t, key, v, hash, &strhash, &streql);
   return true;
 }
 
@@ -3660,11 +3673,12 @@ bool upb_strtable_lookup2(const upb_strtable *t, const char *key, size_t len,
   return lookup(&t->t, strkey2(key, len), v, hash, &streql);
 }
 
-bool upb_strtable_remove(upb_strtable *t, const char *key, upb_value *val) {
+bool upb_strtable_remove2(upb_strtable *t, const char *key, size_t len,
+                         upb_value *val) {
   uint32_t hash = MurmurHash2(key, strlen(key), 0);
   upb_tabkey tabkey;
-  if (rm(&t->t, strkey(key), val, &tabkey, hash, &streql)) {
-    free((void*)tabkey.str);
+  if (rm(&t->t, strkey2(key, len), val, &tabkey, hash, &streql)) {
+    free((void*)tabkey.s.str);
     return true;
   } else {
     return false;
@@ -3693,7 +3707,12 @@ bool upb_strtable_done(const upb_strtable_iter *i) {
 
 const char *upb_strtable_iter_key(upb_strtable_iter *i) {
   assert(!upb_strtable_done(i));
-  return str_tabent(i)->key.str;
+  return str_tabent(i)->key.s.str;
+}
+
+size_t upb_strtable_iter_keylength(upb_strtable_iter *i) {
+  assert(!upb_strtable_done(i));
+  return str_tabent(i)->key.s.length;
 }
 
 upb_value upb_strtable_iter_value(const upb_strtable_iter *i) {

ruby/ext/google/protobuf_c/upb.h

@@ -600,6 +600,9 @@ typedef struct {
 
 // Like strdup(), which isn't always available since it's not ANSI C.
 char *upb_strdup(const char *s);
+// Variant that works with a length-delimited rather than NULL-delimited string,
+// as supported by strtable.
+char *upb_strdup2(const char *s, size_t len);
 
 UPB_INLINE void _upb_value_setval(upb_value *v, _upb_value val,
                                   upb_ctype_t ctype) {
@@ -654,12 +657,24 @@ FUNCS(fptr,     fptr,         upb_func*,    UPB_CTYPE_FPTR);
 
 typedef union {
   uintptr_t num;
-  const char *str;  // We own, nullz.
+  struct {
+    // We own this. NULL-terminated but may also contain binary data; see
+    // explicit length below.
+    // TODO: move the length to the start of the string in order to reduce
+    // tabkey's size (to one machine word) in a way that supports static
+    // initialization.
+    const char *str;
+    size_t length;
+  } s;
 } upb_tabkey;
 
 #define UPB_TABKEY_NUM(n) {n}
 #ifdef UPB_C99
-#define UPB_TABKEY_STR(s) {.str = s}
+// Given that |s| is a string literal, sizeof(s) gives us a
+// compile-time-constant strlen(). We must ensure that this works for static
+// data initializers.
+#define UPB_TABKEY_STR(strval) { .s = { .str = strval,                    \
+                                        .length = sizeof(strval) - 1 } }
 #endif
 // TODO(haberman): C++
 #define UPB_TABKEY_NONE {0}
@@ -765,7 +780,14 @@ UPB_INLINE size_t upb_strtable_count(const upb_strtable *t) {
 // If a table resize was required but memory allocation failed, false is
 // returned and the table is unchanged.
 bool upb_inttable_insert(upb_inttable *t, uintptr_t key, upb_value val);
-bool upb_strtable_insert(upb_strtable *t, const char *key, upb_value val);
+bool upb_strtable_insert2(upb_strtable *t, const char *key, size_t len,
+                          upb_value val);
+
+// For NULL-terminated strings.
+UPB_INLINE bool upb_strtable_insert(upb_strtable *t, const char *key,
+                                    upb_value val) {
+  return upb_strtable_insert2(t, key, strlen(key), val);
+}
 
 // Looks up key in this table, returning "true" if the key was found.
 // If v is non-NULL, copies the value for this key into *v.
@@ -782,7 +804,14 @@ UPB_INLINE bool upb_strtable_lookup(const upb_strtable *t, const char *key,
 // Removes an item from the table.  Returns true if the remove was successful,
 // and stores the removed item in *val if non-NULL.
 bool upb_inttable_remove(upb_inttable *t, uintptr_t key, upb_value *val);
-bool upb_strtable_remove(upb_strtable *t, const char *key, upb_value *val);
+bool upb_strtable_remove2(upb_strtable *t, const char *key, size_t len,
+                          upb_value *val);
+
+// For NULL-terminated strings.
+UPB_INLINE bool upb_strtable_remove(upb_strtable *t, const char *key,
+                                    upb_value *v) {
+  return upb_strtable_remove2(t, key, strlen(key), v);
+}
 
 // Updates an existing entry in an inttable.  If the entry does not exist,
 // returns false and does nothing.  Unlike insert/remove, this does not
@@ -876,6 +905,7 @@ void upb_strtable_begin(upb_strtable_iter *i, const upb_strtable *t);
 void upb_strtable_next(upb_strtable_iter *i);
 bool upb_strtable_done(const upb_strtable_iter *i);
 const char *upb_strtable_iter_key(upb_strtable_iter *i);
+size_t upb_strtable_iter_keylength(upb_strtable_iter *i);
 upb_value upb_strtable_iter_value(const upb_strtable_iter *i);
 void upb_strtable_iter_setdone(upb_strtable_iter *i);
 bool upb_strtable_iter_isequal(const upb_strtable_iter *i1,
@@ -1777,6 +1807,10 @@ UPB_DEFINE_DEF(upb::MessageDef, msgdef, MSG, UPB_QUOTE(
   // just be moved into symtab.c?
   MessageDef* Dup(const void* owner) const;
 
+  // Is this message a map entry?
+  void setmapentry(bool map_entry);
+  bool mapentry() const;
+
   // Iteration over fields.  The order is undefined.
   class iterator : public std::iterator<std::forward_iterator_tag, FieldDef*> {
    public:
@@ -1823,6 +1857,11 @@ UPB_DEFINE_STRUCT(upb_msgdef, upb_def,
   upb_inttable itof;  // int to field
   upb_strtable ntof;  // name to field
 
+  // Is this a map-entry message?
+  // TODO: set this flag properly for static descriptors; regenerate
+  // descriptor.upb.c.
+  bool map_entry;
+
   // TODO(haberman): proper extension ranges (there can be multiple).
 ));
 
@@ -1830,7 +1869,7 @@ UPB_DEFINE_STRUCT(upb_msgdef, upb_def,
                         refs, ref2s)                                          \
   {                                                                           \
     UPB_DEF_INIT(name, UPB_DEF_MSG, refs, ref2s), selector_count,             \
-        submsg_field_count, itof, ntof                                        \
+        submsg_field_count, itof, ntof, false                                 \
   }
 
 UPB_BEGIN_EXTERN_C  // {
@@ -1878,6 +1917,9 @@ UPB_INLINE upb_fielddef *upb_msgdef_ntof_mutable(upb_msgdef *m,
   return (upb_fielddef *)upb_msgdef_ntof(m, name, len);
 }
 
+void upb_msgdef_setmapentry(upb_msgdef *m, bool map_entry);
+bool upb_msgdef_mapentry(const upb_msgdef *m);
+
 // upb_msg_iter i;
 // for(upb_msg_begin(&i, m); !upb_msg_done(&i); upb_msg_next(&i)) {
 //   upb_fielddef *f = upb_msg_iter_field(&i);
@@ -2331,6 +2373,12 @@ inline const FieldDef *MessageDef::FindFieldByName(const char *name,
 inline MessageDef* MessageDef::Dup(const void *owner) const {
   return upb_msgdef_dup(this, owner);
 }
+inline void MessageDef::setmapentry(bool map_entry) {
+  upb_msgdef_setmapentry(this, map_entry);
+}
+inline bool MessageDef::mapentry() const {
+  return upb_msgdef_mapentry(this);
+}
 inline MessageDef::iterator MessageDef::begin() { return iterator(this); }
 inline MessageDef::iterator MessageDef::end() { return iterator::end(this); }
 inline MessageDef::const_iterator MessageDef::begin() const {

ruby/tests/basic.rb

@@ -36,23 +36,43 @@ module BasicTest
     add_message "TestMessage2" do
       optional :foo, :int32, 1
     end
+
     add_message "Recursive1" do
       optional :foo, :message, 1, "Recursive2"
     end
     add_message "Recursive2" do
       optional :foo, :message, 1, "Recursive1"
     end
+
     add_enum "TestEnum" do
       value :Default, 0
       value :A, 1
       value :B, 2
       value :C, 3
     end
+
     add_message "BadFieldNames" do
       optional :dup, :int32, 1
       optional :class, :int32, 2
       optional :"a.b", :int32, 3
     end
+
+    add_message "MapMessage" do
+      map :map_string_int32, :string, :int32, 1
+      map :map_string_msg, :string, :message, 2, "TestMessage2"
+    end
+    add_message "MapMessageWireEquiv" do
+      repeated :map_string_int32, :message, 1, "MapMessageWireEquiv_entry1"
+      repeated :map_string_msg, :message, 2, "MapMessageWireEquiv_entry2"
+    end
+    add_message "MapMessageWireEquiv_entry1" do
+      optional :key, :string, 1
+      optional :value, :int32, 2
+    end
+    add_message "MapMessageWireEquiv_entry2" do
+      optional :key, :string, 1
+      optional :value, :message, 2, "TestMessage2"
+    end
   end
 
   TestMessage = pool.lookup("TestMessage").msgclass
@@ -61,6 +81,12 @@ module BasicTest
   Recursive2 = pool.lookup("Recursive2").msgclass
   TestEnum = pool.lookup("TestEnum").enummodule
   BadFieldNames = pool.lookup("BadFieldNames").msgclass
+  MapMessage = pool.lookup("MapMessage").msgclass
+  MapMessageWireEquiv = pool.lookup("MapMessageWireEquiv").msgclass
+  MapMessageWireEquiv_entry1 =
+    pool.lookup("MapMessageWireEquiv_entry1").msgclass
+  MapMessageWireEquiv_entry2 =
+    pool.lookup("MapMessageWireEquiv_entry2").msgclass
 
 # ------------ test cases ---------------
 
@@ -300,7 +326,7 @@ module BasicTest
       l.push :B
       l.push :C
       assert l.count == 3
-      assert_raise NameError do
+      assert_raise RangeError do
         l.push :D
       end
       assert l[0] == :A
@@ -324,12 +350,240 @@ module BasicTest
       end
     end
 
+    def test_map_basic
+      # allowed key types:
+      # :int32, :int64, :uint32, :uint64, :bool, :string, :bytes.
+
+      m = Google::Protobuf::Map.new(:string, :int32)
+      m["asdf"] = 1
+      assert m["asdf"] == 1
+      m["jkl;"] = 42
+      assert m == { "jkl;" => 42, "asdf" => 1 }
+      assert m.has_key?("asdf")
+      assert !m.has_key?("qwerty")
+      assert m.length == 2
+
+      m2 = m.dup
+      assert m == m2
+      assert m.hash != 0
+      assert m.hash == m2.hash
+
+      collected = {}
+      m.each { |k,v| collected[v] = k }
+      assert collected == { 42 => "jkl;", 1 => "asdf" }
+
+      assert m.delete("asdf") == 1
+      assert !m.has_key?("asdf")
+      assert m["asdf"] == nil
+      assert !m.has_key?("asdf")
+
+      # We only assert on inspect value when there is one map entry because the
+      # order in which elements appear is unspecified (depends on the internal
+      # hash function). We don't want a brittle test.
+      assert m.inspect == "{\"jkl;\" => 42}"
+
+      assert m.keys == ["jkl;"]
+      assert m.values == [42]
+
+      m.clear
+      assert m.length == 0
+      assert m == {}
+
+      assert_raise TypeError do
+        m[1] = 1
+      end
+      assert_raise RangeError do
+        m["asdf"] = 0x1_0000_0000
+      end
+    end
+
+    def test_map_ctor
+      m = Google::Protobuf::Map.new(:string, :int32,
+                                    {"a" => 1, "b" => 2, "c" => 3})
+      assert m == {"a" => 1, "c" => 3, "b" => 2}
+    end
+
+    def test_map_keytypes
+      m = Google::Protobuf::Map.new(:int32, :int32)
+      m[1] = 42
+      m[-1] = 42
+      assert_raise RangeError do
+        m[0x8000_0000] = 1
+      end
+      assert_raise TypeError do
+        m["asdf"] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:int64, :int32)
+      m[0x1000_0000_0000_0000] = 1
+      assert_raise RangeError do
+        m[0x1_0000_0000_0000_0000] = 1
+      end
+      assert_raise TypeError do
+        m["asdf"] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:uint32, :int32)
+      m[0x8000_0000] = 1
+      assert_raise RangeError do
+        m[0x1_0000_0000] = 1
+      end
+      assert_raise RangeError do
+        m[-1] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:uint64, :int32)
+      m[0x8000_0000_0000_0000] = 1
+      assert_raise RangeError do
+        m[0x1_0000_0000_0000_0000] = 1
+      end
+      assert_raise RangeError do
+        m[-1] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:bool, :int32)
+      m[true] = 1
+      m[false] = 2
+      assert_raise TypeError do
+        m[1] = 1
+      end
+      assert_raise TypeError do
+        m["asdf"] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:string, :int32)
+      m["asdf"] = 1
+      assert_raise TypeError do
+        m[1] = 1
+      end
+      assert_raise TypeError do
+        bytestring = ["FFFF"].pack("H*")
+        m[bytestring] = 1
+      end
+
+      m = Google::Protobuf::Map.new(:bytes, :int32)
+      bytestring = ["FFFF"].pack("H*")
+      m[bytestring] = 1
+      assert_raise TypeError do
+        m["asdf"] = 1
+      end
+      assert_raise TypeError do
+        m[1] = 1
+      end
+    end
+
+    def test_map_msg_enum_valuetypes
+      m = Google::Protobuf::Map.new(:string, :message, TestMessage)
+      m["asdf"] = TestMessage.new
+      assert_raise TypeError do
+        m["jkl;"] = TestMessage2.new
+      end
+
+      m = Google::Protobuf::Map.new(:string, :message, TestMessage,
+                                    { "a" => TestMessage.new(:optional_int32 => 42),
+                                      "b" => TestMessage.new(:optional_int32 => 84) })
+      assert m.length == 2
+      assert m.values.map{|msg| msg.optional_int32}.sort == [42, 84]
+
+      m = Google::Protobuf::Map.new(:string, :enum, TestEnum,
+                                    { "x" => :A, "y" => :B, "z" => :C })
+      assert m.length == 3
+      assert m["z"] == :C
+      m["z"] = 2
+      assert m["z"] == :B
+      m["z"] = 4
+      assert m["z"] == 4
+      assert_raise RangeError do
+        m["z"] = :Z
+      end
+      assert_raise TypeError do
+        m["z"] = "z"
+      end
+    end
+
+    def test_map_dup_deep_copy
+      m = Google::Protobuf::Map.new(:string, :message, TestMessage,
+                                    { "a" => TestMessage.new(:optional_int32 => 42),
+                                      "b" => TestMessage.new(:optional_int32 => 84) })
+
+      m2 = m.dup
+      assert m == m2
+      assert m.object_id != m2.object_id
+      assert m["a"].object_id == m2["a"].object_id
+      assert m["b"].object_id == m2["b"].object_id
+
+      m2 = Google::Protobuf.deep_copy(m)
+      assert m == m2
+      assert m.object_id != m2.object_id
+      assert m["a"].object_id != m2["a"].object_id
+      assert m["b"].object_id != m2["b"].object_id
+    end
+
+    def test_map_field
+      m = MapMessage.new
+      assert m.map_string_int32 == {}
+      assert m.map_string_msg == {}
+
+      m = MapMessage.new(:map_string_int32 => {"a" => 1, "b" => 2},
+                         :map_string_msg => {"a" => TestMessage2.new(:foo => 1),
+                                             "b" => TestMessage2.new(:foo => 2)})
+      assert m.map_string_int32.keys.sort == ["a", "b"]
+      assert m.map_string_int32["a"] == 1
+      assert m.map_string_msg["b"].foo == 2
+
+      m.map_string_int32["c"] = 3
+      assert m.map_string_int32["c"] == 3
+      m.map_string_msg["c"] = TestMessage2.new(:foo => 3)
+      assert m.map_string_msg["c"] == TestMessage2.new(:foo => 3)
+      m.map_string_msg.delete("b")
+      m.map_string_msg.delete("c")
+      assert m.map_string_msg == { "a" => TestMessage2.new(:foo => 1) }
+
+      assert_raise TypeError do
+        m.map_string_msg["e"] = TestMessage.new # wrong value type
+      end
+      # ensure nothing was added by the above
+      assert m.map_string_msg == { "a" => TestMessage2.new(:foo => 1) }
+
+      m.map_string_int32 = Google::Protobuf::Map.new(:string, :int32)
+      assert_raise TypeError do
+        m.map_string_int32 = Google::Protobuf::Map.new(:string, :int64)
+      end
+      assert_raise TypeError do
+        m.map_string_int32 = {}
+      end
+
+      assert_raise TypeError do
+        m = MapMessage.new(:map_string_int32 => { 1 => "I am not a number" })
+      end
+    end
+
+    def test_map_encode_decode
+      m = MapMessage.new(:map_string_int32 => {"a" => 1, "b" => 2},
+                         :map_string_msg => {"a" => TestMessage2.new(:foo => 1),
+                                             "b" => TestMessage2.new(:foo => 2)})
+      m2 = MapMessage.decode(MapMessage.encode(m))
+      assert m == m2
+
+      m3 = MapMessageWireEquiv.decode(MapMessage.encode(m))
+      assert m3.map_string_int32.length == 2
+
+      kv = {}
+      m3.map_string_int32.map { |msg| kv[msg.key] = msg.value }
+      assert kv == {"a" => 1, "b" => 2}
+
+      kv = {}
+      m3.map_string_msg.map { |msg| kv[msg.key] = msg.value }
+      assert kv == {"a" => TestMessage2.new(:foo => 1),
+                    "b" => TestMessage2.new(:foo => 2)}
+    end
+
     def test_enum_field
       m = TestMessage.new
       assert m.optional_enum == :Default
       m.optional_enum = :A
       assert m.optional_enum == :A
-      assert_raise NameError do
+      assert_raise RangeError do
         m.optional_enum = :ASDF
       end
       m.optional_enum = 1