From 178e9a80d67069ac5deb0b0e12f5670d640a8594 Mon Sep 17 00:00:00 2001
From: Sokwhan Huh <sokwhan@google.com>
Date: Wed, 28 Feb 2024 17:35:06 -0800
Subject: [PATCH] Fix optional list/map indexing with unknowns. Fix optional
 map indexing via field selection.

PiperOrigin-RevId: 611285327
---
 .../java/dev/cel/checker/ExprChecker.java     |  2 +-
 .../test/java/dev/cel/extensions/BUILD.bazel  |  1 +
 .../extensions/CelOptionalLibraryTest.java    | 47 +++++++++++++++++++
 .../dev/cel/runtime/DefaultInterpreter.java   |  4 +-
 4 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/checker/src/main/java/dev/cel/checker/ExprChecker.java b/checker/src/main/java/dev/cel/checker/ExprChecker.java
index 2506466ce..cbebcd238 100644
--- a/checker/src/main/java/dev/cel/checker/ExprChecker.java
+++ b/checker/src/main/java/dev/cel/checker/ExprChecker.java
@@ -705,7 +705,7 @@ private CelExpr visitOptionalCall(CelExpr expr, CelExpr.CelCall call) {
     CelExpr visitedOperand = visit(operand);
     if (namespacedDeclarations && !operand.equals(visitedOperand)) {
       // Subtree has been rewritten. Replace the operand.
-      expr = replaceSelectOperandSubtree(expr, visitedOperand);
+      expr = replaceCallArgumentSubtree(expr, visitedOperand, 0);
     }
     CelType resultType = visitSelectField(expr, operand, field.constant().stringValue(), true);
     env.setType(expr, resultType);
diff --git a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
index c71df2d8a..7f572f58f 100644
--- a/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
+++ b/extensions/src/test/java/dev/cel/extensions/BUILD.bazel
@@ -26,6 +26,7 @@ java_library(
         "//extensions:strings",
         "//parser:macro",
         "//runtime",
+        "//runtime:interpreter_util",
         "@maven//:com_google_guava_guava",
         "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_google_protobuf_protobuf_java_util",
diff --git a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
index 8cb7a11f7..31df7a325 100644
--- a/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
+++ b/extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java
@@ -46,6 +46,7 @@
 import dev.cel.runtime.CelEvaluationException;
 import dev.cel.runtime.CelRuntime;
 import dev.cel.runtime.CelRuntime.CelFunctionBinding;
+import dev.cel.runtime.InterpreterUtil;
 import dev.cel.testing.testdata.proto2.TestAllTypesProto.TestAllTypes;
 import dev.cel.testing.testdata.proto2.TestAllTypesProto.TestAllTypes.NestedMessage;
 import java.util.List;
@@ -766,6 +767,38 @@ public void optionalIndex_onMap_returnsOptionalValue() throws Exception {
     assertThat(result).isEqualTo(Optional.of("goodbye"));
   }
 
+  @Test
+  @TestParameters("{source: '{?x: optional.of(1)}'}")
+  @TestParameters("{source: '{?1: x}'}")
+  @TestParameters("{source: '{?x: x}'}")
+  public void optionalIndex_onMapWithUnknownInput_returnsUnknownResult(String source)
+      throws Exception {
+    Cel cel = newCelBuilder().addVar("x", OptionalType.create(SimpleType.INT)).build();
+    CelAbstractSyntaxTree ast = cel.compile(source).getAst();
+
+    Object result = cel.createProgram(ast).eval();
+
+    assertThat(InterpreterUtil.isUnknown(result)).isTrue();
+  }
+
+  @Test
+  public void optionalIndex_onOptionalMapUsingFieldSelection_returnsOptionalValue()
+      throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .addVar(
+                "m",
+                MapType.create(
+                    SimpleType.STRING, MapType.create(SimpleType.STRING, SimpleType.STRING)))
+            .setResultType(OptionalType.create(SimpleType.STRING))
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("{?'key': optional.of('test')}.?key").getAst();
+
+    Object result = cel.createProgram(ast).eval();
+
+    assertThat(result).isEqualTo(Optional.of("test"));
+  }
+
   @Test
   public void optionalIndex_onList_returnsOptionalEmpty() throws Exception {
     Cel cel =
@@ -824,6 +857,20 @@ public void optionalIndex_onOptionalList_returnsOptionalValue() throws Exception
     assertThat(result).isEqualTo(Optional.of("hello"));
   }
 
+  @Test
+  public void optionalIndex_onListWithUnknownInput_returnsUnknownResult() throws Exception {
+    Cel cel =
+        newCelBuilder()
+            .addVar("x", OptionalType.create(SimpleType.INT))
+            .setResultType(ListType.create(SimpleType.INT))
+            .build();
+    CelAbstractSyntaxTree ast = cel.compile("[?x]").getAst();
+
+    Object result = cel.createProgram(ast).eval();
+
+    assertThat(InterpreterUtil.isUnknown(result)).isTrue();
+  }
+
   @Test
   public void traditionalIndex_onOptionalList_returnsOptionalEmpty() throws Exception {
     Cel cel =
diff --git a/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java b/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
index 3c4c0003a..bf6af55ee 100644
--- a/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
+++ b/runtime/src/main/java/dev/cel/runtime/DefaultInterpreter.java
@@ -672,7 +672,7 @@ private IntermediateResult evalList(
 
         argChecker.checkArg(evaluatedElement);
         Object value = evaluatedElement.value();
-        if (optionalIndicesSet.contains(i)) {
+        if (optionalIndicesSet.contains(i) && !isUnknownValue(value)) {
           Optional<?> optionalVal = (Optional<?>) value;
           if (!optionalVal.isPresent()) {
             continue;
@@ -712,7 +712,7 @@ private IntermediateResult evalMap(ExecutionFrame frame, CelCreateMap mapExpr)
         }
 
         Object value = valueResult.value();
-        if (entry.optionalEntry()) {
+        if (entry.optionalEntry() && !isUnknownValue(value)) {
           Optional<?> optionalVal = (Optional<?>) value;
           if (!optionalVal.isPresent()) {
             // This is a no-op currently but will be semantically correct when extended proto