From cdad3a11d55a73377501b6cc7a7ee87d246a93ba Mon Sep 17 00:00:00 2001
From: Justin King <jcking@google.com>
Date: Wed, 14 Oct 2026 13:39:43 -0700
Subject: [PATCH] Allow updating `cel::UnknownValue` in-place

PiperOrigin-RevId: 999588834
---
 common/internal/unknowns.cc | 10 ++++------
 common/internal/unknowns.h  | 37 ++++++++++++++++++++++++++++++-------
 common/value.h              | 13 +++++++++----
 3 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/common/internal/unknowns.cc b/common/internal/unknowns.cc
index 34fc05ed6..9f9061585 100644
--- a/common/internal/unknowns.cc
+++ b/common/internal/unknowns.cc
@@ -159,7 +159,7 @@ std::string UnknownTree::ToString(
   return result;
 }
 
-UnknownSet UnknownTree::GraftImpl(const UnknownSet& in) {
+void UnknownTree::GraftImpl(const UnknownSet& in, UnknownSet& out) {
   ABSL_DCHECK_NE(in.root, RootIfPresent());
   // This should only happen in rare circumstances where somebody passed an
   // UnknownValue from one expression evaluation into another.
@@ -184,11 +184,9 @@ UnknownSet UnknownTree::GraftImpl(const UnknownSet& in) {
     std::sort(unknown_function_set->nodes,
               unknown_function_set->nodes + unknown_function_set->size);
   }
-  return UnknownSet{
-      .root = Root(),
-      .attributes = unknown_attribute_set,
-      .functions = unknown_function_set,
-  };
+  out.root = Root();
+  out.attributes = unknown_attribute_set;
+  out.functions = unknown_function_set;
 }
 
 UnknownAttributeNode* UnknownTree::GraftImpl(const UnknownAttributeNode* in) {
diff --git a/common/internal/unknowns.h b/common/internal/unknowns.h
index d23517265..b0bcc7dba 100644
--- a/common/internal/unknowns.h
+++ b/common/internal/unknowns.h
@@ -22,6 +22,7 @@
 #include <string>
 #include <variant>
 
+#include "absl/base/attributes.h"
 #include "absl/base/nullability.h"
 #include "absl/base/optimization.h"
 #include "absl/container/flat_hash_set.h"
@@ -258,7 +259,9 @@ class UnknownTree {
   UnknownTree(const UnknownTree&) = delete;
   UnknownTree& operator=(const UnknownTree&) = delete;
 
-  explicit UnknownTree(google::protobuf::Arena* absl_nonnull arena) : arena_(arena) {}
+  explicit UnknownTree(google::protobuf::Arena* absl_nonnull arena) : arena_(arena) {
+    ABSL_DCHECK(arena != nullptr);
+  }
 
   void MaybeReset(google::protobuf::Arena* absl_nonnull arena) {
     ABSL_DCHECK(arena != nullptr);
@@ -345,21 +348,41 @@ class UnknownTree {
   // address for efficiency. This function does that grafting when it is
   // required.
   [[nodiscard]]
-  UnknownSet Graft(const UnknownSet& in) {
+  ABSL_ATTRIBUTE_ALWAYS_INLINE UnknownSet Graft(const UnknownSet& in) {
     // Fast path, should be taken almost always.
+    UnknownSet out;
+    if (ABSL_PREDICT_TRUE(in.root == nullptr || in.root == RootIfPresent())) {
+      ABSL_DCHECK(in.root != nullptr || in.attributes == nullptr);
+      ABSL_DCHECK(in.root != nullptr || in.functions == nullptr);
+      out = in;
+    } else {
+      // Slow path, should be taken extremely rarely.
+      GraftImpl(in, out);
+    }
+    return out;
+  }
+
+  // In the rare event that somebody manages to pass an unknown set from one
+  // expression into another, we need to be able to graft those nodes into our
+  // tree. The unknown attribute set and unknown function set sort nodes by
+  // address for efficiency. This function does that grafting when it is
+  // required.
+  ABSL_ATTRIBUTE_ALWAYS_INLINE void GraftAt(UnknownSet& in) {
     if (ABSL_PREDICT_TRUE(in.root == nullptr || in.root == RootIfPresent())) {
-      return in;
+      // Fast path, should be taken almost always.
+      ABSL_DCHECK(in.root != nullptr || in.attributes == nullptr);
+      ABSL_DCHECK(in.root != nullptr || in.functions == nullptr);
+    } else {
+      // Slow path, should be taken extremely rarely.
+      GraftImpl(in, in);
     }
-    // Slow path, should be taken extremely rarely.
-    return GraftImpl(in);
   }
 
  private:
   friend class UnknownAttributeNode;
   friend class UnknownRoot;
 
-  [[nodiscard]]
-  UnknownSet GraftImpl(const UnknownSet& in);
+  ABSL_ATTRIBUTE_NOINLINE void GraftImpl(const UnknownSet& in, UnknownSet& out);
 
   [[nodiscard]]
   UnknownAttributeNode* GraftImpl(const UnknownAttributeNode* in);
diff --git a/common/value.h b/common/value.h
index b8801c00e..fc7992c49 100644
--- a/common/value.h
+++ b/common/value.h
@@ -1020,9 +1020,7 @@ class Value final : private common_internal::ValueMixin<Value> {
   // Performs a checked cast from a value to an unknown value,
   // returning a non-empty optional with either a value or reference to the
   // unknown value. Otherwise an empty optional is returned.
-  optional_ref<const UnknownValue> AsUnknown() & ABSL_ATTRIBUTE_LIFETIME_BOUND {
-    return std::as_const(*this).AsUnknown();
-  }
+  optional_ref<UnknownValue> AsUnknown() & ABSL_ATTRIBUTE_LIFETIME_BOUND;
   optional_ref<const UnknownValue> AsUnknown()
       const& ABSL_ATTRIBUTE_LIFETIME_BOUND;
   absl::optional<UnknownValue> AsUnknown() &&;
@@ -1666,7 +1664,7 @@ class Value final : private common_internal::ValueMixin<Value> {
   // `AsUnknown()`.
   template <typename T>
       std::enable_if_t<std::is_same_v<UnknownValue, T>,
-                       optional_ref<const UnknownValue>>
+                       optional_ref<UnknownValue>>
       As() & ABSL_ATTRIBUTE_LIFETIME_BOUND {
     return AsUnknown();
   }
@@ -3193,6 +3191,13 @@ inline optional_ref<const UnknownValue> Value::AsUnknown() const& {
   return std::nullopt;
 }
 
+inline optional_ref<UnknownValue> Value::AsUnknown() & {
+  if (auto* alternative = variant_.As<UnknownValue>(); alternative != nullptr) {
+    return *alternative;
+  }
+  return std::nullopt;
+}
+
 inline absl::optional<UnknownValue> Value::AsUnknown() && {
   if (auto* alternative = variant_.As<UnknownValue>(); alternative != nullptr) {
     return std::move(*alternative);