Consolidate struct value builder

PiperOrigin-RevId: 690350923

Author: jcking

common/type_reflector.cc

@@ -2640,7 +2640,7 @@ class ValueBuilder {
 
   virtual absl::Status SetFieldByNumber(int64_t number, Value value) = 0;
 
-  virtual Value Build() && = 0;
+  virtual absl::StatusOr<Value> Build() && = 0;
 };
 
 using ValueBuilderPtr = std::unique_ptr<ValueBuilder>;

common/value.h

@@ -41,6 +41,7 @@
 #include "common/type_introspector.h"
 #include "common/type_reflector.h"
 #include "common/value.h"
+#include "common/value_factory.h"
 #include "common/value_kind.h"
 #include "common/value_manager.h"
 #include "extensions/protobuf/internal/map_reflection.h"
@@ -129,28 +130,9 @@ class CompatTypeReflector final : public TypeReflector {
   absl::StatusOr<absl::optional<Value>> DeserializeValueImpl(
       ValueFactory& value_factory, absl::string_view type_url,
       const absl::Cord& value) const override {
-    absl::string_view type_name;
-    if (!ParseTypeUrl(type_url, &type_name)) {
-      return absl::InvalidArgumentError("invalid type URL");
-    }
-    const auto* descriptor =
-        descriptor_pool()->FindMessageTypeByName(type_name);
-    if (descriptor == nullptr) {
-      return absl::nullopt;
-    }
-    const auto* prototype = message_factory()->GetPrototype(descriptor);
-    if (prototype == nullptr) {
-      return absl::nullopt;
-    }
-    absl::Nullable<google::protobuf::Arena*> arena =
-        value_factory.GetMemoryManager().arena();
-    auto message = WrapShared(prototype->New(arena), arena);
-    if (!message->ParsePartialFromCord(value)) {
-      return absl::InvalidArgumentError(
-          absl::StrCat("failed to parse `", type_url, "`"));
-    }
-    return Value::Message(WrapShared(prototype->New(arena), arena), pool_,
-                          factory_);
+    // This should not be reachable, as we provide both the pool and the factory
+    // which should trigger DeserializeValue to handle the call and not call us.
+    return absl::nullopt;
   }
 
  private:
@@ -1011,9 +993,9 @@ GetProtoRepeatedFieldFromValueMutator(
   }
 }
 
