Support oneofs in MRI Ruby C extension.

ruby/ext/google/protobuf_c/encode_decode.c

@@ -59,6 +59,30 @@ static const void *newsubmsghandlerdata(upb_handlers* h, uint32_t ofs,
   return hd;
 }
 
+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
+} oneof_handlerdata_t;
+
+static const void *newoneofhandlerdata(upb_handlers *h,
+                                       uint32_t ofs,
+                                       uint32_t case_ofs,
+                                       const upb_fielddef *f) {
+  oneof_handlerdata_t *hd = ALLOC(oneof_handlerdata_t);
+  hd->ofs = ofs;
+  hd->case_ofs = case_ofs;
+  hd->tag = upb_fielddef_number(f);
+  if (upb_fielddef_type(f) == UPB_TYPE_MESSAGE) {
+    hd->md = upb_fielddef_msgsubdef(f);
+  } else {
+    hd->md = NULL;
+  }
+  upb_handlers_addcleanup(h, hd, free);
+  return hd;
+}
+
 // A handler that starts a repeated field.  Gets the Repeated*Field instance for
 // this field (such an instance always exists even in an empty message).
 static void *startseq_handler(void* closure, const void* hd) {
@@ -67,8 +91,7 @@ static void *startseq_handler(void* closure, const void* hd) {
   return (void*)DEREF(msg, *ofs, VALUE);
 }
 
-// Handlers that append primitive values to a repeated field (a regular Ruby
-// array for now).
+// Handlers that append primitive values to a repeated field.
 #define DEFINE_APPEND_HANDLER(type, ctype)                 \
   static bool append##type##_handler(void *closure, const void *hd, \
                                      ctype val) {                   \
@@ -85,7 +108,7 @@ DEFINE_APPEND_HANDLER(int64,  int64_t)
 DEFINE_APPEND_HANDLER(uint64, uint64_t)
 DEFINE_APPEND_HANDLER(double, double)
 
-// Appends a string to a repeated field (a regular Ruby array for now).
+// Appends a string to a repeated field.
 static void* appendstr_handler(void *closure,
                                const void *hd,
                                size_t size_hint) {
@@ -96,7 +119,7 @@ static void* appendstr_handler(void *closure,
   return (void*)str;
 }
 
-// Appends a 'bytes' string to a repeated field (a regular Ruby array for now).
+// Appends a 'bytes' string to a repeated field.
 static void* appendbytes_handler(void *closure,
                                  const void *hd,
                                  size_t size_hint) {
@@ -266,6 +289,75 @@ static map_handlerdata_t* new_map_handlerdata(
   return hd;
 }
 
+// Handlers that set primitive values in oneofs.
+#define DEFINE_ONEOF_HANDLER(type, ctype)                           \
+  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->ofs, ctype) = val;                    \
+    return true;                                                    \
+  }
+
+DEFINE_ONEOF_HANDLER(bool,   bool)
+DEFINE_ONEOF_HANDLER(int32,  int32_t)
+DEFINE_ONEOF_HANDLER(uint32, uint32_t)
+DEFINE_ONEOF_HANDLER(float,  float)
+DEFINE_ONEOF_HANDLER(int64,  int64_t)
+DEFINE_ONEOF_HANDLER(uint64, uint64_t)
+DEFINE_ONEOF_HANDLER(double, double)
+
+#undef DEFINE_ONEOF_HANDLER
+
+// Handlers for strings in a oneof.
+static void *oneofstr_handler(void *closure,
+                              const void *hd,
+                              size_t size_hint) {
+  MessageHeader* msg = 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->ofs, VALUE) = str;
+  return (void*)str;
+}
+
+static void *oneofbytes_handler(void *closure,
+                                const void *hd,
+                                size_t size_hint) {
+  MessageHeader* msg = 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->ofs, VALUE) = str;
+  return (void*)str;
+}
+
+// Handler for a submessage field in a oneof.
+static void *oneofsubmsg_handler(void *closure,
+                                 const void *hd) {
+  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;
+
+  VALUE subdesc =
+      get_def_obj((void*)oneofdata->md);
+  VALUE subklass = Descriptor_msgclass(subdesc);
+
+  if (oldcase != oneofdata->tag ||
+      DEREF(msg, oneofdata->ofs, VALUE) == Qnil) {
+    DEREF(msg, oneofdata->ofs, VALUE) =
+        rb_class_new_instance(0, NULL, subklass);
+  }
+
+  VALUE submsg_rb = DEREF(msg, oneofdata->ofs, VALUE);
+  MessageHeader* submsg;
+  TypedData_Get_Struct(submsg_rb, MessageHeader, &Message_type, submsg);
+  return submsg;
+}
+
 // Set up handlers for a repeated field.
 static void add_handlers_for_repeated_field(upb_handlers *h,
                                             const upb_fielddef *f,
@@ -383,6 +475,53 @@ static void add_handlers_for_mapentry(const upb_msgdef* msgdef,
       offsetof(map_parse_frame_t, value_storage));
 }
 
