Dedupe enum types by enum descriptor's fully qualified name in ProtoM…

…essageTypeProvider PiperOrigin-RevId: 626125941
google · Apr 18, 2024 · c7c2447 · c7c2447
1 parent 141dda8
commit c7c2447
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/bundle/src/test/java/dev/cel/bundle/CelImplTest.java b/bundle/src/test/java/dev/cel/bundle/CelImplTest.java
@@ -41,12 +41,15 @@
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.protobuf.Any;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.DescriptorProtos.FileDescriptorProto;
 import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
+import com.google.protobuf.Descriptors.FileDescriptor;
 import com.google.protobuf.Duration;
 import com.google.protobuf.FieldMask;
 import com.google.protobuf.Message;
 import com.google.protobuf.NullValue;
 import com.google.protobuf.Struct;
+import com.google.protobuf.TextFormat;
 import com.google.protobuf.Timestamp;
 import com.google.protobuf.util.Timestamps;
 import com.google.rpc.context.AttributeContext;
@@ -1021,6 +1024,34 @@ public void compile_enumTypeTransitiveResolutionFailure() {
     assertThat(e).hasMessageThat().contains("undeclared reference to 'NullValue'");
   }
 
+  @Test
+  public void compile_multipleInstancesOfEnumDescriptor_dedupedByFullName() throws Exception {
+    String enumTextProto =
+        "name: \"standalone_global_enum.proto\"\n"
+            + "package: \"dev.cel.testing.testdata.proto3\"\n"
+            + "enum_type {\n"
+            + "  name: \"StandaloneGlobalEnum\"\n"
+            + "  value {\n"
+            + "    name: \"SGOO\"\n"
+            + "    number: 0\n"
+            + "  }\n"
+            + "}\n"
+            + "syntax: \"proto3\"\n";
+    FileDescriptorProto enumFileDescriptorProto =
+        TextFormat.parse(enumTextProto, FileDescriptorProto.class);
+    FileDescriptor enumFileDescriptor =
+        FileDescriptor.buildFrom(enumFileDescriptorProto, new FileDescriptor[] {});
+    Cel cel =
+        standardCelBuilderWithMacros()
+            .setContainer("dev.cel.testing.testdata")
+            .addFileTypes(enumFileDescriptor)
+            .addFileTypes(StandaloneGlobalEnum.getDescriptor().getFile())
+            .build();
+
+    assertThat(cel.compile("dev.cel.testing.testdata.proto3.StandaloneGlobalEnum.SGOO").getAst())
+        .isNotNull();
+  }
+
   @Test
   public void program_customVarResolver() throws Exception {
     Cel cel =

diff --git a/common/src/main/java/dev/cel/common/types/ProtoMessageTypeProvider.java b/common/src/main/java/dev/cel/common/types/ProtoMessageTypeProvider.java
@@ -142,16 +142,19 @@ private ImmutableMap<String, CelType> createProtoMessageTypes(
 
   private ImmutableMap<String, CelType> createEnumTypes(
       Collection<EnumDescriptor> enumDescriptors) {
-    ImmutableMap.Builder<String, CelType> enumTypes = ImmutableMap.builder();
+    HashMap<String, CelType> enumTypes = new HashMap<>();
     for (EnumDescriptor enumDescriptor : enumDescriptors) {
+      if (enumTypes.containsKey(enumDescriptor.getFullName())) {
+        continue;
+      }
       ImmutableMap<String, Integer> values =
           enumDescriptor.getValues().stream()
               .collect(
                   toImmutableMap(EnumValueDescriptor::getName, EnumValueDescriptor::getNumber));
       enumTypes.put(
           enumDescriptor.getFullName(), EnumType.create(enumDescriptor.getFullName(), values));
     }
-    return enumTypes.buildOrThrow();
+    return ImmutableMap.copyOf(enumTypes);
   }
 
   private static class FieldResolver {