Fix optionals to properly error on invalid qualification

PiperOrigin-RevId: 875419705

Author: l46kok

common/src/main/java/dev/cel/common/exceptions/CelInvalidArgumentException.java

@@ -24,4 +24,8 @@ public final class CelInvalidArgumentException extends CelRuntimeException {
   public CelInvalidArgumentException(Throwable cause) {
     super(cause, CelErrorCode.INVALID_ARGUMENT);
   }
+
+  public CelInvalidArgumentException(String message) {
+    super(message, CelErrorCode.INVALID_ARGUMENT);
+  }
 }

conformance/src/test/java/dev/cel/conformance/BUILD.bazel

@@ -175,11 +175,6 @@ _TESTS_TO_SKIP_PLANNER = [
     "timestamps/timestamp_range/sub_time_duration_under",
 
     # Skip until fixed.
-    "fields/qualified_identifier_resolution/map_key_float",
-    "fields/qualified_identifier_resolution/map_key_null",
-    "fields/qualified_identifier_resolution/map_value_repeat_key_heterogeneous",
-    "optionals/optionals/map_null_entry_no_such_key",
-    "optionals/optionals/map_present_key_invalid_field",
     "parse/receiver_function_names",
     "proto2/extensions_get/package_scoped_test_all_types_ext",
     "proto2/extensions_get/package_scoped_repeated_test_all_types",

runtime/src/main/java/dev/cel/runtime/planner/BUILD.bazel

@@ -334,6 +334,7 @@ java_library(
         ":localized_evaluation_exception",
         ":planned_interpretable",
         "//common/exceptions:duplicate_key",
+        "//common/exceptions:invalid_argument",
         "//runtime:evaluation_exception",
         "//runtime:interpretable",
         "@maven//:com_google_errorprone_error_prone_annotations",

runtime/src/main/java/dev/cel/runtime/planner/EvalCreateMap.java

@@ -17,8 +17,10 @@
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
+import com.google.common.primitives.UnsignedLong;
 import com.google.errorprone.annotations.Immutable;
 import dev.cel.common.exceptions.CelDuplicateKeyException;
+import dev.cel.common.exceptions.CelInvalidArgumentException;
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.GlobalResolver;
 import java.util.HashSet;
@@ -46,13 +48,37 @@ public Object eval(GlobalResolver resolver, ExecutionFrame frame) throws CelEval
     HashSet<Object> keysSeen = Sets.newHashSetWithExpectedSize(keys.length);
 
     for (int i = 0; i < keys.length; i++) {
-      Object key = keys[i].eval(resolver, frame);
-      Object val = values[i].eval(resolver, frame);
+      PlannedInterpretable keyInterpretable = keys[i];
+      Object key = keyInterpretable.eval(resolver, frame);
+      if (!(key instanceof String
+          || key instanceof Long
+          || key instanceof UnsignedLong
+          || key instanceof Boolean)) {
+        throw new LocalizedEvaluationException(
+            new CelInvalidArgumentException("Unsupported key type: " + key),
+            keyInterpretable.exprId());
+      }
 
-      if (!keysSeen.add(key)) {
-        throw new LocalizedEvaluationException(CelDuplicateKeyException.of(key), keys[i].exprId());
+      boolean isDuplicate = !keysSeen.add(key);
+      if (!isDuplicate) {
+        if (key instanceof Long) {
+          long longVal = (Long) key;
+          if (longVal >= 0) {
+            isDuplicate = keysSeen.contains(UnsignedLong.valueOf(longVal));
+          }
+        } else if (key instanceof UnsignedLong) {
+          UnsignedLong ulongVal = (UnsignedLong) key;
+          isDuplicate = keysSeen.contains(ulongVal.longValue());
+        }
       }
 
+      if (isDuplicate) {
+        throw new LocalizedEvaluationException(
+            CelDuplicateKeyException.of(key), keyInterpretable.exprId());
+      }
+
+      Object val = values[i].eval(resolver, frame);
+
       if (isOptional[i]) {
         if (!(val instanceof Optional)) {
           throw new IllegalArgumentException(

runtime/src/main/java/dev/cel/runtime/planner/StringQualifier.java

@@ -15,6 +15,7 @@
 package dev.cel.runtime.planner;
 
 import dev.cel.common.exceptions.CelAttributeNotFoundException;
+import dev.cel.common.values.OptionalValue;
 import dev.cel.common.values.SelectableValue;
 import java.util.Map;
 
@@ -31,21 +32,34 @@ public String value() {
   @Override
   @SuppressWarnings("unchecked") // Qualifications on maps/structs must be a string
   public Object qualify(Object obj) {
+    if (obj instanceof OptionalValue) {
+      OptionalValue<?, ?> opt = (OptionalValue<?, ?>) obj;
+      if (!opt.isZeroValue()) {
+        Object inner = opt.value();
+        if (!(inner instanceof SelectableValue) && !(inner instanceof Map)) {
+          throw CelAttributeNotFoundException.forFieldResolution(value);
+        }
+      }
+    }
+
     if (obj instanceof SelectableValue) {
       return ((SelectableValue<String>) obj).select(value);
-    } else if (obj instanceof Map) {
-      Map<String, Object> map = (Map<String, Object>) obj;
-      if (!map.containsKey(value)) {
-        throw CelAttributeNotFoundException.forMissingMapKey(value);
-      }
+    }
 
+    if (obj instanceof Map) {
+      Map<?, ?> map = (Map<?, ?>) obj;
       Object mapVal = map.get(value);
 
-      if (mapVal == null) {
-        throw CelAttributeNotFoundException.of(
-            String.format("Map value cannot be null for key: %s", value));
+      if (mapVal != null) {
+        return mapVal;
+      }
+
+      if (!map.containsKey(value)) {
+        throw CelAttributeNotFoundException.forMissingMapKey(value);
       }
-      return map.get(value);
+
+      throw CelAttributeNotFoundException.of(
+          String.format("Map value cannot be null for key: %s", value));
     }
 
     throw CelAttributeNotFoundException.forFieldResolution(value);

runtime/src/test/java/dev/cel/runtime/planner/ProgramPlannerTest.java

@@ -390,6 +390,17 @@ public void plan_createMap_containsDuplicateKey_throws() throws Exception {
         .contains("evaluation error at <input>:20: duplicate map key [true]");
   }
 
+  @Test
+  public void plan_createMap_unsupportedKeyType_throws() throws Exception {
+    CelAbstractSyntaxTree ast = compile("{1.0: 'foo'}");
+    Program program = PLANNER.plan(ast);
+
+    CelEvaluationException e = assertThrows(CelEvaluationException.class, program::eval);
+    assertThat(e)
+        .hasMessageThat()
+        .contains("evaluation error at <input>:1: Unsupported key type: 1.0");
+  }
+
   @Test
   public void plan_createStruct() throws Exception {
     CelAbstractSyntaxTree ast = compile("cel.expr.conformance.proto3.TestAllTypes{}");