+// Set up handlers for a oneof field.
+static void add_handlers_for_oneof_field(upb_handlers *h,
+                                         const upb_fielddef *f,
+                                         size_t offset,
+                                         size_t oneof_case_offset) {
+
+  upb_handlerattr attr = UPB_HANDLERATTR_INITIALIZER;
+  upb_handlerattr_sethandlerdata(
+      &attr, newoneofhandlerdata(h, offset, oneof_case_offset, f));
+
+  switch (upb_fielddef_type(f)) {
+
+#define SET_HANDLER(utype, ltype)                                 \
+  case utype:                                                     \
+    upb_handlers_set##ltype(h, f, oneof##ltype##_handler, &attr); \
+    break;
+
+    SET_HANDLER(UPB_TYPE_BOOL,   bool);
+    SET_HANDLER(UPB_TYPE_INT32,  int32);
+    SET_HANDLER(UPB_TYPE_UINT32, uint32);
+    SET_HANDLER(UPB_TYPE_ENUM,   int32);
+    SET_HANDLER(UPB_TYPE_FLOAT,  float);
+    SET_HANDLER(UPB_TYPE_INT64,  int64);
+    SET_HANDLER(UPB_TYPE_UINT64, uint64);
+    SET_HANDLER(UPB_TYPE_DOUBLE, double);
+
+#undef SET_HANDLER
+
+    case UPB_TYPE_STRING:
+    case UPB_TYPE_BYTES: {
+      bool is_bytes = upb_fielddef_type(f) == UPB_TYPE_BYTES;
+      upb_handlers_setstartstr(h, f, is_bytes ?
+                               oneofbytes_handler : oneofstr_handler,
+                               &attr);
+      upb_handlers_setstring(h, f, stringdata_handler, NULL);
+      break;
+    }
+    case UPB_TYPE_MESSAGE: {
+      upb_handlers_setstartsubmsg(h, f, oneofsubmsg_handler, &attr);
+      break;
+    }
+  }
+
+  upb_handlerattr_uninit(&attr);
+}
+
+
 static void add_handlers_for_message(const void *closure, upb_handlers *h) {
   const upb_msgdef* msgdef = upb_handlers_msgdef(h);
   Descriptor* desc = ruby_to_Descriptor(get_def_obj((void*)msgdef));
@@ -402,16 +541,20 @@ static void add_handlers_for_message(const void *closure, upb_handlers *h) {
     desc->layout = create_layout(desc->msgdef);
   }
 
-  upb_msg_iter i;
-
-  for (upb_msg_begin(&i, desc->msgdef);
-       !upb_msg_done(&i);
-       upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for (upb_msg_field_begin(&i, desc->msgdef);
+       !upb_msg_field_done(&i);
+       upb_msg_field_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->fields[upb_fielddef_index(f)].offset +
+        sizeof(MessageHeader);
+    size_t oneof_case_offset =
+        desc->layout->fields[upb_fielddef_index(f)].case_offset +
         sizeof(MessageHeader);
 
-    if (is_map_field(f)) {
+    if (upb_fielddef_containingoneof(f)) {
+      add_handlers_for_oneof_field(h, f, offset, oneof_case_offset);
+    } else if (is_map_field(f)) {
       add_handlers_for_mapfield(h, f, offset, desc);
     } else if (upb_fielddef_isseq(f)) {
       add_handlers_for_repeated_field(h, f, offset);
@@ -804,13 +947,28 @@ static void putmsg(VALUE msg_rb, const Descriptor* desc,
   MessageHeader* msg;
   TypedData_Get_Struct(msg_rb, MessageHeader, &Message_type, msg);
 
-  upb_msg_iter i;
-  for (upb_msg_begin(&i, desc->msgdef);
-       !upb_msg_done(&i);
-       upb_msg_next(&i)) {
+  upb_msg_field_iter i;
+  for (upb_msg_field_begin(&i, desc->msgdef);
+       !upb_msg_field_done(&i);
+       upb_msg_field_next(&i)) {
     upb_fielddef *f = upb_msg_iter_field(&i);
     uint32_t offset =
-        desc->layout->offsets[upb_fielddef_index(f)] + sizeof(MessageHeader);
+        desc->layout->fields[upb_fielddef_index(f)].offset +
+        sizeof(MessageHeader);
+    uint32_t oneof_case_offset =
+        desc->layout->fields[upb_fielddef_index(f)].case_offset +
+        sizeof(MessageHeader);
+
+    if (upb_fielddef_containingoneof(f)) {
+      // For a oneof, check that this field is actually present -- skip all the
+      // below if not.
+      if (DEREF(msg, oneof_case_offset, uint32_t) !=
+          upb_fielddef_number(f)) {
+        continue;
+      }
+      // Otherwise, fall through to the appropriate singular-field handler
+      // below.
+    }
 
     if (is_map_field(f)) {
       VALUE map = DEREF(msg, offset, VALUE);

ruby/ext/google/protobuf_c/protobuf.c

@@ -77,8 +77,10 @@ void Init_protobuf_c() {
   DescriptorPool_register(protobuf);
   Descriptor_register(protobuf);
   FieldDescriptor_register(protobuf);
+  OneofDescriptor_register(protobuf);
   EnumDescriptor_register(protobuf);
   MessageBuilderContext_register(internal);
+  OneofBuilderContext_register(internal);
   EnumBuilderContext_register(internal);
   Builder_register(internal);
   RepeatedField_register(protobuf);

ruby/ext/google/protobuf_c/protobuf.h

@@ -43,6 +43,7 @@ struct Descriptor;
 struct FieldDescriptor;
 struct EnumDescriptor;
 struct MessageLayout;
+struct MessageField;
 struct MessageHeader;
 struct MessageBuilderContext;
 struct EnumBuilderContext;
@@ -51,10 +52,13 @@ struct Builder;
 typedef struct DescriptorPool DescriptorPool;
 typedef struct Descriptor Descriptor;
 typedef struct FieldDescriptor FieldDescriptor;
+typedef struct OneofDescriptor OneofDescriptor;
 typedef struct EnumDescriptor EnumDescriptor;
 typedef struct MessageLayout MessageLayout;
+typedef struct MessageField MessageField;
 typedef struct MessageHeader MessageHeader;
 typedef struct MessageBuilderContext MessageBuilderContext;
+typedef struct OneofBuilderContext OneofBuilderContext;
 typedef struct EnumBuilderContext EnumBuilderContext;
 typedef struct Builder Builder;
 
@@ -120,6 +124,10 @@ struct FieldDescriptor {
   const upb_fielddef* fielddef;
 };
 
+struct OneofDescriptor {
+  const upb_oneofdef* oneofdef;
+};
+
 struct EnumDescriptor {
   const upb_enumdef* enumdef;
   VALUE module;  // begins as nil
@@ -130,6 +138,11 @@ struct MessageBuilderContext {
   VALUE builder;
 };
 
+struct OneofBuilderContext {
+  VALUE descriptor;
+  VALUE builder;
+};
+
 struct EnumBuilderContext {
   VALUE enumdesc;
 };
@@ -144,6 +157,7 @@ extern VALUE cDescriptor;
 extern VALUE cFieldDescriptor;
 extern VALUE cEnumDescriptor;
 extern VALUE cMessageBuilderContext;
+extern VALUE cOneofBuilderContext;
 extern VALUE cEnumBuilderContext;
 extern VALUE cBuilder;
 
@@ -175,6 +189,9 @@ VALUE Descriptor_name_set(VALUE _self, VALUE str);
 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_lookup_oneof(VALUE _self, VALUE name);
 VALUE Descriptor_msgclass(VALUE _self);
 extern const rb_data_type_t _Descriptor_type;
 
@@ -199,6 +216,16 @@ VALUE FieldDescriptor_set(VALUE _self, VALUE msg_rb, VALUE value);
 upb_fieldtype_t ruby_to_fieldtype(VALUE type);
 VALUE fieldtype_to_ruby(upb_fieldtype_t type);
 
+void OneofDescriptor_mark(void* _self);
+void OneofDescriptor_free(void* _self);
+VALUE OneofDescriptor_alloc(VALUE klass);
+void OneofDescriptor_register(VALUE module);
+OneofDescriptor* ruby_to_OneofDescriptor(VALUE value);
+VALUE OneofDescriptor_name(VALUE _self);
+VALUE OneofDescriptor_name_set(VALUE _self, VALUE value);
+VALUE OneofDescriptor_add_field(VALUE _self, VALUE field);
+VALUE OneofDescriptor_each(VALUE _self, VALUE field);
+
 void EnumDescriptor_mark(void* _self);
 void EnumDescriptor_free(void* _self);
 VALUE EnumDescriptor_alloc(VALUE klass);
@@ -225,6 +252,17 @@ 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);
+VALUE MessageBuilderContext_oneof(VALUE _self, VALUE name);
+
+void OneofBuilderContext_mark(void* _self);
+void OneofBuilderContext_free(void* _self);
+VALUE OneofBuilderContext_alloc(VALUE klass);
+void OneofBuilderContext_register(VALUE module);
+OneofBuilderContext* ruby_to_OneofBuilderContext(VALUE value);
+VALUE OneofBuilderContext_initialize(VALUE _self,
+                                     VALUE descriptor,
+                                     VALUE builder);
+VALUE OneofBuilderContext_optional(int argc, VALUE* argv, VALUE _self);
 
 void EnumBuilderContext_mark(void* _self);
 void EnumBuilderContext_free(void* _self);
@@ -247,7 +285,7 @@ VALUE Builder_finalize_to_pool(VALUE _self, VALUE pool_rb);
 // Native slot storage abstraction.
 // -----------------------------------------------------------------------------
 
-#define NATIVE_SLOT_MAX_SIZE sizeof(void*)
+#define NATIVE_SLOT_MAX_SIZE sizeof(uint64_t)
 
 size_t native_slot_size(upb_fieldtype_t type);
 void native_slot_set(upb_fieldtype_t type,
@@ -384,9 +422,16 @@ VALUE Map_iter_value(Map_iter* iter);
 // Message layout / storage.
 // -----------------------------------------------------------------------------
 
+#define MESSAGE_FIELD_NO_CASE ((size_t)-1)
+
+struct MessageField {
+  size_t offset;
+  size_t case_offset;  // for oneofs, a uint32. Else, MESSAGE_FIELD_NO_CASE.
+};
+
 struct MessageLayout {
   const upb_msgdef* msgdef;
-  size_t* offsets;
+  MessageField* fields;
   size_t size;
 };
 

ruby/tests/basic.rb

@@ -73,6 +73,15 @@ module BasicTest
       optional :key, :string, 1
       optional :value, :message, 2, "TestMessage2"
     end
+
+    add_message "OneofMessage" do
+      oneof :my_oneof do
+        optional :a, :string, 1
+        optional :b, :int32, 2
+        optional :c, :message, 3, "TestMessage2"
+        optional :d, :enum, 4, "TestEnum"
+      end
+    end
   end
 
   TestMessage = pool.lookup("TestMessage").msgclass
@@ -87,6 +96,7 @@ module BasicTest
     pool.lookup("MapMessageWireEquiv_entry1").msgclass
   MapMessageWireEquiv_entry2 =
     pool.lookup("MapMessageWireEquiv_entry2").msgclass
+  OneofMessage = pool.lookup("OneofMessage").msgclass
 
 # ------------ test cases ---------------
 
@@ -582,6 +592,80 @@ module BasicTest
                     "b" => TestMessage2.new(:foo => 2)}
     end
 
+    def test_oneof_descriptors
+      d = OneofMessage.descriptor
+      o = d.lookup_oneof("my_oneof")
+      assert o != nil
+      assert o.class == Google::Protobuf::OneofDescriptor
+      assert o.name == "my_oneof"
+      assert d.oneofs == [o]
+      assert o.count == 4
+      field_names = o.map{|f| f.name}.sort
+      assert field_names == ["a", "b", "c", "d"]
+    end
+
+    def test_oneof
+      d = OneofMessage.new
+      assert d.a == nil
+      assert d.b == nil
+      assert d.c == nil
+      assert d.d == nil
+
+      d.a = "hi"
+      assert d.a == "hi"
+      assert d.b == nil
+      assert d.c == nil
+      assert d.d == nil
+
+      d.b = 42
+      assert d.a == nil
+      assert d.b == 42
+      assert d.c == nil
+      assert d.d == nil
+
+      d.c = TestMessage2.new(:foo => 100)
+      assert d.a == nil
+      assert d.b == nil
+      assert d.c.foo == 100
+      assert d.d == nil
+
+      d.d = :C
+      assert d.a == nil
+      assert d.b == nil
+      assert d.c == nil
+      assert d.d == :C
+
+      d2 = OneofMessage.decode(OneofMessage.encode(d))
+      assert d2 == d
+
+      encoded_field_a = OneofMessage.encode(OneofMessage.new(:a => "string"))
+      encoded_field_b = OneofMessage.encode(OneofMessage.new(:b => 1000))
+      encoded_field_c = OneofMessage.encode(
+        OneofMessage.new(:c => TestMessage2.new(:foo => 1)))
+      encoded_field_d = OneofMessage.encode(OneofMessage.new(:d => :B))
+
+      d3 = OneofMessage.decode(
+        encoded_field_c + encoded_field_a + encoded_field_d)
+      assert d3.a == nil
+      assert d3.b == nil
+      assert d3.c == nil
+      assert d3.d == :B
+
+      d4 = OneofMessage.decode(
+        encoded_field_c + encoded_field_a + encoded_field_d +
+        encoded_field_c)
+      assert d4.a == nil
+      assert d4.b == nil
+      assert d4.c.foo == 1
+      assert d4.d == nil
+
+      d5 = OneofMessage.new(:a => "hello")
+      assert d5.a != nil
+      d5.a = nil
+      assert d5.a == nil
+      assert OneofMessage.encode(d5) == ''
+    end
+
     def test_enum_field
       m = TestMessage.new
       assert m.optional_enum == :Default
@@ -621,6 +705,14 @@ module BasicTest
       assert m.repeated_msg[0].object_id != m2.repeated_msg[0].object_id
     end
 
+    def test_eq
+      m = TestMessage.new(:optional_int32 => 42,
+                          :repeated_int32 => [1, 2, 3])
+      m2 = TestMessage.new(:optional_int32 => 43,
+                           :repeated_int32 => [1, 2, 3])
+      assert m != m2
+    end
+
     def test_enum_lookup
       assert TestEnum::A == 1
       assert TestEnum::B == 2