Migrate gaer::CelProtoWrapper::CreateMessage to just call the modern

eval/public/structs/BUILD

@@ -30,10 +30,17 @@ cc_library(
     deps = [
         ":cel_proto_wrap_util",
         ":proto_message_type_adapter",
+        ":trivial_legacy_type_info_internal",
+        "//common:value",
         "//eval/public:cel_value",
         "//eval/public:message_wrapper",
         "//internal:proto_time_encoding",
-        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_protobuf//:duration_cc_proto",
         "@com_google_protobuf//:protobuf",
         "@com_google_protobuf//:timestamp_cc_proto",

eval/public/structs/cel_proto_wrapper.cc

@@ -14,11 +14,21 @@
 
 #include "eval/public/structs/cel_proto_wrapper.h"
 
-#include "absl/types/optional.h"
+#include <optional>
+
+#include "absl/base/no_destructor.h"
+#include "absl/base/nullability.h"
+#include "absl/log/absl_check.h"
+#include "absl/log/absl_log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "common/legacy_value.h"
+#include "common/value.h"
 #include "eval/public/cel_value.h"
 #include "eval/public/message_wrapper.h"
 #include "eval/public/structs/cel_proto_wrap_util.h"
 #include "eval/public/structs/proto_message_type_adapter.h"
+#include "eval/public/structs/trivial_legacy_type_info_internal.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
@@ -27,9 +37,36 @@ namespace google::api::expr::runtime {
 
 namespace {
 
+using ::cel::interop_internal::TrivialTypeInfo;
 using ::google::protobuf::Arena;
 using ::google::protobuf::Descriptor;
+using ::google::protobuf::DescriptorPool;
 using ::google::protobuf::Message;
+using ::google::protobuf::MessageFactory;
+
+// Returns the arena for the given message, or the fallback arena if the
+// message does not have an arena.
+//
+// A global fallback arena is used to avoid allocation when the arena is not
+// specified. This is effectively a memory leak, but won't trigger leak
+// check analyzers.
+//
+// This emulates the old behavior of tolerating a nullptr arena without
+// triggering a crash.
+google::protobuf::Arena* GetArena(const Message* absl_nonnull message,
+                        google::protobuf::Arena* absl_nullable arena) {
+  if (arena != nullptr) {
+    return arena;
+  }
+  if (message->GetArena() != nullptr) {
+    return message->GetArena();
+  }
+  static absl::NoDestructor<google::protobuf::Arena> fallback_arena;
+  ABSL_LOG(WARNING)
+      << "CelValue: using fallback global arena for wrapping message: "
+      << message->GetTypeName();
+  return fallback_arena.get();
+}
 
 }  // namespace
 
@@ -38,14 +75,50 @@ CelValue CelProtoWrapper::InternalWrapMessage(const Message* message) {
       MessageWrapper(message, &GetGenericProtoTypeInfoInstance()));
 }
 
-// CreateMessage creates CelValue from google::protobuf::Message.
-// As some of CEL basic types are subclassing google::protobuf::Message,
-// this method contains type checking and downcasts.
-CelValue CelProtoWrapper::CreateMessage(const Message* value, Arena* arena) {
-  return internal::UnwrapMessageToValue(value, &InternalWrapMessage, arena);
+CelValue CelProtoWrapper::CreateMessage(
+    const Message* absl_nonnull value,
+    const google::protobuf::DescriptorPool* absl_nonnull pool,
+    MessageFactory* absl_nonnull factory, Arena* absl_nonnull arena) {
+  ABSL_DCHECK(value != nullptr);
+  if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) {
+    // This only happens for custom google::protobuf::Message subclasses that CEL can't
+    // support.
+    return CelValue::CreateMessageWrapper(
+        MessageWrapper(value, TrivialTypeInfo::GetInstance()));
+  }
+
+  auto modern_value =
+      cel::Value::WrapMessageUnsafe(value, pool, factory, arena);
+
+  absl::StatusOr<CelValue> cel_value = cel::LegacyValue(arena, modern_value);
+  if (!cel_value.ok()) {
+    // This only happens for custom google::protobuf::Message subclasses that CEL can't
+    // support.
+    auto* status =
+        google::protobuf::Arena::Create<absl::Status>(arena, cel_value.status());
+    return CelValue::CreateError(status);
+  }
+  return *cel_value;
+}
+
+CelValue CelProtoWrapper::CreateMessage(const Message* absl_nullable value,
+                                        Arena* absl_nullable arena) {
+  if (value == nullptr) {
+    return CelValue::CreateNull();
+  }
+
+  if (value->GetDescriptor() == nullptr || value->GetReflection() == nullptr) {
+    // This only happens for custom messages subclasses that CEL can't support.
+    return CelValue::CreateMessageWrapper(
+        MessageWrapper(value, TrivialTypeInfo::GetInstance()));
+  }
+  const auto* pool = value->GetDescriptor()->file()->pool();
+  auto* factory = value->GetReflection()->GetMessageFactory();
+  arena = GetArena(value, arena);
+  return CreateMessage(value, pool, factory, arena);
 }
 
-absl::optional<CelValue> CelProtoWrapper::MaybeWrapValue(
+std::optional<CelValue> CelProtoWrapper::MaybeWrapValue(
     const Descriptor* descriptor, google::protobuf::MessageFactory* factory,
     const CelValue& value, Arena* arena) {
   const Message* msg =

eval/public/structs/cel_proto_wrapper.h

@@ -1,9 +1,11 @@
 #ifndef THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_
 #define THIRD_PARTY_CEL_CPP_EVAL_PUBLIC_STRUCTS_CEL_PROTO_WRAPPER_H_
 
+#include <optional>
+
 #include "google/protobuf/duration.pb.h"
 #include "google/protobuf/timestamp.pb.h"
-#include "absl/types/optional.h"
+#include "absl/base/nullability.h"
 #include "eval/public/cel_value.h"
 #include "internal/proto_time_encoding.h"
 #include "google/protobuf/arena.h"
@@ -17,8 +19,19 @@ class CelProtoWrapper {
   // CreateMessage creates CelValue from google::protobuf::Message.
   // As some of CEL basic types are subclassing google::protobuf::Message,
   // this method contains type checking and downcasts.
-  static CelValue CreateMessage(const google::protobuf::Message* value,
-                                google::protobuf::Arena* arena);
+  static CelValue CreateMessage(const google::protobuf::Message* absl_nonnull value,
+                                const google::protobuf::DescriptorPool* absl_nonnull pool,
+                                google::protobuf::MessageFactory* absl_nonnull factory,
+                                google::protobuf::Arena* absl_nonnull arena);
+
+  // Prefer using the overload that takes an explicit descriptor pool and
+  // message factory instead. This overload will use the ones associated with
+  // the value.
+  //
+  // For backward compatibility, nullptr message is allowed and will result in
+  // the CEL null_type value.
+  static CelValue CreateMessage(const google::protobuf::Message* absl_nullable value,
+                                google::protobuf::Arena* absl_nullable arena);
 
   // Internal utility for creating a CelValue wrapping a user defined type.
   // Assumes that the message has been properly unpacked.
@@ -43,7 +56,7 @@ class CelProtoWrapper {
   // message to native CelValue representation during a protobuf field read.
   // Just as CreateMessage should only be used when reading protobuf values,
   // MaybeWrapValue should only be used when assigning protobuf fields.
-  static absl::optional<CelValue> MaybeWrapValue(
+  static std::optional<CelValue> MaybeWrapValue(
       const google::protobuf::Descriptor* descriptor, google::protobuf::MessageFactory* factory,
       const CelValue& value, google::protobuf::Arena* arena);
 };

eval/public/structs/cel_proto_wrapper_test.cc

@@ -104,7 +104,7 @@ class CelProtoWrapperTest : public ::testing::Test {
 
     T dyn_value;
     CelValue cel_dyn_value =
-        CelProtoWrapper::CreateMessage(ReflectedCopy(message).get(), arena());
+        CelProtoWrapper::CreateMessage(ReflectedCopy(message), arena());
     EXPECT_THAT(cel_dyn_value.type(), Eq(cel_value.type()));
     EXPECT_TRUE(cel_dyn_value.GetValue(&dyn_value));
     EXPECT_THAT(value, Eq(dyn_value));
@@ -121,10 +121,9 @@ class CelProtoWrapperTest : public ::testing::Test {
     EXPECT_THAT(cel_value.MessageOrDie(), testutil::EqualsProto(*result));
   }
 
-  std::unique_ptr<google::protobuf::Message> ReflectedCopy(
-      const google::protobuf::Message& message) {
-    std::unique_ptr<google::protobuf::Message> dynamic_value(
-        factory_.GetPrototype(message.GetDescriptor())->New());
+  google::protobuf::Message* ReflectedCopy(const google::protobuf::Message& message) {
+    google::protobuf::Message* dynamic_value =
+        factory_.GetPrototype(message.GetDescriptor())->New(&arena_);
     dynamic_value->CopyFrom(message);
     return dynamic_value;
   }
@@ -213,7 +212,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueNull) {
   value_msg.set_null_value(protobuf::NULL_VALUE);
 
   CelValue value =
-      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
+      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
   EXPECT_TRUE(value.IsNull());
 }
 
@@ -314,8 +313,8 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicStruct) {
   const std::string kFieldBool = "field_bool";
   (*struct_msg.mutable_fields())[kFieldInt].set_number_value(1.);
   (*struct_msg.mutable_fields())[kFieldBool].set_bool_value(true);
-  CelValue value =
-      CelProtoWrapper::CreateMessage(ReflectedCopy(struct_msg).get(), arena());
+  auto reflected_copy = ReflectedCopy(struct_msg);
+  CelValue value = CelProtoWrapper::CreateMessage(reflected_copy, arena());
   EXPECT_TRUE(value.IsMap());
   const CelMap* cel_map = value.MapOrDie();
   ASSERT_TRUE(cel_map != nullptr);
@@ -355,7 +354,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueStruct) {
       .set_number_value(2);
 
   CelValue value =
-      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
+      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
   EXPECT_TRUE(value.IsMap());
   EXPECT_TRUE(
       (*value.MapOrDie())[CelValue::CreateString(&kField1)].has_value());
@@ -398,7 +397,7 @@ TEST_F(CelProtoWrapperTest, UnwrapDynamicValueListValue) {
   value_msg.mutable_list_value()->add_values()->set_number_value(2.);
 
   CelValue value =
-      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg).get(), arena());
+      CelProtoWrapper::CreateMessage(ReflectedCopy(value_msg), arena());
   EXPECT_TRUE(value.IsList());
   EXPECT_THAT((*value.ListOrDie())[0].DoubleOrDie(), testing::DoubleEq(1));
   EXPECT_THAT((*value.ListOrDie())[1].DoubleOrDie(), testing::DoubleEq(2));
@@ -426,6 +425,47 @@ TEST_F(CelProtoWrapperTest, UnwrapInvalidAny) {
   ASSERT_TRUE(CelProtoWrapper::CreateMessage(&any, arena()).IsError());
 }
 
+TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactory) {
+  TestMessage test_message;
+  test_message.set_string_value("test");
+
+  CelValue value = CelProtoWrapper::CreateMessage(
+      &test_message, google::protobuf::DescriptorPool::generated_pool(),
+      google::protobuf::MessageFactory::generated_factory(), arena());
+  ASSERT_TRUE(value.IsMessage());
+  EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message));
+}
+
+TEST_F(CelProtoWrapperTest, CreateMessageExplicitPoolAndFactoryUnpackAny) {
+  TestMessage test_message;
+  test_message.set_string_value("test");
+
+  Any any;
+  any.PackFrom(test_message);
+
+  google::protobuf::DynamicMessageFactory factory(
+      google::protobuf::DescriptorPool::generated_pool());
+  CelValue value = CelProtoWrapper::CreateMessage(
+      &any, google::protobuf::DescriptorPool::generated_pool(), &factory, arena());
+  ASSERT_TRUE(value.IsMessage());
+  EXPECT_THAT(value.MessageOrDie(), testutil::EqualsProto(test_message));
+}
+
+TEST_F(CelProtoWrapperTest,
+       CreateMessageExplicitPoolAndFactoryUnpackAnyNotFound) {
+  TestMessage test_message;
+  test_message.set_string_value("test");
+
+  Any any;
+  any.PackFrom(test_message);
+
+  google::protobuf::DescriptorPool empty_pool;
+  google::protobuf::DynamicMessageFactory factory(&empty_pool);
+  CelValue value =
+      CelProtoWrapper::CreateMessage(&any, &empty_pool, &factory, arena());
+  EXPECT_TRUE(value.IsError());
+}
+
 // Test support of google.protobuf.<Type>Value wrappers in CelValue.
 TEST_F(CelProtoWrapperTest, UnwrapBoolWrapper) {
   bool value = true;