Properly retain celStandardDeclarations when invoking toBuilder on checker

PiperOrigin-RevId: 802639372

Author: l46kok

checker/src/main/java/dev/cel/checker/CelCheckerLegacyImpl.java

@@ -55,6 +55,7 @@
 import java.util.List;
 import java.util.SortedSet;
 import java.util.TreeSet;
+import org.jspecify.annotations.Nullable;
 
 /**
  * {@code CelChecker} implementation which uses the original CEL-Java APIs to provide a simple,
@@ -128,6 +129,10 @@ public CelCheckerBuilder toCheckerBuilder() {
       builder.setResultType(expectedResultType.get());
     }
 
+    if (overriddenStandardDeclarations != null) {
+      builder.setStandardDeclarations(overriddenStandardDeclarations);
+    }
+
     return builder;
   }
 
@@ -177,6 +182,7 @@ public static CelCheckerBuilder newBuilder() {
   /** Builder class for the legacy {@code CelChecker} implementation. */
   public static final class Builder implements CelCheckerBuilder {
 
+    // LINT.IfChange
     private final ImmutableSet.Builder<CelIdentDecl> identDeclarations;
     private final ImmutableSet.Builder<CelFunctionDecl> functionDeclarations;
     private final ImmutableSet.Builder<ProtoTypeMask> protoTypeMasks;
@@ -191,6 +197,8 @@ public static final class Builder implements CelCheckerBuilder {
     private boolean standardEnvironmentEnabled;
     private CelStandardDeclarations standardDeclarations;
 
+    // LINT.ThenChange(CelCheckerLegacyImplTest.java)
+
     @Override
     public CelCheckerBuilder setOptions(CelOptions celOptions) {
       this.celOptions = checkNotNull(celOptions);
@@ -379,38 +387,53 @@ Builder addIdentDeclarations(ImmutableSet<CelIdentDecl> identDeclarations) {
       return this;
     }
 
-    // The following getters exist for asserting immutability for collections held by this builder,
-    // and shouldn't be exposed to the public.
+    // The following getters marked @VisibleForTesting exist for testing toCheckerBuilder copies
+    // over all properties. Do not expose these to public
     @VisibleForTesting
-    ImmutableSet.Builder<CelFunctionDecl> getFunctionDecls() {
+    ImmutableSet.Builder<CelFunctionDecl> functionDecls() {
       return this.functionDeclarations;
     }
 
     @VisibleForTesting
-    ImmutableSet.Builder<CelIdentDecl> getIdentDecls() {
+    ImmutableSet.Builder<CelIdentDecl> identDecls() {
       return this.identDeclarations;
     }
 
     @VisibleForTesting
-    ImmutableSet.Builder<ProtoTypeMask> getProtoTypeMasks() {
+    ImmutableSet.Builder<ProtoTypeMask> protoTypeMasks() {
       return this.protoTypeMasks;
     }
 
     @VisibleForTesting
-    ImmutableSet.Builder<Descriptor> getMessageTypes() {
+    ImmutableSet.Builder<Descriptor> messageTypes() {
       return this.messageTypes;
     }
 
     @VisibleForTesting
-    ImmutableSet.Builder<FileDescriptor> getFileTypes() {
+    ImmutableSet.Builder<FileDescriptor> fileTypes() {
       return this.fileTypes;
     }
 
     @VisibleForTesting
-    ImmutableSet.Builder<CelCheckerLibrary> getCheckerLibraries() {
+    ImmutableSet.Builder<CelCheckerLibrary> checkerLibraries() {
       return this.celCheckerLibraries;
     }
 
+    @VisibleForTesting
+    CelStandardDeclarations standardDeclarations() {
+      return this.standardDeclarations;
+    }
+
+    @VisibleForTesting
+    CelOptions options() {
+      return this.celOptions;
+    }
+
+    @VisibleForTesting
+    CelTypeProvider celTypeProvider() {
+      return this.celTypeProvider;
+    }
+
     @Override
     @CheckReturnValue
     public CelCheckerLegacyImpl build() {
@@ -506,7 +529,7 @@ private CelCheckerLegacyImpl(
       TypeProvider typeProvider,
       CelTypeProvider celTypeProvider,
       boolean standardEnvironmentEnabled,
-      CelStandardDeclarations overriddenStandardDeclarations,
+      @Nullable CelStandardDeclarations overriddenStandardDeclarations,
       ImmutableSet<CelCheckerLibrary> checkerLibraries,
       ImmutableSet<FileDescriptor> fileDescriptors,
       ImmutableSet<ProtoTypeMask> protoTypeMasks) {

checker/src/test/java/dev/cel/checker/CelCheckerLegacyImplTest.java

@@ -16,12 +16,20 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.ImmutableList;
+import dev.cel.checker.CelStandardDeclarations.StandardFunction;
+import dev.cel.common.CelContainer;
 import dev.cel.common.CelFunctionDecl;
+import dev.cel.common.CelOptions;
 import dev.cel.common.CelOverloadDecl;
 import dev.cel.common.CelVarDecl;
+import dev.cel.common.types.CelType;
+import dev.cel.common.types.CelTypeProvider;
 import dev.cel.common.types.SimpleType;
 import dev.cel.compiler.CelCompilerFactory;
 import dev.cel.expr.conformance.proto3.TestAllTypes;
+import java.util.Optional;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -49,7 +57,45 @@ public void toCheckerBuilder_isImmutable() {
     CelCheckerLegacyImpl.Builder newCheckerBuilder =
         (CelCheckerLegacyImpl.Builder) celChecker.toCheckerBuilder();
 
-    assertThat(newCheckerBuilder.getCheckerLibraries().build()).isEmpty();
+    assertThat(newCheckerBuilder.checkerLibraries().build()).isEmpty();
+  }
+
+  @Test
+  public void toCheckerBuilder_singularFields_copied() {
+    CelStandardDeclarations subsetDecls =
+        CelStandardDeclarations.newBuilder().includeFunctions(StandardFunction.BOOL).build();
+    CelOptions celOptions = CelOptions.current().enableTimestampEpoch(true).build();
+    CelContainer celContainer = CelContainer.ofName("foo");
+    CelType expectedResultType = SimpleType.BOOL;
+    CelTypeProvider customTypeProvider =
+        new CelTypeProvider() {
+          @Override
+          public ImmutableCollection<CelType> types() {
+            return ImmutableList.of();
+          }
+
+          @Override
+          public Optional<CelType> findType(String typeName) {
+            return Optional.empty();
+          }
+        };
+    CelCheckerBuilder celCheckerBuilder =
+        CelCompilerFactory.standardCelCheckerBuilder()
+            .setOptions(celOptions)
+            .setContainer(celContainer)
+            .setResultType(expectedResultType)
+            .setTypeProvider(customTypeProvider)
+            .setStandardEnvironmentEnabled(false)
+            .setStandardDeclarations(subsetDecls);
+    CelCheckerLegacyImpl celChecker = (CelCheckerLegacyImpl) celCheckerBuilder.build();
+
+    CelCheckerLegacyImpl.Builder newCheckerBuilder =
+        (CelCheckerLegacyImpl.Builder) celChecker.toCheckerBuilder();
+
+    assertThat(newCheckerBuilder.standardDeclarations()).isEqualTo(subsetDecls);
+    assertThat(newCheckerBuilder.options()).isEqualTo(celOptions);
+    assertThat(newCheckerBuilder.container()).isEqualTo(celContainer);
+    assertThat(newCheckerBuilder.celTypeProvider()).isEqualTo(customTypeProvider);
   }
 
   @Test
@@ -70,12 +116,12 @@ public void toCheckerBuilder_collectionProperties_copied() {
     CelCheckerLegacyImpl.Builder newCheckerBuilder =
         (CelCheckerLegacyImpl.Builder) celChecker.toCheckerBuilder();
 
-    assertThat(newCheckerBuilder.getFunctionDecls().build()).hasSize(1);
-    assertThat(newCheckerBuilder.getIdentDecls().build()).hasSize(1);
-    assertThat(newCheckerBuilder.getProtoTypeMasks().build()).hasSize(1);
-    assertThat(newCheckerBuilder.getFileTypes().build())
+    assertThat(newCheckerBuilder.functionDecls().build()).hasSize(1);
+    assertThat(newCheckerBuilder.identDecls().build()).hasSize(1);
+    assertThat(newCheckerBuilder.protoTypeMasks().build()).hasSize(1);
+    assertThat(newCheckerBuilder.fileTypes().build())
         .hasSize(1); // MessageTypes and FileTypes deduped into the same file descriptor
-    assertThat(newCheckerBuilder.getCheckerLibraries().build()).hasSize(1);
+    assertThat(newCheckerBuilder.checkerLibraries().build()).hasSize(1);
   }
 
   @Test
@@ -96,11 +142,11 @@ public void toCheckerBuilder_collectionProperties_areImmutable() {
         ProtoTypeMask.ofAllFields("cel.expr.conformance.proto3.TestAllTypes"));
     celCheckerBuilder.addLibraries(new CelCheckerLibrary() {});
 
-    assertThat(newCheckerBuilder.getFunctionDecls().build()).isEmpty();
-    assertThat(newCheckerBuilder.getIdentDecls().build()).isEmpty();
-    assertThat(newCheckerBuilder.getProtoTypeMasks().build()).isEmpty();
-    assertThat(newCheckerBuilder.getMessageTypes().build()).isEmpty();
-    assertThat(newCheckerBuilder.getFileTypes().build()).isEmpty();
-    assertThat(newCheckerBuilder.getCheckerLibraries().build()).isEmpty();
+    assertThat(newCheckerBuilder.functionDecls().build()).isEmpty();
+    assertThat(newCheckerBuilder.identDecls().build()).isEmpty();
+    assertThat(newCheckerBuilder.protoTypeMasks().build()).isEmpty();
+    assertThat(newCheckerBuilder.messageTypes().build()).isEmpty();
+    assertThat(newCheckerBuilder.fileTypes().build()).isEmpty();
+    assertThat(newCheckerBuilder.checkerLibraries().build()).isEmpty();
   }
 }

checker/src/test/java/dev/cel/checker/CelCompilerImplTest.java

@@ -42,4 +42,6 @@ public void toCompilerBuilder_isImmutable() {
 
     assertThat(newCompilerBuilder).isNotEqualTo(celCompilerBuilder);
   }
+
+
 }