Implement CelMutableExpr

common/ast/BUILD.bazel

@@ -32,3 +32,8 @@ java_library(
     name = "expr_util",
     exports = ["//common/src/main/java/dev/cel/common/ast:expr_util"],
 )
+
+java_library(
+    name = "mutable_ast",
+    exports = ["//common/src/main/java/dev/cel/common/ast:mutable_ast"],
+)

common/src/main/java/dev/cel/common/ast/BUILD.bazel

@@ -31,6 +31,11 @@ EXPR_FACTORY_SOURCES = [
     "CelExprIdGeneratorFactory.java",
 ]
 
+# keep sorted
+MUTABLE_AST_SOURCES = [
+    "CelMutableExpr.java",
+]
+
 java_library(
     name = "ast",
     srcs = AST_SOURCES,
@@ -109,3 +114,14 @@ java_library(
         "@maven//:com_google_guava_guava",
     ],
 )
+
+java_library(
+    name = "mutable_ast",
+    srcs = MUTABLE_AST_SOURCES,
+    tags = [
+    ],
+    deps = [
+        ":ast",
+        "@maven//:com_google_guava_guava",
+    ],
+)

common/src/main/java/dev/cel/common/ast/CelMutableExpr.java

@@ -0,0 +1,157 @@
+// 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.
+
+package dev.cel.common.ast;
+
+import static com.google.common.base.Preconditions.checkArgument;
+
+import dev.cel.common.ast.CelExpr.CelNotSet;
+import dev.cel.common.ast.CelExpr.ExprKind;
+import dev.cel.common.ast.CelExpr.ExprKind.Kind;
+
+/**
+ * An abstract representation of a common expression that allows mutation in any of its properties.
+ * The expressions are semantically the same as that of the immutable {@link CelExpr}.
+ *
+ * <p>This allows for an efficient optimization of an AST without having to traverse and rebuild the
+ * entire tree.
+ *
+ * <p>This class is not thread-safe by design.
+ */
+public final class CelMutableExpr {
+  private long id;
+  private ExprKind.Kind exprKind;
+  private CelNotSet notSet;
+  private CelConstant constant;
+  private int hash = 0;
+
+  public long id() {
+    return id;
+  }
+
+  public void setId(long id) {
+    this.id = id;
+  }
+
+  public ExprKind.Kind getKind() {
+    return exprKind;
+  }
+
+  public CelNotSet notSet() {
+    checkExprKind(Kind.NOT_SET);
+    return notSet;
+  }
+
+  public CelConstant constant() {
+    checkExprKind(Kind.CONSTANT);
+    return constant;
+  }
+
+  public void setConstant(CelConstant constant) {
+    this.exprKind = ExprKind.Kind.CONSTANT;
+    this.constant = constant;
+  }
+
+  public static CelMutableExpr ofConstant(CelConstant constant) {
+    return ofConstant(0L, constant);
+  }
+
+  public static CelMutableExpr ofConstant(long id, CelConstant constant) {
+    return new CelMutableExpr(id, constant);
+  }
+
+  public static CelMutableExpr ofNotSet() {
+    return ofNotSet(0L);
+  }
+
+  public static CelMutableExpr ofNotSet(long id) {
+    return new CelMutableExpr(id);
+  }
+
+  private CelMutableExpr(long id, CelConstant mutableConstant) {
+    this.id = id;
+    setConstant(mutableConstant);
+  }
+
+  private CelMutableExpr(long id) {
+    this();
+    this.id = id;
+  }
+
+  private CelMutableExpr() {
+    this.notSet = CelExpr.newBuilder().build().exprKind().notSet();
+    this.exprKind = ExprKind.Kind.NOT_SET;
+  }
+
+  private Object exprValue() {
+    switch (this.exprKind) {
+      case NOT_SET:
+        return notSet();
+      case CONSTANT:
+        return constant();
+      case IDENT:
+      case SELECT:
+      case CALL:
+      case CREATE_LIST:
+      case CREATE_STRUCT:
+      case CREATE_MAP:
+      case COMPREHENSION:
+        // fall-through (not implemented yet)
+    }
+
+    throw new IllegalStateException("Unexpected expr kind: " + this.exprKind);
+  }
+
+  private void checkExprKind(ExprKind.Kind exprKind) {
+    checkArgument(this.exprKind.equals(exprKind), "Invalid ExprKind: %s", exprKind);
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj == this) {
+      return true;
+    }
+    if (obj instanceof CelMutableExpr) {
+      CelMutableExpr that = (CelMutableExpr) obj;
+      if (this.id != that.id() || !this.exprKind.equals(that.getKind())) {
+        return false;
+      }
+      // When both objects' hashes are cached and they do not match, they can never be equal.
+      if (this.hash != 0 && that.hash != 0 && this.hash != that.hash) {
+        return false;
+      }
+      return this.exprValue().equals(that.exprValue());
+    }
+
+    return false;
+  }
+
+  @Override
+  public int hashCode() {
+    if (hash == 0) {
+      int h = 1;
+      h *= 1000003;
+      h ^= (int) ((id >>> 32) ^ id);
+      h *= 1000003;
+      h ^= this.exprValue().hashCode();
+
+      if (h == 0) {
+        h = 1;
+      }
+      hash = h;
+    }
+
+    return hash;
+  }
+}