-class StructValueBuilderImpl final : public StructValueBuilder {
+class MessageValueBuilderImpl {
  public:
-  StructValueBuilderImpl(
+  MessageValueBuilderImpl(
       absl::Nullable<google::protobuf::Arena*> arena,
       absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
       absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
@@ -1025,13 +1007,13 @@ class StructValueBuilderImpl final : public StructValueBuilder {
         descriptor_(message_->GetDescriptor()),
         reflection_(message_->GetReflection()) {}
 
-  ~StructValueBuilderImpl() override {
+  ~MessageValueBuilderImpl() {
     if (arena_ == nullptr && message_ != nullptr) {
       delete message_;
     }
   }
 
-  absl::Status SetFieldByName(absl::string_view name, Value value) override {
+  absl::Status SetFieldByName(absl::string_view name, Value value) {
     const auto* field = descriptor_->FindFieldByName(name);
     if (field == nullptr) {
       field = descriptor_pool_->FindExtensionByPrintableName(descriptor_, name);
@@ -1042,7 +1024,7 @@ class StructValueBuilderImpl final : public StructValueBuilder {
     return SetField(field, std::move(value));
   }
 
-  absl::Status SetFieldByNumber(int64_t number, Value value) override {
+  absl::Status SetFieldByNumber(int64_t number, Value value) {
     if (number < std::numeric_limits<int32_t>::min() ||
         number > std::numeric_limits<int32_t>::max()) {
       return NoSuchFieldError(absl::StrCat(number)).NativeValue();
@@ -1055,7 +1037,12 @@ class StructValueBuilderImpl final : public StructValueBuilder {
     return SetField(field, std::move(value));
   }
 
-  absl::StatusOr<StructValue> Build() && override {
+  absl::StatusOr<Value> Build() && {
+    return Value::Message(WrapShared(std::exchange(message_, nullptr)),
+                          descriptor_pool_, message_factory_);
+  }
+
+  absl::StatusOr<StructValue> BuildStruct() && {
     return ParsedMessageValue(
         WrapShared(std::exchange(message_, nullptr), Allocator(arena_)));
   }
@@ -1519,19 +1506,91 @@ class StructValueBuilderImpl final : public StructValueBuilder {
   well_known_types::Reflection well_known_types_;
 };
 
+class ValueBuilderImpl final : public ValueBuilder {
+ public:
+  ValueBuilderImpl(absl::Nullable<google::protobuf::Arena*> arena,
+                   absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
+                   absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
+                   absl::Nonnull<google::protobuf::Message*> message)
+      : builder_(arena, descriptor_pool, message_factory, message) {}
+
+  absl::Status SetFieldByName(absl::string_view name, Value value) override {
+    return builder_.SetFieldByName(name, std::move(value));
+  }
+
+  absl::Status SetFieldByNumber(int64_t number, Value value) override {
+    return builder_.SetFieldByNumber(number, std::move(value));
+  }
+
+  absl::StatusOr<Value> Build() && override {
+    return std::move(builder_).Build();
+  }
+
+ private:
+  MessageValueBuilderImpl builder_;
+};
+
+class StructValueBuilderImpl final : public StructValueBuilder {
+ public:
+  StructValueBuilderImpl(
+      absl::Nullable<google::protobuf::Arena*> arena,
+      absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
+      absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
+      absl::Nonnull<google::protobuf::Message*> message)
+      : builder_(arena, descriptor_pool, message_factory, message) {}
+
+  absl::Status SetFieldByName(absl::string_view name, Value value) override {
+    return builder_.SetFieldByName(name, std::move(value));
+  }
+
+  absl::Status SetFieldByNumber(int64_t number, Value value) override {
+    return builder_.SetFieldByNumber(number, std::move(value));
+  }
+
+  absl::StatusOr<StructValue> Build() && override {
+    return std::move(builder_).BuildStruct();
+  }
+
+ private:
+  MessageValueBuilderImpl builder_;
+};
+
 }  // namespace
 
-absl::StatusOr<absl::Nonnull<cel::StructValueBuilderPtr>> NewStructValueBuilder(
+absl::StatusOr<absl::Nullable<cel::ValueBuilderPtr>> NewValueBuilder(
+    Allocator<> allocator,
+    absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
+    absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
+    absl::string_view name) {
+  absl::Nullable<const google::protobuf::Descriptor*> descriptor =
+      descriptor_pool->FindMessageTypeByName(name);
+  if (descriptor == nullptr) {
+    return nullptr;
+  }
+  absl::Nullable<const google::protobuf::Message*> prototype =
+      message_factory->GetPrototype(descriptor);
+  if (prototype == nullptr) {
+    return absl::NotFoundError(absl::StrCat(
+        "unable to get prototype for descriptor: ", descriptor->full_name()));
+  }
+  return std::make_unique<ValueBuilderImpl>(allocator.arena(), descriptor_pool,
+                                            message_factory,
+                                            prototype->New(allocator.arena()));
+}
+
+absl::StatusOr<absl::Nullable<cel::StructValueBuilderPtr>>
+NewStructValueBuilder(
     Allocator<> allocator,
     absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
     absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
     absl::string_view name) {
-  const auto* descriptor = descriptor_pool->FindMessageTypeByName(name);
+  absl::Nullable<const google::protobuf::Descriptor*> descriptor =
+      descriptor_pool->FindMessageTypeByName(name);
   if (descriptor == nullptr) {
-    return absl::NotFoundError(
-        absl::StrCat("unable to find descriptor for type: ", name));
+    return nullptr;
   }
-  const auto* prototype = message_factory->GetPrototype(descriptor);
+  absl::Nullable<const google::protobuf::Message*> prototype =
+      message_factory->GetPrototype(descriptor);
   if (prototype == nullptr) {
     return absl::NotFoundError(absl::StrCat(
         "unable to get prototype for descriptor: ", descriptor->full_name()));

common/values/struct_value_builder.cc

@@ -23,20 +23,15 @@
 #include "google/protobuf/descriptor.h"
 #include "google/protobuf/message.h"
 
-namespace cel {
+namespace cel::common_internal {
 
-class ValueFactory;
-
-namespace common_internal {
-
-absl::StatusOr<absl::Nonnull<cel::StructValueBuilderPtr>> NewStructValueBuilder(
+absl::StatusOr<absl::Nullable<cel::StructValueBuilderPtr>>
+NewStructValueBuilder(
     Allocator<> allocator,
     absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
     absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
     absl::string_view name);
 
-}  // namespace common_internal
-
-}  // namespace cel
+}  // namespace cel::common_internal
 
 #endif  // THIRD_PARTY_CEL_CPP_COMMON_VALUES_STRUCT_VALUE_BUILDER_H_

common/values/struct_value_builder.h

@@ -0,0 +1,37 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef THIRD_PARTY_CEL_CPP_COMMON_VALUES_VALUE_BUILDER_H_
+#define THIRD_PARTY_CEL_CPP_COMMON_VALUES_VALUE_BUILDER_H_
+
+#include "absl/base/nullability.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "common/allocator.h"
+#include "common/value.h"
+#include "google/protobuf/descriptor.h"
+#include "google/protobuf/message.h"
+
+namespace cel::common_internal {
+
+// Like NewStructValueBuilder, but deals with well known types.
+absl::StatusOr<absl::Nullable<cel::ValueBuilderPtr>> NewValueBuilder(
+    Allocator<> allocator,
+    absl::Nonnull<const google::protobuf::DescriptorPool*> descriptor_pool,
+    absl::Nonnull<google::protobuf::MessageFactory*> message_factory,
+    absl::string_view name);
+
+}  // namespace cel::common_internal
+
+#endif  // THIRD_PARTY_CEL_CPP_COMMON_VALUES_VALUE_BUILDER_H_

common/values/value_builder.h

@@ -35,24 +35,19 @@ def _expand_tests_to_skip(tests_to_skip):
             result.append(test_to_skip[0:slash] + part)
     return result
 
-def _conformance_test_name(name, modern, arena, optimize, recursive, skip_check):
+def _conformance_test_name(name, optimize, recursive):
     return "_".join(
         [
             name,
-            "arena" if arena else "refcount",
             "optimized" if optimize else "unoptimized",
             "recursive" if recursive else "iterative",
         ],
     )
 
-def _conformance_test_args(modern, arena, optimize, recursive, skip_check, skip_tests, dashboard):
+def _conformance_test_args(modern, optimize, recursive, skip_check, skip_tests, dashboard):
     args = []
     if modern:
         args.append("--modern")
-    elif not arena:
-        fail("arena must be true for legacy")
-    if not modern or arena:
-        args.append("--arena")
     if optimize:
         args.append("--opt")
     if recursive:
@@ -66,10 +61,10 @@ def _conformance_test_args(modern, arena, optimize, recursive, skip_check, skip_
         args.append("--dashboard")
     return args
 
-def _conformance_test(name, data, modern, arena, optimize, recursive, skip_check, skip_tests, tags, dashboard):
+def _conformance_test(name, data, modern, optimize, recursive, skip_check, skip_tests, tags, dashboard):
     native.cc_test(
-        name = _conformance_test_name(name, modern, arena, optimize, recursive, skip_check),
-        args = _conformance_test_args(modern, arena, optimize, recursive, skip_check, skip_tests, dashboard) + ["$(location " + test + ")" for test in data],
+        name = _conformance_test_name(name, optimize, recursive),
+        args = _conformance_test_args(modern, optimize, recursive, skip_check, skip_tests, dashboard) + ["$(location " + test + ")" for test in data],
         data = data,
         deps = ["//conformance:run"],
         tags = tags,
@@ -89,15 +84,12 @@ def gen_conformance_tests(name, data, modern = False, checked = False, dashboard
         dashboard: enable dashboard mode
     """
     skip_check = not checked
-
-    # TODO: enable refcount mode for modern.
     for optimize in (True, False):
         for recursive in (True, False):
             _conformance_test(
                 name,
                 data,
                 modern = modern,
-                arena = True,
                 optimize = optimize,
                 recursive = recursive,
                 skip_check = skip_check,

conformance/run.bzl

@@ -57,9 +57,6 @@ ABSL_FLAG(bool, opt, false, "Enable optimizations (constant folding)");
 ABSL_FLAG(
     bool, modern, false,
     "Use modern cel::Value APIs implementation of the conformance service.");
-ABSL_FLAG(bool, arena, false,
-          "Use arena memory manager (default: global heap ref-counted). Only "
-          "affects the modern implementation");
 ABSL_FLAG(bool, recursive, false,
           "Enable recursive plans. Depth limited to slightly more than the "
           "default nesting limit.");
@@ -279,7 +276,6 @@ NewConformanceServiceFromFlags() {
       cel_conformance::ConformanceServiceOptions{
           .optimize = absl::GetFlag(FLAGS_opt),
           .modern = absl::GetFlag(FLAGS_modern),
-          .arena = absl::GetFlag(FLAGS_arena),
           .recursive = absl::GetFlag(FLAGS_recursive)});
   ABSL_CHECK_OK(status_or_service);
   return std::shared_ptr<cel_conformance::ConformanceServiceInterface>(