Consistency fixes for legacy adapter types:

- make direct CelMap::Get calls return error as the legacy implementations did
- avoid copying messages already marked as unsafe when converting to legacy

PiperOrigin-RevId: 972686690

Author: jnthntatum

common/legacy_value.cc

@@ -260,6 +260,11 @@ CelValue LegacyTrivialStructValue(google::protobuf::Arena* absl_nonnull arena,
   }
   if (auto parsed_message_value = value.AsParsedMessage();
       parsed_message_value) {
+    if (interop_internal::IsUnsafeParsedMessageValue(*parsed_message_value)) {
+      return CelValue::CreateMessageWrapper(
+          AsMessageWrapper(cel::to_address(*parsed_message_value),
+                           &GetGenericProtoTypeInfoInstance()));
+    }
     auto maybe_cloned = parsed_message_value->Clone(arena);
     return CelValue::CreateMessageWrapper(MessageWrapper(
         cel::to_address(maybe_cloned), &GetGenericProtoTypeInfoInstance()));

common/legacy_value.h

@@ -65,6 +65,10 @@ class MessageFactory;
 
 namespace cel::interop_internal {
 
+inline bool IsUnsafeParsedMessageValue(const cel::ParsedMessageValue& value) {
+  return value.is_unsafe();
+}
+
 // Returns the underlying `google::protobuf::Message` of a `cel::Value` if it is a legacy
 // message with the default type info, or `nullptr` otherwise.
 const google::protobuf::Message* absl_nullable GetLegacyMessage(const Value& value);

common/values/legacy_map_value.cc

@@ -22,7 +22,6 @@
 #include <utility>
 
 #include "absl/base/nullability.h"
-#include "absl/base/optimization.h"
 #include "absl/log/absl_check.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
@@ -218,21 +217,33 @@ class LegacyParsedMapFieldMapValue final
     if (arena == nullptr) {
       arena = arena_;
     }
-    if (auto status =
-            google::api::expr::runtime::CelValue::CheckMapKeyType(key);
-        !status.ok()) {
-      status.IgnoreError();
-      return std::nullopt;
-    }
     Value modern_key;
-    if (ABSL_PREDICT_FALSE(!ModernValue(arena, key, modern_key).ok())) {
+    if (!ModernValue(arena, key, modern_key).ok()) {
+      // Legacy to modern should succeed for a valid CelValue.
       return std::nullopt;
     }
     Value modern_val;
-    auto status_or_found =
-        Find(modern_key, google::protobuf::DescriptorPool::generated_pool(),
-             google::protobuf::MessageFactory::generated_factory(), arena, &modern_val);
-    if (!status_or_found.ok() || !*status_or_found) {
+    // Call custom map Find directly. MapValue normally handles wrapping
+    // non-ok result to error value types, so emulate that here.
+    //
+    // Use the descriptor pool and message factory from the value. This is not
+    // totally consistent with modern APIs, but this should behave the same as
+    // the legacy map did.
+    const google::protobuf::Message* msg = value_.message_;
+    ABSL_DCHECK(msg->GetDescriptor() != nullptr);
+    ABSL_DCHECK(msg->GetReflection() != nullptr);
+
+    const google::protobuf::DescriptorPool* descriptor_pool =
+        msg->GetDescriptor()->file()->pool();
+    google::protobuf::MessageFactory* message_factory =
+        msg->GetReflection()->GetMessageFactory();
+    auto found =
+        Find(modern_key, descriptor_pool, message_factory, arena, &modern_val);
+    if (!found.ok()) {
+      return google::api::expr::runtime::CreateErrorValue(arena,
+                                                          found.status());
+    }
+    if (!(*found) && !modern_val.IsError()) {
       return std::nullopt;
     }
     return UnsafeLegacyValue(modern_val, /*stable=*/false, arena);
@@ -401,21 +412,25 @@ class LegacyParsedJsonMapValue final
     if (arena == nullptr) {
       arena = arena_;
     }
-    if (auto status =
-            google::api::expr::runtime::CelValue::CheckMapKeyType(key);
-        !status.ok()) {
-      status.IgnoreError();
-      return std::nullopt;
-    }
     Value modern_key;
-    if (ABSL_PREDICT_FALSE(!ModernValue(arena, key, modern_key).ok())) {
+    if (!ModernValue(arena, key, modern_key).ok()) {
+      // Legacy to modern should succeed for a valid CelValue.
       return std::nullopt;
     }
     Value modern_val;
-    auto status_or_found = value_.Find(
-        modern_key, google::protobuf::DescriptorPool::generated_pool(),
-        google::protobuf::MessageFactory::generated_factory(), arena, &modern_val);
-    if (!status_or_found.ok() || !*status_or_found) {
+    // Call custom map Find directly. MapValue normally handles wrapping
+    // non-ok result to error value types, so emulate that here.
+    //
+    // We know that the descriptor pool and message factory aren't needed here,
+    // so fine to use generated.
+    auto found =
+        Find(modern_key, google::protobuf::DescriptorPool::generated_pool(),
+             google::protobuf::MessageFactory::generated_factory(), arena, &modern_val);
+    if (!found.ok()) {
+      return google::api::expr::runtime::CreateErrorValue(arena,
+                                                          found.status());
+    }
+    if (!(*found) && !modern_val.IsError()) {
       return std::nullopt;
     }
     return UnsafeLegacyValue(modern_val, /*stable=*/false, arena);

common/values/legacy_struct_value_test.cc

@@ -190,6 +190,18 @@ TEST_F(LegacyStructValueTest, MapFieldKeyTypeValidation) {
   CelValue str_key = CelValue::CreateString(&str_key_val);
   auto invalid_has_res = cel_map->Has(str_key);
   EXPECT_THAT(invalid_has_res, StatusIs(absl::StatusCode::kInvalidArgument));
+
+  auto invalid_get_res = cel_map->Get(arena(), str_key);
+  ASSERT_TRUE(invalid_get_res.has_value());
+  ASSERT_TRUE(invalid_get_res->IsError());
+  EXPECT_THAT(*invalid_get_res->ErrorOrDie(),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+
+  auto invalid_subscript_res = (*cel_map)[str_key];
+  ASSERT_TRUE(invalid_subscript_res.has_value());
+  ASSERT_TRUE(invalid_subscript_res->IsError());
+  EXPECT_THAT(*invalid_subscript_res->ErrorOrDie(),
+              StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
 TEST_F(LegacyStructValueTest, JsonStructAccess) {

common/values/parsed_map_field_value.h

@@ -47,6 +47,10 @@ class ValueIterator;
 class ListValue;
 class ParsedJsonMapValue;
 
+namespace common_internal {
+class LegacyParsedMapFieldMapValue;
+}  // namespace common_internal
+
 // ParsedMapFieldValue is a MapValue over a map field of a parsed protocol
 // buffer message.
 class ParsedMapFieldValue final
@@ -192,6 +196,7 @@ class ParsedMapFieldValue final
   friend class ParsedJsonMapValue;
   friend class common_internal::ValueMixin<ParsedMapFieldValue>;
   friend class common_internal::MapValueMixin<ParsedMapFieldValue>;
+  friend class common_internal::LegacyParsedMapFieldMapValue;
   friend ParsedMapFieldValue UnsafeParsedMapFieldValue(
       const google::protobuf::Message* absl_nonnull message,
       const google::protobuf::FieldDescriptor* absl_nonnull field);

common/values/parsed_message_value.h

@@ -48,6 +48,10 @@
 
 namespace cel {
 
+namespace interop_internal {
+bool IsUnsafeParsedMessageValue(const ParsedMessageValue& value);
+}
+
 class MessageValue;
 class StructValue;
 class Value;
@@ -189,6 +193,8 @@ class ParsedMessageValue final
   friend class common_internal::StructValueMixin<ParsedMessageValue>;
   friend ParsedMessageValue UnsafeParsedMessageValue(
       const google::protobuf::Message* absl_nonnull value);
+  friend bool interop_internal::IsUnsafeParsedMessageValue(
+      const ParsedMessageValue& value);
 
   explicit ParsedMessageValue(
       const google::protobuf::Message* absl_nonnull value ABSL_ATTRIBUTE_LIFETIME_BOUND)

eval/internal/BUILD

@@ -56,11 +56,10 @@ cc_test(
         "//eval/public/structs:trivial_legacy_type_info",
         "//eval/testutil:test_message_cc_proto",
         "//internal:testing",
-        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/log:absl_check",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
-        "@com_google_absl//absl/types:variant",
         "@com_google_googleapis//google/rpc/context:attribute_context_cc_proto",
         "@com_google_protobuf//:any_cc_proto",
         "@com_google_protobuf//:protobuf",

eval/internal/cel_value_equal_test.cc

@@ -18,20 +18,21 @@
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <tuple>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "google/protobuf/any.pb.h"
 #include "google/rpc/context/attribute_context.pb.h"
 #include "google/protobuf/descriptor.pb.h"
-#include "absl/status/statusor.h"
+#include "absl/log/absl_check.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "absl/types/span.h"
-#include "absl/types/variant.h"
 #include "eval/public/cel_value.h"
 #include "eval/public/containers/container_backed_list_impl.h"
 #include "eval/public/containers/container_backed_map_impl.h"
@@ -131,6 +132,8 @@ const std::vector<CelValue>& ValueExamples1() {
     result->push_back(CelValue::CreateMap(&CelMapExample1()));
     result->push_back(CelValue::CreateCelTypeView("type"));
 
+    ABSL_CHECK_EQ(arena.SpaceUsed(), 0) << "Arena should not be used.";
+
     return result.release();
   }();
   return *examples;
@@ -185,7 +188,7 @@ std::string CelValueEqualTestName(
 }
 
 TEST_P(CelValueEqualImplTypesTest, Basic) {
-  absl::optional<bool> result = CelValueEqualImpl(lhs(), rhs());
+  std::optional<bool> result = CelValueEqualImpl(lhs(), rhs());
 
   if (lhs().IsNull() || rhs().IsNull()) {
     if (lhs().IsNull() && rhs().IsNull()) {
@@ -267,7 +270,7 @@ const std::vector<NumericInequalityTestCase>& NumericValuesNotEqualExample() {
 using NumericInequalityTest = testing::TestWithParam<NumericInequalityTestCase>;
 TEST_P(NumericInequalityTest, NumericValues) {
   NumericInequalityTestCase test_case = GetParam();
-  absl::optional<bool> result = CelValueEqualImpl(test_case.a, test_case.b);
+  std::optional<bool> result = CelValueEqualImpl(test_case.a, test_case.b);
   EXPECT_TRUE(result.has_value());
   EXPECT_EQ(*result, false);
 }
@@ -280,7 +283,7 @@ INSTANTIATE_TEST_SUITE_P(
     });
 
 TEST(CelValueEqualImplTest, LossyNumericEquality) {
-  absl::optional<bool> result = CelValueEqualImpl(
+  std::optional<bool> result = CelValueEqualImpl(
       CelValue::CreateDouble(
           static_cast<double>(std::numeric_limits<int64_t>::max()) - 1),
       CelValue::CreateInt64(std::numeric_limits<int64_t>::max()));