Fix last optional element to be retained as an optional index while folding (#841)

* Fix last optional element to be marked as an optional index while folding

Author: l46kok

cel/folding.go

@@ -248,23 +248,26 @@ func pruneOptionalListElements(ctx *OptimizerContext, e ast.Expr) {
 	}
 	updatedElems := []ast.Expr{}
 	updatedIndices := []int32{}
-	for i, e := range elems {
-		if !l.IsOptional(int32(i)) {
+	newOptIndex := -1
+	for _, e := range elems {
+		newOptIndex++
+		if !l.IsOptional(int32(newOptIndex)) {
 			updatedElems = append(updatedElems, e)
 			continue
 		}
 		if e.Kind() != ast.LiteralKind {
 			updatedElems = append(updatedElems, e)
-			updatedIndices = append(updatedIndices, int32(i))
+			updatedIndices = append(updatedIndices, int32(newOptIndex))
 			continue
 		}
 		optElemVal, ok := e.AsLiteral().(*types.Optional)
 		if !ok {
 			updatedElems = append(updatedElems, e)
-			updatedIndices = append(updatedIndices, int32(i))
+			updatedIndices = append(updatedIndices, int32(newOptIndex))
 			continue
 		}
 		if !optElemVal.HasValue() {
+			newOptIndex-- // Skipping causes the list to get smaller.
 			continue
 		}
 		e.SetKindCase(ctx.NewLiteral(optElemVal.GetValue()))

cel/folding_test.go

@@ -129,10 +129,26 @@ func TestConstantFoldingOptimizer(t *testing.T) {
 			expr:   `[google.expr.proto3.test.TestAllTypes{single_int32: 2 + 3}].map(i, i)[0]`,
 			folded: `google.expr.proto3.test.TestAllTypes{single_int32: 5}`,
 		},
+		{
+			expr:   `[?optional.ofNonZeroValue(0)]`,
+			folded: `[]`,
+		},
 		{
 			expr:   `[1, ?optional.ofNonZeroValue(0)]`,
 			folded: `[1]`,
 		},
+		{
+			expr:   `[optional.none(), ?x]`,
+			folded: `[optional.none(), ?x]`,
+		},
+		{
+			expr:   `[?optional.none(), ?x]`,
+			folded: `[?x]`,
+		},
+		{
+			expr:   `[1, x, ?optional.ofNonZeroValue(0), ?x.?y]`,
+			folded: `[1, x, ?x.?y]`,
+		},
 		{
 			expr:   `[1, x, ?optional.ofNonZeroValue(3), ?x.?y]`,
 			folded: `[1, x, 3, ?x.?y]`,