Inline legacy message field access logic into select/optimized_select implementations.

eval/eval/BUILD

@@ -314,11 +314,15 @@ cc_library(
         ":direct_expression_step",
         ":evaluator_core",
         ":expression_step_base",
+        "//common:memory",
         "//common:type",
         "//common:value",
         "//common:value_kind",
+        "//eval/public:cel_value",
+        "//eval/public/structs:proto_message_type_adapter",
         "//internal:status_macros",
         "//runtime:runtime_options",
+        "@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",

eval/eval/select_step.cc

@@ -2,23 +2,28 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 
+#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 "absl/strings/string_view.h"
 #include "absl/types/optional.h"
 #include "common/legacy_value.h"
+#include "common/memory.h"
 #include "common/type.h"
 #include "common/value.h"
 #include "common/value_kind.h"
 #include "eval/eval/attribute_trail.h"
 #include "eval/eval/direct_expression_step.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_base.h"
+#include "eval/public/cel_value.h"
+#include "eval/public/structs/proto_message_type_adapter.h"
 #include "internal/status_macros.h"
 #include "runtime/runtime_options.h"
 #include "google/protobuf/arena.h"
@@ -74,6 +79,29 @@ absl::optional<Value> CheckForMarkedAttributes(const AttributeTrail& trail,
   return std::nullopt;
 }
 
+// Helper for StructValue::GetFieldByName. Used for opting out of old reflection
+// implementation.
+absl::Status WrappedStructGet(
+    const Value& target, absl::string_view field,
+    ProtoWrapperTypeOptions unboxing_option,
+    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) {
+  if (const google::protobuf::Message* message =
+          cel::interop_internal::GetLegacyMessage(target);
+      message != nullptr) {
+    CelValue::MessageWrapper message_wrapper(
+        message, &GetGenericProtoTypeInfoInstance());
+    CEL_ASSIGN_OR_RETURN(CelValue cel_value,
+                         internal::GetGenericProtoAccessApisInstance().GetField(
+                             field, message_wrapper, unboxing_option,
+                             cel::MemoryManagerRef::Pooling(arena)));
+    return cel::ModernValue(arena, cel_value, *result);
+  }
+  return target.GetStruct().GetFieldByName(
+      field, unboxing_option, descriptor_pool, message_factory, arena, result);
+}
+
 absl::Status PerformHas(const Value& target, absl::string_view field,
                         const StringValue& field_value,
                         const google::protobuf::DescriptorPool* descriptor_pool,
@@ -115,9 +143,9 @@ absl::Status PerformGet(const Value& target, absl::string_view field,
       return absl::OkStatus();
     }
     case ValueKind::kStruct: {
-      auto status = target.GetStruct().GetFieldByName(
-          field, unboxing_option, descriptor_pool, message_factory, arena,
-          &result);
+      auto status =
+          WrappedStructGet(target, field, unboxing_option, descriptor_pool,
+                           message_factory, arena, &result);
       if (!status.ok()) {
         result = ErrorValue(std::move(status));
       }
@@ -154,9 +182,9 @@ absl::Status PerformOptionalGet(const Value& target, absl::string_view field,
         result = OptionalValue::None();
         return absl::OkStatus();
       }
-      CEL_RETURN_IF_ERROR(target.GetStruct().GetFieldByName(
-          field, unboxing_option, descriptor_pool, message_factory, arena,
-          &result));
+      CEL_RETURN_IF_ERROR(WrappedStructGet(target, field, unboxing_option,
+                                           descriptor_pool, message_factory,
+                                           arena, &result));
 
       ABSL_DCHECK(!result.IsUnknown());
       result = OptionalValue::Of(std::move(result), arena);

extensions/BUILD

@@ -331,6 +331,7 @@ cc_library(
         "//common:expr",
         "//common:function_descriptor",
         "//common:kind",
+        "//common:memory",
         "//common:native_type",
         "//common:type",
         "//common:value",
@@ -340,10 +341,13 @@ cc_library(
         "//eval/eval:direct_expression_step",
         "//eval/eval:evaluator_core",
         "//eval/eval:expression_step_base",
+        "//eval/public:cel_value",
+        "//eval/public/structs:proto_message_type_adapter",
         "//internal:casts",
         "//internal:number",
         "//internal:status_macros",
         "//runtime:runtime_builder",
+        "//runtime:runtime_options",
         "//runtime/internal:errors",
         "//runtime/internal:runtime_friend_access",
         "//runtime/internal:runtime_impl",
@@ -355,7 +359,6 @@ cc_library(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_absl//absl/types:variant",
         "@com_google_protobuf//:protobuf",

extensions/select_optimization.cc

@@ -18,8 +18,10 @@
 #include <cstdint>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -31,7 +33,6 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
-#include "absl/types/optional.h"
 #include "absl/types/span.h"
 #include "absl/types/variant.h"
 #include "base/attribute.h"
@@ -43,6 +44,8 @@
 #include "common/expr.h"
 #include "common/function_descriptor.h"
 #include "common/kind.h"
+#include "common/legacy_value.h"
+#include "common/memory.h"
 #include "common/native_type.h"
 #include "common/type.h"
 #include "common/value.h"
@@ -52,13 +55,16 @@
 #include "eval/eval/direct_expression_step.h"
 #include "eval/eval/evaluator_core.h"
 #include "eval/eval/expression_step_base.h"
+#include "eval/public/cel_value.h"
+#include "eval/public/structs/proto_message_type_adapter.h"
 #include "internal/casts.h"
 #include "internal/number.h"
 #include "internal/status_macros.h"
 #include "runtime/internal/errors.h"
 #include "runtime/internal/runtime_friend_access.h"
 #include "runtime/internal/runtime_impl.h"
 #include "runtime/runtime_builder.h"
+#include "runtime/runtime_options.h"
 #include "google/protobuf/arena.h"
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
@@ -74,12 +80,15 @@ using ::cel::Expr;
 using ::cel::ExprKind;
 using ::cel::SelectExpr;
 using ::google::api::expr::runtime::AttributeTrail;
+using ::google::api::expr::runtime::CelValue;
 using ::google::api::expr::runtime::DirectExpressionStep;
 using ::google::api::expr::runtime::ExecutionFrame;
 using ::google::api::expr::runtime::ExecutionFrameBase;
 using ::google::api::expr::runtime::ExpressionStepBase;
+using ::google::api::expr::runtime::GetGenericProtoTypeInfoInstance;
 using ::google::api::expr::runtime::PlannerContext;
 using ::google::api::expr::runtime::ProgramOptimizer;
+using ::google::api::expr::runtime::internal::GetGenericProtoAccessApisInstance;
 
 // Represents a single select operation (field access or indexing).
 // For struct-typed field accesses, includes the field name and the field
@@ -267,6 +276,57 @@ absl::StatusOr<Value> MapKeyFromQualifier(const AttributeQualifier& qual,
   }
 }
 
+// Helper for StructValue::GetFieldByName. Used for opting out of old reflection
+// implementation.
+absl::StatusOr<Value> WrappedStructGet(
+    const Value& target, absl::string_view field,
+    const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
+    google::protobuf::MessageFactory* absl_nonnull message_factory,
+    google::protobuf::Arena* absl_nonnull arena) {
+  if (const google::protobuf::Message* message =
+          cel::interop_internal::GetLegacyMessage(target);
+      message != nullptr) {
+    CelValue::MessageWrapper message_wrapper(
+        message, &GetGenericProtoTypeInfoInstance());
+    CEL_ASSIGN_OR_RETURN(
+        CelValue cel_value,
+        GetGenericProtoAccessApisInstance().GetField(
+            field, message_wrapper, ProtoWrapperTypeOptions::kUnsetProtoDefault,
+            MemoryManagerRef::Pooling(arena)));
+    Value result;
+    CEL_RETURN_IF_ERROR(cel::ModernValue(arena, cel_value, result));
+    return result;
+  }
+  return target.GetStruct().GetFieldByName(field, descriptor_pool,
+                                           message_factory, arena);
+}
+
+// Helper for StructValue::Qualify. Used for opting out of old reflection
+// implementation.
+absl::StatusOr<std::pair<Value, int>> WrappedStructQualify(
+    const StructValue& struct_value,
+    absl::Span<const SelectQualifier> qualifiers, bool presence_test,
+    const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
+    google::protobuf::MessageFactory* absl_nonnull message_factory,
+    google::protobuf::Arena* absl_nonnull arena) {
+  if (const google::protobuf::Message* message =
+          cel::interop_internal::GetLegacyMessage(struct_value);
+      message != nullptr) {
+    CelValue::MessageWrapper message_wrapper(
+        message, &GetGenericProtoTypeInfoInstance());
+    CEL_ASSIGN_OR_RETURN(auto legacy_result,
+                         GetGenericProtoAccessApisInstance().Qualify(
+                             qualifiers, message_wrapper, presence_test,
+                             MemoryManagerRef::Pooling(arena)));
+    Value result;
+    CEL_RETURN_IF_ERROR(cel::ModernValue(arena, legacy_result.value, result));
+    return std::pair<Value, int>{std::move(result),
+                                 legacy_result.qualifier_count};
+  }
+  return struct_value.Qualify(qualifiers, presence_test, descriptor_pool,
+                              message_factory, arena);
+}
+
 absl::StatusOr<Value> ApplyQualifier(
     const Value& operand, const SelectQualifier& qualifier,
     const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
@@ -280,8 +340,8 @@ absl::StatusOr<Value> ApplyQualifier(
                   cel::runtime_internal::CreateNoMatchingOverloadError(
                       "<select>"));
             }
-            return operand.GetStruct().GetFieldByName(
-                field_specifier.name, descriptor_pool, message_factory, arena);
+            return WrappedStructGet(operand, field_specifier.name,
+                                    descriptor_pool, message_factory, arena);
           },
           [&](const AttributeQualifier& qualifier) -> absl::StatusOr<Value> {
             if (operand.Is<ListValue>()) {
@@ -632,7 +692,7 @@ absl::StatusOr<Value> OptimizedSelectImpl::ApplySelect(
   auto value_or =
       (options_.force_fallback_implementation)
           ? absl::UnimplementedError("Forced fallback impl")
-          : struct_value.Qualify(select_path_, presence_test_,
+          : WrappedStructQualify(struct_value, select_path_, presence_test_,
                                  frame.descriptor_pool(),
                                  frame.message_factory(), frame.arena());