Implement singular enum accessors

PiperOrigin-RevId: 596984036

rust/test/nested.proto

@@ -15,6 +15,11 @@ message Outer {
       optional bool flag = 1;
     }
 
+    enum InnerEnum {
+      INNER_ENUM_UNSPECIFIED = 0;
+      INNER_ENUM_FOO = 1;
+    }
+
     optional double double = 1;
     optional float float = 2;
     optional int32 int32 = 3;
@@ -31,6 +36,7 @@ message Outer {
     optional string string = 14;
     optional bytes bytes = 15;
     optional InnerSubMsg innersubmsg = 16;
+    optional InnerEnum inner_enum = 17;
 
     message SuperInner {
       message DuperInner {

rust/test/shared/accessors_proto3_test.rs

@@ -11,7 +11,7 @@ use googletest::prelude::*;
 use matchers::{is_set, is_unset};
 use protobuf::Optional;
 use unittest_proto3::proto3_unittest::{TestAllTypes, TestAllTypes_};
-use unittest_proto3_optional::proto2_unittest::TestProto3Optional;
+use unittest_proto3_optional::proto2_unittest::{TestProto3Optional, TestProto3Optional_};
 
 #[test]
 fn test_fixed32_accessors() {
@@ -184,6 +184,72 @@ fn test_optional_string_accessors() {
     assert_that!(msg.optional_string_opt(), eq(Optional::Set("".into())));
 }
 
+#[test]
+fn test_nested_enum_accessors() {
+    use TestAllTypes_::NestedEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Zero));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Zero));
+
+    msg.optional_nested_enum_mut().set(NestedEnum::Baz);
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Baz));
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Baz));
+
+    msg.optional_nested_enum_mut().set(NestedEnum::default());
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Zero));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Zero));
+}
+
+#[test]
+fn test_optional_nested_enum_accessors() {
+    use TestProto3Optional_::NestedEnum;
+
+    let mut msg = TestProto3Optional::new();
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Unspecified));
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Unspecified)));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Unspecified));
+
+    msg.optional_nested_enum_mut().set(NestedEnum::Baz);
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Baz));
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Set(NestedEnum::Baz)));
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Baz));
+
+    msg.optional_nested_enum_mut().set(NestedEnum::default());
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Unspecified));
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Set(NestedEnum::Unspecified)));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Unspecified));
+
+    msg.optional_nested_enum_mut().clear();
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Unspecified));
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Unspecified)));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Unspecified));
+
+    let mut field_mut = msg.optional_nested_enum_mut().or_default();
+    assert_that!(field_mut.get(), eq(NestedEnum::Unspecified));
+    field_mut.set(NestedEnum::Bar);
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Bar));
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Set(NestedEnum::Bar)));
+    assert_that!(msg.optional_nested_enum_mut().get(), eq(NestedEnum::Bar));
+}
+
+#[test]
+fn test_foreign_enum_accessors() {
+    use unittest_proto3::proto3_unittest::ForeignEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignZero));
+    assert_that!(msg.optional_foreign_enum_mut().get(), eq(ForeignEnum::ForeignZero));
+
+    msg.optional_foreign_enum_mut().set(ForeignEnum::ForeignBaz);
+    assert_that!(msg.optional_foreign_enum_mut().get(), eq(ForeignEnum::ForeignBaz));
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignBaz));
+
+    msg.optional_foreign_enum_mut().set(ForeignEnum::default());
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignZero));
+    assert_that!(msg.optional_foreign_enum_mut().get(), eq(ForeignEnum::ForeignZero));
+}
+
 #[test]
 fn test_oneof_accessors() {
     use TestAllTypes_::OneofField::*;

rust/test/shared/accessors_test.rs

@@ -668,7 +668,7 @@ fn test_nonempty_default_string_accessors() {
 
 #[test]
 fn test_singular_msg_field() {
-    use crate::TestAllTypes_::{NestedMessageMut, NestedMessageView};
+    use TestAllTypes_::{NestedMessageMut, NestedMessageView};
 
     let mut msg = TestAllTypes::new();
     let msg_view: NestedMessageView = msg.optional_nested_message();
@@ -680,6 +680,144 @@ fn test_singular_msg_field() {
     assert_that!(msg_mut.bb(), eq(0));
 }
 
+#[test]
+fn test_optional_nested_enum_accessors() {
+    use TestAllTypes_::NestedEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Foo)));
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Foo));
+
+    msg.optional_nested_enum_mut().set(NestedEnum::Neg);
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Set(NestedEnum::Neg)));
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Neg));
+
+    msg.optional_nested_enum_mut().clear();
+    assert_that!(msg.optional_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Foo)));
+    assert_that!(msg.optional_nested_enum(), eq(NestedEnum::Foo));
+}
+
+#[test]
+fn test_default_nested_enum_accessors() {
+    use TestAllTypes_::NestedEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.default_nested_enum(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().get(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Bar)));
+
+    msg.default_nested_enum_mut().set(NestedEnum::Baz);
+    assert_that!(msg.default_nested_enum(), eq(NestedEnum::Baz));
+    assert_that!(msg.default_nested_enum_mut().get(), eq(NestedEnum::Baz));
+    assert_that!(msg.default_nested_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_nested_enum_opt(), eq(Optional::Set(NestedEnum::Baz)));
+
+    msg.default_nested_enum_mut().clear();
+    assert_that!(msg.default_nested_enum(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().get(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_nested_enum_opt(), eq(Optional::Unset(NestedEnum::Bar)));
+
+    msg.default_nested_enum_mut().or_default();
+    assert_that!(msg.default_nested_enum(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().get(), eq(NestedEnum::Bar));
+    assert_that!(msg.default_nested_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_nested_enum_opt(), eq(Optional::Set(NestedEnum::Bar)));
+}
+
+#[test]
+fn test_optional_foreign_enum_accessors() {
+    use unittest_proto::proto2_unittest::ForeignEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.optional_foreign_enum_opt(), eq(Optional::Unset(ForeignEnum::ForeignFoo)));
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignFoo));
+
+    msg.optional_foreign_enum_mut().set(ForeignEnum::ForeignBax);
+    assert_that!(msg.optional_foreign_enum_opt(), eq(Optional::Set(ForeignEnum::ForeignBax)));
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignBax));
+
+    msg.optional_foreign_enum_mut().clear();
+    assert_that!(msg.optional_foreign_enum_opt(), eq(Optional::Unset(ForeignEnum::ForeignFoo)));
+    assert_that!(msg.optional_foreign_enum(), eq(ForeignEnum::ForeignFoo));
+}
+
+#[test]
+fn test_default_foreign_enum_accessors() {
+    use unittest_proto::proto2_unittest::ForeignEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.default_foreign_enum(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().get(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_foreign_enum_opt(), eq(Optional::Unset(ForeignEnum::ForeignBar)));
+
+    msg.default_foreign_enum_mut().set(ForeignEnum::ForeignBaz);
+    assert_that!(msg.default_foreign_enum(), eq(ForeignEnum::ForeignBaz));
+    assert_that!(msg.default_foreign_enum_mut().get(), eq(ForeignEnum::ForeignBaz));
+    assert_that!(msg.default_foreign_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_foreign_enum_opt(), eq(Optional::Set(ForeignEnum::ForeignBaz)));
+
+    msg.default_foreign_enum_mut().clear();
+    assert_that!(msg.default_foreign_enum(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().get(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_foreign_enum_opt(), eq(Optional::Unset(ForeignEnum::ForeignBar)));
+
+    msg.default_foreign_enum_mut().or_default();
+    assert_that!(msg.default_foreign_enum(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().get(), eq(ForeignEnum::ForeignBar));
+    assert_that!(msg.default_foreign_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_foreign_enum_opt(), eq(Optional::Set(ForeignEnum::ForeignBar)));
+}
+
+#[test]
+fn test_optional_import_enum_accessors() {
+    use unittest_proto::proto2_unittest_import::ImportEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.optional_import_enum_opt(), eq(Optional::Unset(ImportEnum::ImportFoo)));
+    assert_that!(msg.optional_import_enum(), eq(ImportEnum::ImportFoo));
+
+    msg.optional_import_enum_mut().set(ImportEnum::ImportBar);
+    assert_that!(msg.optional_import_enum_opt(), eq(Optional::Set(ImportEnum::ImportBar)));
+    assert_that!(msg.optional_import_enum(), eq(ImportEnum::ImportBar));
+
+    msg.optional_import_enum_mut().clear();
+    assert_that!(msg.optional_import_enum_opt(), eq(Optional::Unset(ImportEnum::ImportFoo)));
+    assert_that!(msg.optional_import_enum(), eq(ImportEnum::ImportFoo));
+}
+
+#[test]
+fn test_default_import_enum_accessors() {
+    use unittest_proto::proto2_unittest_import::ImportEnum;
+
+    let mut msg = TestAllTypes::new();
+    assert_that!(msg.default_import_enum(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().get(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_import_enum_opt(), eq(Optional::Unset(ImportEnum::ImportBar)));
+
+    msg.default_import_enum_mut().set(ImportEnum::ImportBaz);
+    assert_that!(msg.default_import_enum(), eq(ImportEnum::ImportBaz));
+    assert_that!(msg.default_import_enum_mut().get(), eq(ImportEnum::ImportBaz));
+    assert_that!(msg.default_import_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_import_enum_opt(), eq(Optional::Set(ImportEnum::ImportBaz)));
+
+    msg.default_import_enum_mut().clear();
+    assert_that!(msg.default_import_enum(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().get(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().is_set(), eq(false));
+    assert_that!(msg.default_import_enum_opt(), eq(Optional::Unset(ImportEnum::ImportBar)));
+
+    msg.default_import_enum_mut().or_default();
+    assert_that!(msg.default_import_enum(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().get(), eq(ImportEnum::ImportBar));
+    assert_that!(msg.default_import_enum_mut().is_set(), eq(true));
+    assert_that!(msg.default_import_enum_opt(), eq(Optional::Set(ImportEnum::ImportBar)));
+}
+
 #[test]
 fn test_oneof_accessors() {
     use TestAllTypes_::OneofField::*;

rust/test/shared/simple_nested_test.rs

@@ -9,6 +9,7 @@ use googletest::prelude::*;
 use nested_proto::nest::Outer;
 use nested_proto::nest::Outer_::InnerMut;
 use nested_proto::nest::Outer_::InnerView;
+use nested_proto::nest::Outer_::Inner_::InnerEnum;
 
 #[test]
 fn test_deeply_nested_message() {
@@ -22,11 +23,12 @@ fn test_deeply_nested_message() {
 
 #[test]
 fn test_deeply_nested_enum() {
-    let deep = nested_proto::nest::Outer_::Inner_::SuperInner_::DuperInner_::EvenMoreInner_
-            ::JustWayTooInner::default();
+    use nested_proto::nest::Outer_::Inner_::SuperInner_::DuperInner_::EvenMoreInner_::JustWayTooInner;
+    let deep = JustWayTooInner::default();
     assert_that!(i32::from(deep), eq(0));
 
-    // TODO: Test deeply nested enum field access.
+    let outer_msg = Outer::new();
+    assert_that!(outer_msg.deep_enum(), eq(JustWayTooInner::Unspecified));
 }
 
 #[test]
@@ -49,6 +51,7 @@ fn test_nested_views() {
     assert_that!(*inner_msg.string().as_bytes(), empty());
     assert_that!(*inner_msg.bytes(), empty());
     assert_that!(inner_msg.innersubmsg().flag(), eq(false));
+    assert_that!(inner_msg.inner_enum(), eq(InnerEnum::Unspecified));
 }
 
 #[test]
@@ -61,7 +64,7 @@ fn test_nested_muts() {
     // new:
     //   set_and_test_mut!(inner_msg => double_mut, 543.21);
     macro_rules! set_and_test_mut {
-    ( $a:expr => $($target_mut:ident, $val:literal;)* ) => {
+    ( $a:expr => $($target_mut:ident, $val:expr;)* ) => {
         $(
           $a.$target_mut().set($val);
           assert_that!($a.$target_mut().get(), eq($val));
@@ -86,7 +89,8 @@ fn test_nested_muts() {
             fixed64(): eq(0),
             sfixed32(): eq(0),
             sfixed64(): eq(0),
-            bool(): eq(false)
+            bool(): eq(false),
+            inner_enum(): eq(InnerEnum::Unspecified)
         })
     );
     assert_that!(inner_msg.string_mut().get(), eq(""));
@@ -106,6 +110,7 @@ fn test_nested_muts() {
         bool_mut, true;
         string_mut, "hi";
         bytes_mut, b"bye";
+        inner_enum_mut, InnerEnum::Foo;
     );
 }
 

src/google/protobuf/compiler/rust/accessors/accessors.cc

@@ -63,6 +63,18 @@ std::unique_ptr<AccessorGenerator> AccessorGeneratorFor(
         return std::make_unique<RepeatedScalar>();
       }
       return std::make_unique<SingularScalar>();
+    case FieldDescriptor::TYPE_ENUM:
+      // TODO: support enums which are defined in other crates.
+      if (!IsInCurrentlyGeneratingCrate(ctx, *field.enum_type())) {
+        return std::make_unique<UnsupportedField>(
+            "enum fields that are imported from another proto_library"
+            " (defined in a separate Rust crate) are not supported");
+      }
+      if (field.is_repeated()) {
+        return std::make_unique<UnsupportedField>(
+            "repeated enum not supported");
+      }
+      return std::make_unique<SingularScalar>();
     case FieldDescriptor::TYPE_BYTES:
     case FieldDescriptor::TYPE_STRING:
       if (field.is_repeated()) {
@@ -73,17 +85,14 @@ std::unique_ptr<AccessorGenerator> AccessorGeneratorFor(
       if (field.is_repeated()) {
         return std::make_unique<UnsupportedField>("repeated msg not supported");
       }
-      if (!ctx.generator_context().is_file_in_current_crate(
-              *field.message_type()->file())) {
+      // TODO: support messages which are defined in other crates.
+      if (!IsInCurrentlyGeneratingCrate(ctx, *field.message_type())) {
         return std::make_unique<UnsupportedField>(
             "message fields that are imported from another proto_library"
             " (defined in a separate Rust crate) are not supported");
       }
       return std::make_unique<SingularMessage>();
 
-    case FieldDescriptor::TYPE_ENUM:
-      return std::make_unique<UnsupportedField>("enum not supported");
-
     case FieldDescriptor::TYPE_GROUP:
       return std::make_unique<UnsupportedField>("group not supported");
   }

src/google/protobuf/compiler/rust/accessors/singular_scalar.cc

@@ -5,6 +5,8 @@
 // license that can be found in the LICENSE file or at
 // https://developers.google.com/open-source/licenses/bsd
 
+#include <string>
+
 #include "absl/strings/string_view.h"
 #include "google/protobuf/compiler/cpp/helpers.h"
 #include "google/protobuf/compiler/rust/accessors/accessor_generator.h"
@@ -117,6 +119,13 @@ void SingularScalar::InMsgImpl(Context& ctx,
 
 void SingularScalar::InExternC(Context& ctx,
                                const FieldDescriptor& field) const {
+  // In order to soundly pass a Rust type to C/C++ as a function argument,
+  // the types must be FFI-compatible.
+  // This requires special consideration for enums, which aren't trivial
+  // primitive types. Rust protobuf enums are defined as `#[repr(transparent)]`
+  // over `i32`, making them ABI-compatible with `int32_t`.
+  // Upb defines enum thunks as taking `int32_t`, and so we can pass Rust enums
+  // directly to thunks without any cast.
   ctx.Emit({{"Scalar", RsTypePath(ctx, field)},
             {"hazzer_thunk", ThunkName(ctx, field, "has")},
             {"getter_thunk", ThunkName(ctx, field, "get")},
@@ -141,8 +150,21 @@ void SingularScalar::InExternC(Context& ctx,
 
 void SingularScalar::InThunkCc(Context& ctx,
                                const FieldDescriptor& field) const {
+  std::string scalar;
+  auto enum_ = field.enum_type();
+  if (enum_ != nullptr) {
+    // The C++ runtime defines its thunks as receiving enum types.
+    // This is fine since:
+    // - the C++ runtime represents enums as `int`
+    // - the C++ runtime guarantees `int` is a `int32_t`.
+    // - Rust gencode defines enums as `#[repr(transparent)]` over `i32`.
+    scalar = cpp::QualifiedClassName(enum_);
+  } else {
+    scalar = cpp::PrimitiveTypeName(field.cpp_type());
+  }
+
   ctx.Emit({{"field", cpp::FieldName(&field)},
-            {"Scalar", cpp::PrimitiveTypeName(field.cpp_type())},
+            {"Scalar", scalar},
             {"QualifiedMsg", cpp::QualifiedClassName(field.containing_type())},
             {"hazzer_thunk", ThunkName(ctx, field, "has")},
             {"getter_thunk", ThunkName(ctx, field, "get")},

src/google/protobuf/compiler/rust/context.cc

@@ -76,6 +76,10 @@ bool IsInCurrentlyGeneratingCrate(Context& ctx, const Descriptor& message) {
   return IsInCurrentlyGeneratingCrate(ctx, *message.file());
 }
 
+bool IsInCurrentlyGeneratingCrate(Context& ctx, const EnumDescriptor& enum_) {
+  return IsInCurrentlyGeneratingCrate(ctx, *enum_.file());
+}
+
 }  // namespace rust
 }  // namespace compiler
 }  // namespace protobuf

src/google/protobuf/compiler/rust/context.h

@@ -117,6 +117,7 @@ class Context {
 
 bool IsInCurrentlyGeneratingCrate(Context& ctx, const FileDescriptor& file);
 bool IsInCurrentlyGeneratingCrate(Context& ctx, const Descriptor& message);
+bool IsInCurrentlyGeneratingCrate(Context& ctx, const EnumDescriptor& enum_);
 
 }  // namespace rust
 }  // namespace compiler

src/google/protobuf/compiler/rust/message.cc

@@ -284,6 +284,12 @@ void GetterForViewOrMut(Context& ctx, const FieldDescriptor& field,
           ? "unsafe { __pb::ProtoStr::from_utf8_unchecked(res).into() }"
           : "res";
 
+  // TODO: support enums which are defined in other crates.
+  auto enum_ = field.enum_type();
+  if (enum_ != nullptr && !IsInCurrentlyGeneratingCrate(ctx, *enum_)) {
+    return;
+  }
+
   ctx.Emit({{"field", fieldName},
             {"getter_thunk", getter_thunk},
             {"setter_thunk", setter_thunk},
@@ -335,9 +341,7 @@ void AccessorsForViewOrMut(Context& ctx, const Descriptor& msg, bool is_mut) {
     // TODO - add cord support
     if (field.options().has_ctype()) continue;
     // TODO
-    if (field.type() == FieldDescriptor::TYPE_ENUM ||
-        field.type() == FieldDescriptor::TYPE_GROUP)
-      continue;
+    if (field.type() == FieldDescriptor::TYPE_GROUP) continue;
     GetterForViewOrMut(ctx, field, is_mut);
     ctx.printer().PrintRaw("\n");
   }