common/src/test/java/dev/cel/common/ast/BUILD.bazel

@@ -18,6 +18,7 @@ java_library(
         "//common/ast:expr_converter",
         "//common/ast:expr_factory",
         "//common/ast:expr_v1alpha1_converter",
+        "//common/ast:mutable_ast",
         "//common/resources/testdata/proto3:test_all_types_java_proto",
         "//common/types",
         "//compiler",
@@ -28,6 +29,7 @@ java_library(
         "@cel_spec//proto/cel/expr:expr_java_proto",
         "@com_google_googleapis//google/api/expr/v1alpha1:expr_java_proto",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_guava_guava_testlib",
         "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "@maven//:junit_junit",

common/src/test/java/dev/cel/common/ast/CelMutableExprTest.java

@@ -0,0 +1,126 @@
+// 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.
+
+package dev.cel.common.ast;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import com.google.common.testing.EqualsTester;
+import com.google.testing.junit.testparameterinjector.TestParameter;
+import com.google.testing.junit.testparameterinjector.TestParameterInjector;
+import dev.cel.common.ast.CelExpr.ExprKind.Kind;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(TestParameterInjector.class)
+public class CelMutableExprTest {
+
+  @Test
+  public void ofNotSet() {
+    CelMutableExpr mutableExpr = CelMutableExpr.ofNotSet();
+
+    assertThat(mutableExpr.id()).isEqualTo(0L);
+    assertThat(mutableExpr.notSet()).isNotNull();
+  }
+
+  @Test
+  public void ofNotSet_withId() {
+    CelMutableExpr mutableExpr = CelMutableExpr.ofNotSet(1L);
+
+    assertThat(mutableExpr.id()).isEqualTo(1L);
+    assertThat(mutableExpr.notSet()).isNotNull();
+  }
+
+  @Test
+  public void ofConstant() {
+    CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5L));
+
+    assertThat(mutableExpr.id()).isEqualTo(0L);
+    assertThat(mutableExpr.constant()).isEqualTo(CelConstant.ofValue(5L));
+  }
+
+  @Test
+  public void ofConstant_withId() {
+    CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(1L, CelConstant.ofValue(5L));
+
+    assertThat(mutableExpr.id()).isEqualTo(1L);
+    assertThat(mutableExpr.constant()).isEqualTo(CelConstant.ofValue(5L));
+  }
+
+  @Test
+  public void setId_success() {
+    CelMutableExpr mutableExpr = CelMutableExpr.ofConstant(CelConstant.ofValue(5L));
+
+    mutableExpr.setId(2L);
+
+    assertThat(mutableExpr.id()).isEqualTo(2L);
+  }
+
+  @Test
+  public void equalityTest() {
+    new EqualsTester()
+        .addEqualityGroup(CelMutableExpr.ofNotSet())
+        .addEqualityGroup(CelMutableExpr.ofNotSet(1L), CelMutableExpr.ofNotSet(1L))
+        .addEqualityGroup(CelMutableExpr.ofConstant(1L, CelConstant.ofValue(2L)))
+        .addEqualityGroup(
+            CelMutableExpr.ofConstant(5L, CelConstant.ofValue("hello")),
+            CelMutableExpr.ofConstant(5L, CelConstant.ofValue("hello")))
+        .testEquals();
+  }
+
+  @SuppressWarnings("Immutable") // Mutable by design
+  private enum MutableExprKindTestCase {
+    NOT_SET(CelMutableExpr.ofNotSet(1L)),
+    CONSTANT(CelMutableExpr.ofConstant(CelConstant.ofValue(2L)));
+
+    private final CelMutableExpr mutableExpr;
+
+    MutableExprKindTestCase(CelMutableExpr mutableExpr) {
+      this.mutableExpr = mutableExpr;
+    }
+  }
+
+  @Test
+  public void getExprValue_invalidKind_throws(@TestParameter MutableExprKindTestCase testCase) {
+    Kind testCaseKind = testCase.mutableExpr.getKind();
+    if (!testCaseKind.equals(Kind.NOT_SET)) {
+      assertThrows(IllegalArgumentException.class, testCase.mutableExpr::notSet);
+    }
+    if (!testCaseKind.equals(Kind.CONSTANT)) {
+      assertThrows(IllegalArgumentException.class, testCase.mutableExpr::constant);
+    }
+  }
+
+  @SuppressWarnings("Immutable") // Mutable by design
+  private enum HashCodeTestCase {
+    NOT_SET(CelMutableExpr.ofNotSet(1L), -722379961),
+    CONSTANT(CelMutableExpr.ofConstant(2L, CelConstant.ofValue("test")), -724279919);
+
+    private final CelMutableExpr mutableExpr;
+    private final int expectedHashCode;
+
+    HashCodeTestCase(CelMutableExpr mutableExpr, int expectedHashCode) {
+      this.mutableExpr = mutableExpr;
+      this.expectedHashCode = expectedHashCode;
+    }
+  }
+
+  @Test
+  public void hashCodeTest(@TestParameter HashCodeTestCase testCase) {
+    assertThat(testCase.mutableExpr.hashCode()).isEqualTo(testCase.expectedHashCode);
+    // Run it twice to ensure cached value is stable
+    assertThat(testCase.mutableExpr.hashCode()).isEqualTo(testCase.expectedHashCode);
+  }
+}