Generate extra enum method helpers for Ruby (#5670)

* example with extra enum method

* update expected test output

* slight simplification

* add test for generated enum helpers

* move const helpers to c extension

* more explicit test

* more explicit test

* indent

* add foo test

* add check for _const suffix

ruby/ext/google/protobuf_c/message.c

@@ -118,7 +118,8 @@ enum {
   METHOD_GETTER = 1,
   METHOD_SETTER = 2,
   METHOD_CLEAR = 3,
-  METHOD_PRESENCE = 4
+  METHOD_PRESENCE = 4,
+  METHOD_ENUM_GETTER = 5
 };
 
 static int extract_method_call(VALUE method_name, MessageHeader* self,
@@ -153,9 +154,34 @@ static int extract_method_call(VALUE method_name, MessageHeader* self,
     accessor_type = METHOD_GETTER;
   }
 
+  bool has_field = upb_msgdef_lookupname(self->descriptor->msgdef, name, name_len,
+			                                   &test_f, &test_o);
+
+  // Look for enum accessor of the form <enum_name>_const
+  if (!has_field && accessor_type == METHOD_GETTER &&
+      name_len > 6 && strncmp(name + name_len - 6, "_const", 6) == 0) {
+
+    // Find enum field name
+    char enum_name[name_len - 5];
+    strncpy(enum_name, name, name_len - 6);
+    enum_name[name_len - 4] = '\0';
+
+    // Check if enum field exists
+    const upb_oneofdef* test_o_enum;
+    const upb_fielddef* test_f_enum;
+    if (upb_msgdef_lookupname(self->descriptor->msgdef, enum_name, name_len - 6,
+			                        &test_f_enum, &test_o_enum) &&
+        upb_fielddef_type(test_f_enum) == UPB_TYPE_ENUM) {
+      // It does exist!
+      has_field = true;
+      accessor_type = METHOD_ENUM_GETTER;
+      test_o = test_o_enum;
+      test_f = test_f_enum;
+    }
+  }
+
   // Verify the name corresponds to a oneof or field in this message.
-  if (!upb_msgdef_lookupname(self->descriptor->msgdef, name, name_len,
-			     &test_f, &test_o)) {
+  if (!has_field) {
     return METHOD_UNKNOWN;
   }
 
@@ -248,6 +274,25 @@ VALUE Message_method_missing(int argc, VALUE* argv, VALUE _self) {
     return Qnil;
   } else if (accessor_type == METHOD_PRESENCE) {
     return layout_has(self->descriptor->layout, Message_data(self), f);
+  } else if (accessor_type == METHOD_ENUM_GETTER) {
+    VALUE enum_type = field_type_class(f);
+    VALUE method = rb_intern("const_get");
+    VALUE raw_value = layout_get(self->descriptor->layout, Message_data(self), f);
+
+    // Map repeated fields to a new type with ints
+    if (upb_fielddef_label(f) == UPB_LABEL_REPEATED) {
+      int array_size = FIX2INT(rb_funcall(raw_value, rb_intern("length"), 0));
+      VALUE array_args[1] = { ID2SYM(rb_intern("int64")) };
+      VALUE array = rb_class_new_instance(1, array_args, CLASS_OF(raw_value));
+      for (int i = 0; i < array_size; i++) {
+        VALUE entry = rb_funcall(enum_type, method, 1, rb_funcall(raw_value,
+                                 rb_intern("at"), 1, INT2NUM(i)));
+        rb_funcall(array, rb_intern("push"), 1, entry);
+      }
+      return array;
+    }
+    // Convert the value for singular fields
+    return rb_funcall(enum_type, method, 1, raw_value);
   } else {
     return layout_get(self->descriptor->layout, Message_data(self), f);
   }

ruby/tests/basic_test.proto

@@ -110,6 +110,15 @@ message Outer {
 message Inner {
 }
 
+message Enumer {
+  TestEnum optional_enum = 1;
+  repeated TestEnum repeated_enum = 2;
+  string a_const = 3;
+  oneof a_oneof {
+    string str = 10;
+    TestEnum const = 11;
+  }
+}
 
 message MyRepeatedStruct {
   repeated MyStruct structs = 1;

ruby/tests/basic_test_proto2.proto

@@ -118,6 +118,16 @@ message OneofMessage {
   }
 }
 
+message Enumer {
+  optional TestEnum optional_enum = 11;
+  repeated TestEnum repeated_enum = 22;
+  optional string a_const = 3;
+  oneof a_oneof {
+    string str = 100;
+    TestEnum const = 101;
+  }
+}
+
 message MyRepeatedStruct {
   repeated MyStruct structs = 1;
 }

ruby/tests/common_tests.rb

@@ -708,6 +708,73 @@ module CommonTests
     assert proto_module::TestEnum::resolve(:C) == 3
   end
 
+  def test_enum_const_get_helpers
+    m = proto_module::TestMessage.new
+    assert_equal proto_module::TestEnum::Default, m.optional_enum_const
+    assert_equal proto_module::TestEnum.const_get(:Default), m.optional_enum_const
+
+    m = proto_module::TestMessage.new({optional_enum: proto_module::TestEnum::A})
+    assert_equal proto_module::TestEnum::A, m.optional_enum_const
+    assert_equal proto_module::TestEnum.const_get(:A), m.optional_enum_const
+
+    m = proto_module::TestMessage.new({optional_enum: proto_module::TestEnum::B})
+    assert_equal proto_module::TestEnum::B, m.optional_enum_const
+    assert_equal proto_module::TestEnum.const_get(:B), m.optional_enum_const
+
+    m = proto_module::TestMessage.new({optional_enum: proto_module::TestEnum::C})
+    assert_equal proto_module::TestEnum::C, m.optional_enum_const
+    assert_equal proto_module::TestEnum.const_get(:C), m.optional_enum_const
+
+    m = proto_module::TestMessage2.new({foo: 2})
+    assert_equal 2, m.foo
+    assert_raise(NoMethodError) { m.foo_ }
+    assert_raise(NoMethodError) { m.foo_X }
+    assert_raise(NoMethodError) { m.foo_XX }
+    assert_raise(NoMethodError) { m.foo_XXX }
+    assert_raise(NoMethodError) { m.foo_XXXX }
+    assert_raise(NoMethodError) { m.foo_XXXXX }
+    assert_raise(NoMethodError) { m.foo_XXXXXX }
+
+    m = proto_module::Enumer.new({optional_enum: :B})
+    assert_equal :B, m.optional_enum
+    assert_raise(NoMethodError) { m.optional_enum_ }
+    assert_raise(NoMethodError) { m.optional_enum_X }
+    assert_raise(NoMethodError) { m.optional_enum_XX }
+    assert_raise(NoMethodError) { m.optional_enum_XXX }
+    assert_raise(NoMethodError) { m.optional_enum_XXXX }
+    assert_raise(NoMethodError) { m.optional_enum_XXXXX }
+    assert_raise(NoMethodError) { m.optional_enum_XXXXXX }
+  end
+
+  def test_enum_getter
+    m = proto_module::Enumer.new(:optional_enum => :B, :repeated_enum => [:A, :C])
+
+    assert_equal :B, m.optional_enum
+    assert_equal 2, m.optional_enum_const
+    assert_equal proto_module::TestEnum::B, m.optional_enum_const
+    assert_equal [:A, :C], m.repeated_enum
+    assert_equal [1, 3], m.repeated_enum_const
+    assert_equal [proto_module::TestEnum::A, proto_module::TestEnum::C], m.repeated_enum_const
+  end
+
+  def test_enum_getter_oneof
+    m = proto_module::Enumer.new(:const => :C)
+
+    assert_equal :C, m.const
+    assert_equal 3, m.const_const
+    assert_equal proto_module::TestEnum::C, m.const_const
+  end
+
+  def test_enum_getter_only_enums
+    m = proto_module::Enumer.new(:optional_enum => :B, :a_const => 'thing')
+
+    assert_equal 'thing', m.a_const
+    assert_equal :B, m.optional_enum
+
+    assert_raise(NoMethodError) { m.a }
+    assert_raise(NoMethodError) { m.a_const_const }
+  end
+  
   def test_repeated_push
     m = proto_module::TestMessage.new