From 33a1bb161d2463261dc844266695156679ffda4e Mon Sep 17 00:00:00 2001
From: CEL Dev Team <cel-dev@google.com>
Date: Fri, 16 Aug 2024 19:27:21 -0700
Subject: [PATCH] Add replace function to CEL strings extension in C++

This extension is implemented in Java CEL (https://github.com/google/cel-java/blob/main/extensions/src/main/java/dev/cel/extensions/CelStringExtensions.java) but supported was never added to C++.

PiperOrigin-RevId: 663961461
---
 extensions/strings.cc      | 18 ++++++++++++++++++
 extensions/strings_test.cc | 28 ++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/extensions/strings.cc b/extensions/strings.cc
index 62d41aae1..06d53cb61 100644
--- a/extensions/strings.cc
+++ b/extensions/strings.cc
@@ -25,6 +25,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/cord.h"
+#include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "common/casting.h"
 #include "common/type.h"
@@ -89,6 +90,17 @@ absl::StatusOr<Value> Join1(ValueManager& value_manager,
   return Join2(value_manager, value, StringValue{});
 }
 
+absl::StatusOr<Value> Replace(ValueManager& value_manager,
+                              const StringValue& string,
+                              const StringValue& original,
+                              const StringValue& replacement) {
+  std::string content = string.NativeString();
+  absl::StrReplaceAll({{original.NativeString(), replacement.NativeString()}},
+                      &content);
+  // We assume the original strings were well-formed.
+  return value_manager.CreateUncheckedStringValue(std::move(content));
+}
+
 struct SplitWithEmptyDelimiter {
   ValueManager& value_manager;
   int64_t& limit;
@@ -230,6 +242,12 @@ absl::Status RegisterStringsFunctions(FunctionRegistry& registry,
           CreateDescriptor("join", /*receiver_style=*/true),
       BinaryFunctionAdapter<absl::StatusOr<Value>, ListValue,
                             StringValue>::WrapFunction(Join2)));
+  CEL_RETURN_IF_ERROR(registry.Register(
+      VariadicFunctionAdapter<
+          absl::StatusOr<Value>, StringValue, StringValue,
+          StringValue>::CreateDescriptor("replace", /*receiver_style=*/true),
+      VariadicFunctionAdapter<absl::StatusOr<Value>, StringValue, StringValue,
+                              StringValue>::WrapFunction(Replace)));
   CEL_RETURN_IF_ERROR(registry.Register(
       BinaryFunctionAdapter<absl::StatusOr<Value>, StringValue, StringValue>::
           CreateDescriptor("split", /*receiver_style=*/true),
diff --git a/extensions/strings_test.cc b/extensions/strings_test.cc
index 070c7a26d..627f4304b 100644
--- a/extensions/strings_test.cc
+++ b/extensions/strings_test.cc
@@ -39,6 +39,34 @@ using ::google::api::expr::v1alpha1::ParsedExpr;
 using ::google::api::expr::parser::Parse;
 using ::google::api::expr::parser::ParserOptions;
 
+TEST(Strings, Replace) {
+  MemoryManagerRef memory_manager = MemoryManagerRef::ReferenceCounting();
+  const auto options = RuntimeOptions{};
+  ASSERT_OK_AND_ASSIGN(auto builder, CreateStandardRuntimeBuilder(options));
+  EXPECT_OK(RegisterStringsFunctions(builder.function_registry(), options));
+
+  ASSERT_OK_AND_ASSIGN(auto runtime, std::move(builder).Build());
+
+  ASSERT_OK_AND_ASSIGN(ParsedExpr expr,
+                       Parse("foo.replace('_', ' ') == 'hello world!'",
+                             "<input>", ParserOptions{}));
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Program> program,
+                       ProtobufRuntimeAdapter::CreateProgram(*runtime, expr));
+
+  common_internal::LegacyValueManager value_factory(memory_manager,
+                                                    runtime->GetTypeProvider());
+
+  Activation activation;
+  activation.InsertOrAssignValue("foo",
+                                 StringValue{absl::Cord("hello_world!")});
+
+  ASSERT_OK_AND_ASSIGN(Value result,
+                       program->Evaluate(activation, value_factory));
+  ASSERT_TRUE(result.Is<BoolValue>());
+  EXPECT_TRUE(result.As<BoolValue>().NativeValue());
+}
+
 TEST(Strings, SplitWithEmptyDelimiterCord) {
   MemoryManagerRef memory_manager = MemoryManagerRef::ReferenceCounting();
   const auto options = RuntimeOptions{};