No public description

common/src/main/java/dev/cel/common/internal/CombinedDescriptorPool.java

@@ -14,6 +14,10 @@
 
 package dev.cel.common.internal;
 
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.Immutable;
 import com.google.protobuf.Descriptors.Descriptor;
@@ -33,8 +37,9 @@
 public final class CombinedDescriptorPool implements CelDescriptorPool {
   private final ImmutableList<CelDescriptorPool> descriptorPools;
 
-  @SuppressWarnings("Immutable") // ExtensionRegistry is immutable, just not marked as such.
-  private final ExtensionRegistry extensionRegistry;
+  // ExtensionRegistry and the memoizing Supplier are effectively immutable.
+  @SuppressWarnings("Immutable")
+  private final Supplier<ExtensionRegistry> extensionRegistry;
 
   public static CombinedDescriptorPool create(ImmutableList<CelDescriptorPool> descriptorPools) {
     return new CombinedDescriptorPool(descriptorPools);
@@ -68,19 +73,21 @@ public Optional<FieldDescriptor> findExtensionDescriptor(
 
   @Override
   public ExtensionRegistry getExtensionRegistry() {
-    return extensionRegistry;
+    return extensionRegistry.get();
   }
 
   private CombinedDescriptorPool(ImmutableList<CelDescriptorPool> descriptorPools) {
-    this.descriptorPools = descriptorPools;
+    this.descriptorPools = checkNotNull(descriptorPools);
     // TODO: Combine the extension registry. This will become necessary once we accept
     // ExtensionRegistry through runtime builder. Ideally, proto team should open source this
     // implementation but we may have to create our own.
     this.extensionRegistry =
-        descriptorPools.stream()
-            .map(CelDescriptorPool::getExtensionRegistry)
-            .filter(e -> !e.equals(ExtensionRegistry.getEmptyRegistry()))
-            .findFirst()
-            .orElse(ExtensionRegistry.getEmptyRegistry());
+        Suppliers.memoize(
+            () ->
+                descriptorPools.stream()
+                    .map(CelDescriptorPool::getExtensionRegistry)
+                    .filter(e -> !e.equals(ExtensionRegistry.getEmptyRegistry()))
+                    .findFirst()
+                    .orElse(ExtensionRegistry.getEmptyRegistry()));
   }
 }

common/src/test/java/dev/cel/common/internal/CombinedDescriptorPoolTest.java

@@ -15,15 +15,18 @@
 package dev.cel.common.internal;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableList;
+import com.google.protobuf.Descriptors.Descriptor;
 import com.google.protobuf.Descriptors.FieldDescriptor;
 import com.google.protobuf.ExtensionRegistry;
 import com.google.protobuf.Value;
 import dev.cel.common.CelDescriptorUtil;
 import dev.cel.expr.conformance.proto2.TestAllTypes;
 import dev.cel.expr.conformance.proto2.TestAllTypesExtensions;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -96,6 +99,11 @@ public void findExtensionDescriptor_returnsEmpty() {
         .isEmpty();
   }
 
+  @Test
+  public void create_nullDescriptorPools_throwsNullPointerException() {
+    assertThrows(NullPointerException.class, () -> CombinedDescriptorPool.create(null));
+  }
+
   @Test
   public void getExtensionRegistry_onDefaultPool_returnsEmptyRegistry() {
     CombinedDescriptorPool combinedDescriptorPool =
@@ -104,4 +112,59 @@ public void getExtensionRegistry_onDefaultPool_returnsEmptyRegistry() {
     assertThat(combinedDescriptorPool.getExtensionRegistry())
         .isEqualTo(ExtensionRegistry.getEmptyRegistry());
   }
+
+  @Test
+  public void findDescriptor_doesNotResolveExtensionRegistry() {
+    AtomicInteger getExtensionRegistryCalls = new AtomicInteger();
+    CombinedDescriptorPool combinedDescriptorPool =
+        CombinedDescriptorPool.create(
+            ImmutableList.of(createTrackingPool(getExtensionRegistryCalls)));
+
+    Optional<Descriptor> descriptor = combinedDescriptorPool.findDescriptor("bogus");
+
+    assertThat(descriptor).isEmpty();
+    assertThat(getExtensionRegistryCalls.get()).isEqualTo(0);
+  }
+
+  @Test
+  public void getExtensionRegistry_memoizesResolvedRegistry() {
+    AtomicInteger getExtensionRegistryCalls = new AtomicInteger();
+    CombinedDescriptorPool combinedDescriptorPool =
+        CombinedDescriptorPool.create(
+            ImmutableList.of(createTrackingPool(getExtensionRegistryCalls)));
+
+    ExtensionRegistry first = combinedDescriptorPool.getExtensionRegistry();
+    ExtensionRegistry second = combinedDescriptorPool.getExtensionRegistry();
+
+    assertThat(first).isNotEqualTo(ExtensionRegistry.getEmptyRegistry());
+    assertThat(first).isSameInstanceAs(second);
+    assertThat(getExtensionRegistryCalls.get()).isEqualTo(1);
+  }
+
+  private static CelDescriptorPool createTrackingPool(AtomicInteger getExtensionRegistryCalls) {
+    // Test-only counter for verifying lazy evaluation on a single thread.
+    @SuppressWarnings("Immutable")
+    CelDescriptorPool trackingPool =
+        new CelDescriptorPool() {
+          @Override
+          public Optional<Descriptor> findDescriptor(String name) {
+            return Optional.empty();
+          }
+
+          @Override
+          public Optional<FieldDescriptor> findExtensionDescriptor(
+              Descriptor containingDescriptor, String fieldName) {
+            return Optional.empty();
+          }
+
+          @Override
+          public ExtensionRegistry getExtensionRegistry() {
+            getExtensionRegistryCalls.incrementAndGet();
+            ExtensionRegistry registry = ExtensionRegistry.newInstance();
+            TestAllTypesExtensions.registerAllExtensions(registry);
+            return registry;
+          }
+        };
+    return trackingPool;
+  }
 }