Merge pull request #6797 from haberman/ruby-lazy-wrappers

Ruby lazy wrappers optimization

ruby/ext/google/protobuf_c/map.c

@@ -559,7 +559,8 @@ VALUE Map_deep_copy(VALUE _self) {
     void* mem = value_memory(&v);
     upb_value dup;
     void* dup_mem = value_memory(&dup);
-    native_slot_deep_copy(self->value_type, dup_mem, mem);
+    native_slot_deep_copy(self->value_type, self->value_type_class, dup_mem,
+                          mem);
 
     if (!upb_strtable_insert2(&new_self->table,
                               upb_strtable_iter_key(&it),
@@ -631,7 +632,8 @@ VALUE Map_eq(VALUE _self, VALUE _other) {
       return Qfalse;
     }
 
-    if (!native_slot_eq(self->value_type, mem, other_mem)) {
+    if (!native_slot_eq(self->value_type, self->value_type_class, mem,
+                        other_mem)) {
       // Present, but value not equal.
       return Qfalse;
     }

ruby/ext/google/protobuf_c/message.c

@@ -62,13 +62,12 @@ VALUE Message_alloc(VALUE klass) {
   Descriptor* desc = ruby_to_Descriptor(descriptor);
   MessageHeader* msg;
   VALUE ret;
-  size_t size;
 
   if (desc->layout == NULL) {
     create_layout(desc);
   }
 
-  msg = ALLOC_N(uint8_t, sizeof(MessageHeader) + desc->layout->size);
+  msg = (void*)ALLOC_N(uint8_t, sizeof(MessageHeader) + desc->layout->size);
   msg->descriptor = desc;
   msg->unknown_fields = NULL;
   memcpy(Message_data(msg), desc->layout->empty_template, desc->layout->size);
@@ -109,30 +108,36 @@ enum {
 };
 
 // Check if the field is a well known wrapper type
-static bool is_wrapper_type_field(const MessageLayout* layout,
-                                  const upb_fielddef* field) {
-  const char* field_type_name = rb_class2name(field_type_class(layout, field));
-
-  return strcmp(field_type_name, "Google::Protobuf::DoubleValue") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::FloatValue") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::Int32Value") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::Int64Value") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::UInt32Value") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::UInt64Value") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::BoolValue") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::StringValue") == 0 ||
-         strcmp(field_type_name, "Google::Protobuf::BytesValue") == 0;
+bool is_wrapper_type_field(const upb_fielddef* field) {
+  const upb_msgdef *m;
+  if (upb_fielddef_type(field) != UPB_TYPE_MESSAGE) {
+    return false;
+  }
+  m = upb_fielddef_msgsubdef(field);
+  switch (upb_msgdef_wellknowntype(m)) {
+    case UPB_WELLKNOWN_DOUBLEVALUE:
+    case UPB_WELLKNOWN_FLOATVALUE:
+    case UPB_WELLKNOWN_INT64VALUE:
+    case UPB_WELLKNOWN_UINT64VALUE:
+    case UPB_WELLKNOWN_INT32VALUE:
+    case UPB_WELLKNOWN_UINT32VALUE:
+    case UPB_WELLKNOWN_STRINGVALUE:
+    case UPB_WELLKNOWN_BYTESVALUE:
+    case UPB_WELLKNOWN_BOOLVALUE:
+      return true;
+    default:
+      return false;
+  }
 }
 
 // Get a new Ruby wrapper type and set the initial value
-static VALUE ruby_wrapper_type(const MessageLayout* layout,
-                               const upb_fielddef* field, const VALUE value) {
-  if (is_wrapper_type_field(layout, field) && value != Qnil) {
+VALUE ruby_wrapper_type(VALUE type_class, VALUE value) {
+  if (value != Qnil) {
     VALUE hash = rb_hash_new();
     rb_hash_aset(hash, rb_str_new2("value"), value);
     {
       VALUE args[1] = {hash};
-      return rb_class_new_instance(1, args, field_type_class(layout, field));
+      return rb_class_new_instance(1, args, type_class);
     }
   }
   return Qnil;
@@ -193,8 +198,7 @@ static int extract_method_call(VALUE method_name, MessageHeader* self,
     // Check if field exists and is a wrapper type
     if (upb_msgdef_lookupname(self->descriptor->msgdef, wrapper_field_name,
                               name_len - 9, &test_f_wrapper, &test_o_wrapper) &&
-        upb_fielddef_type(test_f_wrapper) == UPB_TYPE_MESSAGE &&
-        is_wrapper_type_field(self->descriptor->layout, test_f_wrapper)) {
+        is_wrapper_type_field(test_f_wrapper)) {
       // It does exist!
       has_field = true;
       if (accessor_type == METHOD_SETTER) {
@@ -329,12 +333,17 @@ VALUE Message_method_missing(int argc, VALUE* argv, VALUE _self) {
     return layout_has(self->descriptor->layout, Message_data(self), f);
   } else if (accessor_type == METHOD_WRAPPER_GETTER) {
     VALUE value = layout_get(self->descriptor->layout, Message_data(self), f);
-    if (value != Qnil) {
-      value = rb_funcall(value, rb_intern("value"), 0);
+    switch (TYPE(value)) {
+      case T_DATA:
+        return rb_funcall(value, rb_intern("value"), 0);
+      case T_NIL:
+        return Qnil;
+      default:
+        return value;
     }
-    return value;
   } else if (accessor_type == METHOD_WRAPPER_SETTER) {
-    VALUE wrapper = ruby_wrapper_type(self->descriptor->layout, f, argv[1]);
+    VALUE wrapper = ruby_wrapper_type(
+        field_type_class(self->descriptor->layout, f), argv[1]);
     layout_set(self->descriptor->layout, Message_data(self), f, wrapper);
     return Qnil;
   } else if (accessor_type == METHOD_ENUM_GETTER) {

ruby/ext/google/protobuf_c/protobuf.h

@@ -363,8 +363,10 @@ VALUE native_slot_get(upb_fieldtype_t type,
 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);
-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_deep_copy(upb_fieldtype_t type, VALUE type_class, void* to,
+                           void* from);
+bool native_slot_eq(upb_fieldtype_t type, VALUE type_class, void* mem1,
+                    void* mem2);
 
 VALUE native_slot_encode_and_freeze_string(upb_fieldtype_t type, VALUE value);
 void native_slot_check_int_range_precision(const char* name, upb_fieldtype_t type, VALUE value);
@@ -556,6 +558,9 @@ VALUE layout_eq(MessageLayout* layout, void* msg1, void* msg2);
 VALUE layout_hash(MessageLayout* layout, void* storage);
 VALUE layout_inspect(MessageLayout* layout, void* storage);
 
+bool is_wrapper_type_field(const upb_fielddef* field);
+VALUE ruby_wrapper_type(VALUE type_class, VALUE value);
+
 // -----------------------------------------------------------------------------
 // Message class creation.
 // -----------------------------------------------------------------------------

ruby/ext/google/protobuf_c/repeated_field.c

@@ -378,7 +378,7 @@ VALUE RepeatedField_deep_copy(VALUE _self) {
   for (i = 0; i < self->size; i++, off += elem_size) {
     void* to_mem = (uint8_t *)new_rptfield_self->elements + off;
     void* from_mem = (uint8_t *)self->elements + off;
-    native_slot_deep_copy(field_type, to_mem, from_mem);
+    native_slot_deep_copy(field_type, self->field_type_class, to_mem, from_mem);
     new_rptfield_self->size++;
   }
 
@@ -451,7 +451,8 @@ VALUE RepeatedField_eq(VALUE _self, VALUE _other) {
     for (i = 0; i < self->size; i++, off += elem_size) {
       void* self_mem = ((uint8_t *)self->elements) + off;
       void* other_mem = ((uint8_t *)other->elements) + off;
-      if (!native_slot_eq(field_type, self_mem, other_mem)) {
+      if (!native_slot_eq(field_type, self->field_type_class, self_mem,
+                          other_mem)) {
         return Qfalse;
       }
     }

ruby/ext/google/protobuf_c/storage.c

@@ -294,8 +294,20 @@ VALUE native_slot_get(upb_fieldtype_t type,
       return DEREF(memory, int8_t) ? Qtrue : Qfalse;
     case UPB_TYPE_STRING:
     case UPB_TYPE_BYTES:
-    case UPB_TYPE_MESSAGE:
       return DEREF(memory, VALUE);
+    case UPB_TYPE_MESSAGE: {
+      VALUE val = DEREF(memory, VALUE);
+
+      // Lazily expand wrapper type if necessary.
+      int type = TYPE(val);
+      if (type != T_DATA && type != T_NIL) {
+        // This must be a wrapper type.
+        val = ruby_wrapper_type(type_class, val);
+        DEREF(memory, VALUE) = val;
+      }
+
+      return val;
+    }
     case UPB_TYPE_ENUM: {
       int32_t val = DEREF(memory, int32_t);
       VALUE symbol = enum_lookup(type_class, INT2NUM(val));
@@ -372,7 +384,8 @@ void native_slot_dup(upb_fieldtype_t type, void* to, void* from) {
   memcpy(to, from, native_slot_size(type));
 }
 
-void native_slot_deep_copy(upb_fieldtype_t type, void* to, void* from) {
+void native_slot_deep_copy(upb_fieldtype_t type, VALUE type_class, void* to,
+                           void* from) {
   switch (type) {
     case UPB_TYPE_STRING:
     case UPB_TYPE_BYTES: {
@@ -382,7 +395,7 @@ void native_slot_deep_copy(upb_fieldtype_t type, void* to, void* from) {
       break;
     }
     case UPB_TYPE_MESSAGE: {
-      VALUE from_val = DEREF(from, VALUE);
+      VALUE from_val = native_slot_get(type, type_class, from);
       DEREF(to, VALUE) = (from_val != Qnil) ?
           Message_deep_copy(from_val) : Qnil;
       break;
@@ -392,13 +405,14 @@ 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) {
+bool native_slot_eq(upb_fieldtype_t type, VALUE type_class, void* mem1,
+                    void* mem2) {
   switch (type) {
     case UPB_TYPE_STRING:
     case UPB_TYPE_BYTES:
     case UPB_TYPE_MESSAGE: {
-      VALUE val1 = DEREF(mem1, VALUE);
-      VALUE val2 = DEREF(mem2, VALUE);
+      VALUE val1 = native_slot_get(type, type_class, mem1);
+      VALUE val2 = native_slot_get(type, type_class, mem2);
       VALUE ret = rb_funcall(val1, rb_intern("=="), 1, val2);
       return ret == Qtrue;
     }
@@ -1025,7 +1039,9 @@ void layout_deep_copy(MessageLayout* layout, void* to, void* from) {
       if (slot_read_oneof_case(layout, from, oneof) ==
           upb_fielddef_number(field)) {
         *to_oneof_case = *from_oneof_case;
-        native_slot_deep_copy(upb_fielddef_type(field), to_memory, from_memory);
+        native_slot_deep_copy(upb_fielddef_type(field),
+                              field_type_class(layout, field), to_memory,
+                              from_memory);
       }
     } else if (is_map_field(field)) {
       DEREF(to_memory, VALUE) =
@@ -1039,7 +1055,9 @@ void layout_deep_copy(MessageLayout* layout, void* to, void* from) {
         slot_set_hasbit(layout, to, field);
       }
 
-      native_slot_deep_copy(upb_fielddef_type(field), to_memory, from_memory);
+      native_slot_deep_copy(upb_fielddef_type(field),
+                            field_type_class(layout, field), to_memory,
+                            from_memory);
     }
   }
 }
@@ -1061,7 +1079,8 @@ VALUE layout_eq(MessageLayout* layout, void* msg1, void* msg2) {
       if (*msg1_oneof_case != *msg2_oneof_case ||
           (slot_read_oneof_case(layout, msg1, oneof) ==
                upb_fielddef_number(field) &&
-           !native_slot_eq(upb_fielddef_type(field), msg1_memory,
+           !native_slot_eq(upb_fielddef_type(field),
+                           field_type_class(layout, field), msg1_memory,
                            msg2_memory))) {
         return Qfalse;
       }
@@ -1078,7 +1097,9 @@ VALUE layout_eq(MessageLayout* layout, void* msg1, void* msg2) {
     } else {
       if (slot_is_hasbit_set(layout, msg1, field) !=
               slot_is_hasbit_set(layout, msg2, field) ||
-          !native_slot_eq(upb_fielddef_type(field), msg1_memory, msg2_memory)) {
+          !native_slot_eq(upb_fielddef_type(field),
+                          field_type_class(layout, field), msg1_memory,
+                          msg2_memory)) {
         return Qfalse;
       }
     }

ruby/tests/basic.rb

@@ -234,6 +234,48 @@ module BasicTest
       m.map_string_int32['aaa'] = 3
     end
 
+    def test_map_wrappers
+      run_asserts = ->(m) {
+        assert_equal 2.0, m.map_double[0].value
+        assert_equal 4.0, m.map_float[0].value
+        assert_equal 3, m.map_int32[0].value
+        assert_equal 4, m.map_int64[0].value
+        assert_equal 5, m.map_uint32[0].value
+        assert_equal 6, m.map_uint64[0].value
+        assert_equal true, m.map_bool[0].value
+        assert_equal 'str', m.map_string[0].value
+        assert_equal 'fun', m.map_bytes[0].value
+      }
+
+      m = proto_module::Wrapper.new(
+        map_double: {0 => Google::Protobuf::DoubleValue.new(value: 2.0)},
+        map_float: {0 => Google::Protobuf::FloatValue.new(value: 4.0)},
+        map_int32: {0 => Google::Protobuf::Int32Value.new(value: 3)},
+        map_int64: {0 => Google::Protobuf::Int64Value.new(value: 4)},
+        map_uint32: {0 => Google::Protobuf::UInt32Value.new(value: 5)},
+        map_uint64: {0 => Google::Protobuf::UInt64Value.new(value: 6)},
+        map_bool: {0 => Google::Protobuf::BoolValue.new(value: true)},
+        map_string: {0 => Google::Protobuf::StringValue.new(value: 'str')},
+        map_bytes: {0 => Google::Protobuf::BytesValue.new(value: 'fun')},
+      )
+
+      run_asserts.call(m)
+      serialized = proto_module::Wrapper::encode(m)
+      m2 = proto_module::Wrapper::decode(serialized)
+      run_asserts.call(m2)
+
+      # Test the case where we are serializing directly from the parsed form
+      # (before anything lazy is materialized).
+      m3 = proto_module::Wrapper::decode(serialized)
+      serialized2 = proto_module::Wrapper::encode(m3)
+      m4 = proto_module::Wrapper::decode(serialized2)
+      run_asserts.call(m4)
+
+      # Test that the lazy form compares equal to the expanded form.
+      m5 = proto_module::Wrapper::decode(serialized2)
+      assert_equal m5, m
+    end
+
     def test_concurrent_decoding
       o = Outer.new
       o.items[0] = Inner.new

ruby/tests/basic_test.proto

@@ -126,7 +126,46 @@ message Wrapper {
   google.protobuf.BytesValue bytes = 9;
   string real_string = 100;
   oneof a_oneof {
-    string oneof_string = 10;
+    string string_in_oneof = 10;
+  }
+
+  // Repeated wrappers don't make sense, but we still need to make sure they
+  // work and don't crash.
+  repeated google.protobuf.DoubleValue repeated_double = 11;
+  repeated google.protobuf.FloatValue repeated_float = 12;
+  repeated google.protobuf.Int32Value repeated_int32 = 13;
+  repeated google.protobuf.Int64Value repeated_int64 = 14;
+  repeated google.protobuf.UInt32Value repeated_uint32 = 15;
+  repeated google.protobuf.UInt64Value repeated_uint64 = 16;
+  repeated google.protobuf.BoolValue repeated_bool = 17;
+  repeated google.protobuf.StringValue repeated_string = 18;
+  repeated google.protobuf.BytesValue repeated_bytes = 19;
+
+  // Wrappers as map keys don't make sense, but we still need to make sure they
+  // work and don't crash.
+  map<int32, google.protobuf.DoubleValue> map_double = 21;
+  map<int32, google.protobuf.FloatValue> map_float = 22;
+  map<int32, google.protobuf.Int32Value> map_int32 = 23;
+  map<int32, google.protobuf.Int64Value> map_int64 = 24;
+  map<int32, google.protobuf.UInt32Value> map_uint32 = 25;
+  map<int32, google.protobuf.UInt64Value> map_uint64 = 26;
+  map<int32, google.protobuf.BoolValue> map_bool = 27;
+  map<int32, google.protobuf.StringValue> map_string = 28;
+  map<int32, google.protobuf.BytesValue> map_bytes = 29;
+
+  // Wrappers in oneofs don't make sense, but we still need to make sure they
+  // work and don't crash.
+  oneof wrapper_oneof {
+    google.protobuf.DoubleValue oneof_double = 31;
+    google.protobuf.FloatValue oneof_float = 32;
+    google.protobuf.Int32Value oneof_int32 = 33;
+    google.protobuf.Int64Value oneof_int64 = 34;
+    google.protobuf.UInt32Value oneof_uint32 = 35;
+    google.protobuf.UInt64Value oneof_uint64 = 36;
+    google.protobuf.BoolValue oneof_bool = 37;
+    google.protobuf.StringValue oneof_string = 38;
+    google.protobuf.BytesValue oneof_bytes = 39;
+    string oneof_plain_string = 101;
   }
 }
 

ruby/tests/basic_test_proto2.proto

@@ -133,7 +133,34 @@ message Wrapper {
   optional google.protobuf.BytesValue bytes = 9;
   optional string real_string = 100;
   oneof a_oneof {
-    string oneof_string = 10;
+    string string_in_oneof = 10;
+  }
+
+  // Repeated wrappers don't really make sense, but we still need to make sure
+  // they work and don't crash.
+  repeated google.protobuf.DoubleValue repeated_double = 11;
+  repeated google.protobuf.FloatValue repeated_float = 12;
+  repeated google.protobuf.Int32Value repeated_int32 = 13;
+  repeated google.protobuf.Int64Value repeated_int64 = 14;
+  repeated google.protobuf.UInt32Value repeated_uint32 = 15;
+  repeated google.protobuf.UInt64Value repeated_uint64 = 16;
+  repeated google.protobuf.BoolValue repeated_bool = 17;
+  repeated google.protobuf.StringValue repeated_string = 18;
+  repeated google.protobuf.BytesValue repeated_bytes = 19;
+
+  // Wrappers in oneofs don't make sense, but we still need to make sure they
+  // work and don't crash.
+  oneof wrapper_oneof {
+    google.protobuf.DoubleValue oneof_double = 31;
+    google.protobuf.FloatValue oneof_float = 32;
+    google.protobuf.Int32Value oneof_int32 = 33;
+    google.protobuf.Int64Value oneof_int64 = 34;
+    google.protobuf.UInt32Value oneof_uint32 = 35;
+    google.protobuf.UInt64Value oneof_uint64 = 36;
+    google.protobuf.BoolValue oneof_bool = 37;
+    google.protobuf.StringValue oneof_string = 38;
+    google.protobuf.BytesValue oneof_bytes = 39;
+    string oneof_plain_string = 101;
   }
 }