Migrate repeated/map field accesses to use modern implementation.

PiperOrigin-RevId: 966689473

Author: jnthntatum

common/BUILD

@@ -865,6 +865,8 @@ cc_test(
         ":value_kind",
         ":value_testing",
         "//base:attributes",
+        "//eval/public:cel_value",
+        "//eval/public/structs:proto_message_type_adapter",
         "//internal:parse_text_proto",
         "//internal:status_macros",
         "//internal:testing",

common/legacy_value.cc

@@ -43,16 +43,15 @@
 #include "common/unknown.h"
 #include "common/value.h"
 #include "common/value_kind.h"
+#include "common/values/legacy_list_value.h"
+#include "common/values/legacy_map_value.h"
 #include "common/values/list_value_builder.h"
 #include "common/values/map_value_builder.h"
 #include "common/values/values.h"
 #include "eval/internal/cel_value_equal.h"
 #include "eval/public/cel_value.h"
-#include "eval/public/containers/field_backed_list_impl.h"
-#include "eval/public/containers/field_backed_map_impl.h"
 #include "eval/public/message_wrapper.h"
 #include "eval/public/structs/cel_proto_wrap_util.h"
-#include "eval/public/structs/legacy_type_adapter.h"
 #include "eval/public/structs/legacy_type_info_apis.h"
 #include "eval/public/structs/proto_message_type_adapter.h"
 #include "eval/public/structs/trivial_legacy_type_info_internal.h"
@@ -77,12 +76,9 @@ using ::google::api::expr::runtime::CelList;
 using ::google::api::expr::runtime::CelMap;
 using ::google::api::expr::runtime::CelValue;
 using ::google::api::expr::runtime::CreateCelValueFromField;
-using ::google::api::expr::runtime::FieldBackedListImpl;
-using ::google::api::expr::runtime::FieldBackedMapImpl;
 using ::google::api::expr::runtime::GetGenericProtoTypeInfoInstance;
 using ::google::api::expr::runtime::LegacyTypeInfoApis;
 using ::google::api::expr::runtime::MessageWrapper;
-using ::google::api::expr::runtime::internal::GetGenericProtoAccessApisInstance;
 using ::google::api::expr::runtime::internal::MaybeWrapValueToMessage;
 
 absl::Status InvalidMapKeyTypeError(ValueKind kind) {
@@ -284,19 +280,17 @@ CelValue LegacyTrivialListValue(google::protobuf::Arena* absl_nonnull arena,
   }
   if (auto parsed_repeated_field_value = value.AsParsedRepeatedField();
       parsed_repeated_field_value) {
-    auto maybe_cloned = parsed_repeated_field_value->Clone(arena);
-    return CelValue::CreateList(google::protobuf::Arena::Create<FieldBackedListImpl>(
-        arena, &maybe_cloned.message(), maybe_cloned.field(), arena));
+    auto wrapped = common_internal::WrapLegacyParsedRepeatedField(
+        *parsed_repeated_field_value, arena);
+    return CelValue::CreateList(
+        common_internal::AsLegacyListValue(wrapped)->cel_list());
   }
   if (auto parsed_json_list_value = value.AsParsedJsonList();
       parsed_json_list_value) {
-    auto maybe_cloned = parsed_json_list_value->Clone(arena);
-    return CelValue::CreateList(google::protobuf::Arena::Create<FieldBackedListImpl>(
-        arena, cel::to_address(maybe_cloned),
-        well_known_types::GetListValueReflectionOrDie(
-            maybe_cloned->GetDescriptor())
-            .GetValuesDescriptor(),
-        arena));
+    auto wrapped = common_internal::WrapLegacyParsedJsonList(
+        *parsed_json_list_value, arena);
+    return CelValue::CreateList(
+        common_internal::AsLegacyListValue(wrapped)->cel_list());
   }
   if (auto custom_list_value = value.AsCustomList(); custom_list_value) {
     auto status_or_compat_list = common_internal::MakeCompatListValue(
@@ -322,19 +316,17 @@ CelValue LegacyTrivialMapValue(google::protobuf::Arena* absl_nonnull arena,
   }
   if (auto parsed_map_field_value = value.AsParsedMapField();
       parsed_map_field_value) {
-    auto maybe_cloned = parsed_map_field_value->Clone(arena);
-    return CelValue::CreateMap(google::protobuf::Arena::Create<FieldBackedMapImpl>(
-        arena, &maybe_cloned.message(), maybe_cloned.field(), arena));
+    auto wrapped = common_internal::WrapLegacyParsedMapField(
+        *parsed_map_field_value, arena);
+    return CelValue::CreateMap(
+        common_internal::AsLegacyMapValue(wrapped)->cel_map());
   }
   if (auto parsed_json_map_value = value.AsParsedJsonMap();
       parsed_json_map_value) {
-    auto maybe_cloned = parsed_json_map_value->Clone(arena);
-    return CelValue::CreateMap(google::protobuf::Arena::Create<FieldBackedMapImpl>(
-        arena, cel::to_address(maybe_cloned),
-        well_known_types::GetStructReflectionOrDie(
-            maybe_cloned->GetDescriptor())
-            .GetFieldsDescriptor(),
-        arena));
+    auto wrapped =
+        common_internal::WrapLegacyParsedJsonMap(*parsed_json_map_value, arena);
+    return CelValue::CreateMap(
+        common_internal::AsLegacyMapValue(wrapped)->cel_map());
   }
   if (auto custom_map_value = value.AsCustomMap(); custom_map_value) {
     auto status_or_compat_map = common_internal::MakeCompatMapValue(
@@ -352,6 +344,25 @@ CelValue LegacyTrivialMapValue(google::protobuf::Arena* absl_nonnull arena,
                  value.GetRuntimeType().DebugString()))));
 }
 
+LegacyStructValue ParsedMessageToLegacyStructValue(
+    const ParsedMessageValue& parsed_message) {
+  return LegacyStructValue(cel::to_address(parsed_message),
+                           &GetGenericProtoTypeInfoInstance());
+}
+
+LegacyStructValue MakeLegacyStructValue(
+    const google::protobuf::Message* absl_nonnull message,
+    const LegacyTypeInfoApis* legacy_type_info) {
+  // Guard against edge cases where a custom implementation of Message
+  // misbehaves.
+  // Modern value handles this with DCHECKs on value creation, legacy value
+  // would allow it and just report an ErrorValue on accesses.
+  if (message->GetReflection() == nullptr || legacy_type_info == nullptr) {
+    legacy_type_info = TrivialTypeInfo::GetInstance();
+  }
+  return LegacyStructValue(message, legacy_type_info);
+}
+
 }  // namespace
 
 google::api::expr::runtime::CelValue UnsafeLegacyValue(
@@ -394,10 +405,6 @@ google::api::expr::runtime::CelValue UnsafeLegacyValue(
   }
 }
 
-}  // namespace common_internal
-
-namespace common_internal {
-
 std::string LegacyListValue::DebugString() const {
   return CelValue::CreateList(impl_).DebugString();
 }
@@ -837,10 +844,8 @@ absl::Status LegacyStructValue::SerializeTo(
   ABSL_DCHECK(message_factory != nullptr);
   ABSL_DCHECK(output != nullptr);
 
-  auto message_wrapper = AsMessageWrapper(message_ptr_, legacy_type_info_);
   if (ABSL_PREDICT_TRUE(
-          message_wrapper.message_ptr()->SerializePartialToZeroCopyStream(
-              output))) {
+          message_ptr_->SerializePartialToZeroCopyStream(output))) {
     return absl::OkStatus();
   }
   return absl::UnknownError("failed to serialize protocol buffer message");
@@ -918,16 +923,38 @@ absl::Status LegacyStructValue::GetFieldByName(
     const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
     google::protobuf::MessageFactory* absl_nonnull message_factory,
     google::protobuf::Arena* absl_nonnull arena, Value* absl_nonnull result) const {
-  auto message_wrapper = AsMessageWrapper(message_ptr_, legacy_type_info_);
   if (ABSL_PREDICT_FALSE(legacy_type_info_ == TrivialTypeInfo::GetInstance())) {
     *result = NoSuchFieldError(name);
     return absl::OkStatus();
   }
-  CEL_ASSIGN_OR_RETURN(auto cel_value,
-                       GetGenericProtoAccessApisInstance().GetField(
-                           name, message_wrapper, unboxing_options,
-                           MemoryManagerRef::Pooling(arena)));
-  CEL_RETURN_IF_ERROR(ModernValue(arena, cel_value, *result));
+
+  ParsedMessageValue parsed_message = UnsafeParsedMessageValue(message_ptr_);
+  const auto* descriptor = parsed_message.GetDescriptor();
+  const auto* field = descriptor->FindFieldByName(name);
+  if (field == nullptr) {
+    field = descriptor->file()->pool()->FindExtensionByPrintableName(descriptor,
+                                                                     name);
+    if (field == nullptr) {
+      *result = NoSuchFieldError(name);
+      return absl::OkStatus();
+    }
+  }
+
+  if (field->is_map()) {
+    *result = WrapLegacyParsedMapField(
+        UnsafeParsedMapFieldValue(message_ptr_, field), arena);
+    return absl::OkStatus();
+  }
+  if (field->is_repeated()) {
+    *result = WrapLegacyParsedRepeatedField(
+        UnsafeParsedRepeatedFieldValue(message_ptr_, field), arena);
+    return absl::OkStatus();
+  }
+
+  CEL_RETURN_IF_ERROR(parsed_message.GetField(field, unboxing_options,
+                                              descriptor_pool, message_factory,
+                                              arena, result));
+  interop_internal::WrapLegacyFieldAccessResult(arena, result);
   return absl::OkStatus();
 }
 
@@ -980,7 +1007,6 @@ absl::Status LegacyStructValue::Qualify(
   if (ABSL_PREDICT_FALSE(qualifiers.empty())) {
     return absl::InvalidArgumentError("invalid select qualifier path.");
   }
-  auto message_wrapper = AsMessageWrapper(message_ptr_, legacy_type_info_);
   if (ABSL_PREDICT_FALSE(legacy_type_info_ == TrivialTypeInfo::GetInstance())) {
     absl::string_view field_name = absl::visit(
         absl::Overload(
@@ -995,12 +1021,13 @@ absl::Status LegacyStructValue::Qualify(
     *count = -1;
     return absl::OkStatus();
   }
-  CEL_ASSIGN_OR_RETURN(auto legacy_result,
-                       GetGenericProtoAccessApisInstance().Qualify(
-                           qualifiers, message_wrapper, presence_test,
-                           MemoryManager::Pooling(arena)));
-  CEL_RETURN_IF_ERROR(ModernValue(arena, legacy_result.value, *result));
-  *count = legacy_result.qualifier_count;
+
+  ParsedMessageValue parsed_message = UnsafeParsedMessageValue(message_ptr_);
+  CEL_RETURN_IF_ERROR(parsed_message.Qualify(qualifiers, presence_test,
+                                             descriptor_pool, message_factory,
+                                             arena, result, count));
+
+  interop_internal::WrapLegacyFieldAccessResult(arena, result);
   return absl::OkStatus();
 }
 
@@ -1035,7 +1062,7 @@ absl::Status ModernValue(google::protobuf::Arena* arena,
       return absl::OkStatus();
     case CelValue::Type::kMessage: {
       auto message_wrapper = legacy_value.MessageWrapperOrDie();
-      result = common_internal::LegacyStructValue(
+      result = common_internal::MakeLegacyStructValue(
           google::protobuf::DownCastMessage<google::protobuf::Message>(
               message_wrapper.message_ptr()),
           message_wrapper.legacy_type_info());
@@ -1153,7 +1180,7 @@ absl::StatusOr<Value> FromLegacyValue(google::protobuf::Arena* arena,
                         legacy_value.BytesOrDie().value());
     case CelValue::Type::kMessage: {
       auto message_wrapper = legacy_value.MessageWrapperOrDie();
-      return common_internal::LegacyStructValue(
+      return common_internal::MakeLegacyStructValue(
           google::protobuf::DownCastMessage<google::protobuf::Message>(
               message_wrapper.message_ptr()),
           message_wrapper.legacy_type_info());
@@ -1262,6 +1289,33 @@ google::api::expr::runtime::CelValue ModernValueToLegacyValueOrDie(
   return std::move(*status_or_value);
 }
 
+void WrapLegacyFieldAccessResult(google::protobuf::Arena* absl_nonnull arena,
+                                 Value* absl_nonnull result) {
+  if (result->IsParsedMessage()) {
+    *result = common_internal::ParsedMessageToLegacyStructValue(
+        result->GetParsedMessage());
+  } else if (result->IsParsedRepeatedField()) {
+    *result =
+        WrapLegacyParsedRepeatedField(result->GetParsedRepeatedField(), arena);
+  } else if (result->IsParsedJsonList()) {
+    *result = WrapLegacyParsedJsonList(result->GetParsedJsonList(), arena);
+  } else if (result->IsParsedMapField()) {
+    *result = WrapLegacyParsedMapField(result->GetParsedMapField(), arena);
+  } else if (result->IsParsedJsonMap()) {
+    *result = WrapLegacyParsedJsonMap(result->GetParsedJsonMap(), arena);
+  } else if (result->IsList()) {
+    auto is_empty = result->GetList().IsEmpty();
+    if (is_empty.ok() && *is_empty) {
+      *result = CustomListValue(common_internal::EmptyCompatListValue(), arena);
+    }
+  } else if (result->IsMap()) {
+    auto is_empty = result->GetMap().IsEmpty();
+    if (is_empty.ok() && *is_empty) {
+      *result = CustomMapValue(common_internal::EmptyCompatMapValue(), arena);
+    }
+  }
+}
+
 TypeValue CreateTypeValueFromView(google::protobuf::Arena* arena,
                                   absl::string_view input) {
   return TypeValue(common_internal::LegacyRuntimeType(input));

common/legacy_value.h

@@ -65,6 +65,13 @@ namespace cel::interop_internal {
 // message with the default type info, or `nullptr` otherwise.
 const google::protobuf::Message* absl_nullable GetLegacyMessage(const Value& value);
 
+// Helper for wrapping a field accesses for the legacy runtime.
+//
+// Adapts the output to avoid further allocations when converting to a legacy
+// value when possible.
+void WrapLegacyFieldAccessResult(google::protobuf::Arena* absl_nonnull arena,
+                                 Value* absl_nonnull result);
+
 // Access a field on a legacy message value, writing the result to `out`.
 // Prefers wrapping legacy values instead of using the modern value
 // representation.

common/value.cc

@@ -1512,6 +1512,13 @@ Value WrapFieldImpl(
   ABSL_DCHECK(!IsWellKnownMessageType(message->GetDescriptor()));
 
   const auto* reflection = message->GetReflection();
+  if (ABSL_PREDICT_FALSE(reflection == nullptr)) {
+    // This only happens for special implementations of Message that
+    // should not normally be used with CEL.
+    return ErrorValue(absl::InvalidArgumentError(
+        absl::StrCat("failed to get reflection for message type: ",
+                     message->GetDescriptor()->full_name())));
+  }
   if (field->is_map()) {
     if (reflection->FieldSize(*message, field) == 0) {
       return MapValue();
@@ -1653,6 +1660,13 @@ Value WrapRepeatedFieldImpl(
   ABSL_DCHECK(arena != nullptr);
 
   const auto* reflection = message->GetReflection();
+  if (ABSL_PREDICT_FALSE(reflection == nullptr)) {
+    // This only happens for special implementations of Message that
+    // should not normally be used with CEL.
+    return ErrorValue(absl::InvalidArgumentError(
+        absl::StrCat("failed to get reflection for message type: ",
+                     message->GetDescriptor()->full_name())));
+  }
   const int size = reflection->FieldSize(*message, field);
   if (ABSL_PREDICT_FALSE(index < 0 || index >= size)) {
     return ErrorValue(absl::InvalidArgumentError(

common/values/custom_map_value.cc

@@ -678,13 +678,22 @@ absl::StatusOr<bool> CustomMapValue::Find(
     CustomMapValueInterface::Content content =
         content_.To<CustomMapValueInterface::Content>();
     ABSL_DCHECK(content.interface != nullptr);
-    CEL_ASSIGN_OR_RETURN(
-        ok, content.interface->Find(key, descriptor_pool, message_factory,
-                                    arena, result));
+    auto status_or_found = content.interface->Find(
+        key, descriptor_pool, message_factory, arena, result);
+    if (!status_or_found.ok()) {
+      *result = ErrorValue(std::move(status_or_found).status());
+      return false;
+    }
+    ok = *status_or_found;
   } else {
-    CEL_ASSIGN_OR_RETURN(
-        ok, dispatcher_->find(dispatcher_, content_, key, descriptor_pool,
-                              message_factory, arena, result));
+    auto status_or_found =
+        dispatcher_->find(dispatcher_, content_, key, descriptor_pool,
+                          message_factory, arena, result);
+    if (!status_or_found.ok()) {
+      *result = ErrorValue(std::move(status_or_found).status());
+      return false;
+    }
+    ok = *status_or_found;
   }
   if (ok) {
     return true;
@@ -726,12 +735,21 @@ absl::Status CustomMapValue::Has(
     CustomMapValueInterface::Content content =
         content_.To<CustomMapValueInterface::Content>();
     ABSL_DCHECK(content.interface != nullptr);
-    CEL_ASSIGN_OR_RETURN(has, content.interface->Has(key, descriptor_pool,
-                                                     message_factory, arena));
+    auto status_or_has =
+        content.interface->Has(key, descriptor_pool, message_factory, arena);
+    if (!status_or_has.ok()) {
+      *result = ErrorValue(std::move(status_or_has).status());
+      return absl::OkStatus();
+    }
+    has = *status_or_has;
   } else {
-    CEL_ASSIGN_OR_RETURN(
-        has, dispatcher_->has(dispatcher_, content_, key, descriptor_pool,
-                              message_factory, arena));
+    auto status_or_has = dispatcher_->has(
+        dispatcher_, content_, key, descriptor_pool, message_factory, arena);
+    if (!status_or_has.ok()) {
+      *result = ErrorValue(std::move(status_or_has).status());
+      return absl::OkStatus();
+    }
+    has = *status_or_has;
   }
   *result = BoolValue(has);
   return absl::OkStatus();

common/values/custom_map_value.h

@@ -54,6 +54,12 @@ class CustomMapValueInterfaceKeysIterator;
 class CustomMapValue;
 using CustomMapValueContent = CustomValueContent;
 
+// Dispatch table for `CustomMapValue`.
+//
+// See the documentation for `CustomMapValueInterface` for more details on
+// composite functions.
+//
+// See documentation for `UnsafeCustomMapValue` on how to use this class.
 struct CustomMapValueDispatcher {
   using GetTypeId =
       NativeTypeId (*)(const CustomMapValueDispatcher* absl_nonnull dispatcher,
@@ -253,6 +259,9 @@ class CustomMapValueInterface {
       google::protobuf::MessageFactory* absl_nonnull message_factory,
       google::protobuf::Arena* absl_nonnull arena, Value* absl_nonnull result) const = 0;
 
+  // Whether the map has the given key.
+  //
+  // A non-ok status is coerced to an ErrorValue (e.g. wrong key type).
   virtual absl::StatusOr<bool> Has(
       const Value& key,
       const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,