Skip to content

Commit 827eea1

Browse files
jnthntatumcopybara-github
authored andcommitted
Avoid copying complex target in optMap/optFlatMap
cross ref: cel-expr/cel-go#1387 PiperOrigin-RevId: 954914524
1 parent 01ac8a5 commit 827eea1

3 files changed

Lines changed: 131 additions & 8 deletions

File tree

extensions/src/main/java/dev/cel/extensions/CelOptionalLibrary.java

Lines changed: 67 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ static CelExtensionLibrary<CelOptionalLibrary> library() {
297297
public static final CelOptionalLibrary INSTANCE = CelOptionalLibrary.library().latest();
298298

299299
private static final String UNUSED_ITER_VAR = "#unused";
300+
private static final String OPTIONAL_MAP_VAR = "@target";
300301

301302
private final int version;
302303
private final ImmutableSet<CelFunctionDecl> functions;
@@ -524,21 +525,51 @@ private static Optional<CelExpr> expandOptMap(
524525
CelExpr mapExpr = checkNotNull(arguments.get(1));
525526
String varName = varIdent.ident().name();
526527

527-
return Optional.of(
528+
if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
529+
return Optional.of(
530+
exprFactory.newGlobalCall(
531+
Operator.CONDITIONAL.getFunction(),
532+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
533+
exprFactory.newGlobalCall(
534+
OPTIONAL_OF.getFunction(),
535+
exprFactory.fold(
536+
UNUSED_ITER_VAR,
537+
exprFactory.newList(),
538+
varName,
539+
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
540+
exprFactory.newBoolLiteral(true),
541+
exprFactory.newIdentifier(varName),
542+
mapExpr)),
543+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
544+
}
545+
546+
CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
547+
CelExpr localVarCopy = exprFactory.copy(localVar);
548+
CelExpr conditionalExpr =
528549
exprFactory.newGlobalCall(
529550
Operator.CONDITIONAL.getFunction(),
530-
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
551+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
531552
exprFactory.newGlobalCall(
532553
OPTIONAL_OF.getFunction(),
533554
exprFactory.fold(
534555
UNUSED_ITER_VAR,
535556
exprFactory.newList(),
536557
varName,
537-
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
558+
exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
538559
exprFactory.newBoolLiteral(true),
539560
exprFactory.newIdentifier(varName),
540561
mapExpr)),
541-
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
562+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
563+
564+
return Optional.of(
565+
exprFactory.fold(
566+
UNUSED_ITER_VAR,
567+
exprFactory.newList(),
568+
OPTIONAL_MAP_VAR,
569+
target,
570+
exprFactory.newBoolLiteral(false),
571+
exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
572+
conditionalExpr));
542573
}
543574

544575
private static Optional<CelExpr> expandOptFlatMap(
@@ -558,19 +589,47 @@ private static Optional<CelExpr> expandOptFlatMap(
558589
CelExpr mapExpr = checkNotNull(arguments.get(1));
559590
String varName = varIdent.ident().name();
560591

561-
return Optional.of(
592+
if (target.exprKind().getKind() == CelExpr.ExprKind.Kind.IDENT) {
593+
return Optional.of(
594+
exprFactory.newGlobalCall(
595+
Operator.CONDITIONAL.getFunction(),
596+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
597+
exprFactory.fold(
598+
UNUSED_ITER_VAR,
599+
exprFactory.newList(),
600+
varName,
601+
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
602+
exprFactory.newBoolLiteral(true),
603+
exprFactory.newIdentifier(varName),
604+
mapExpr),
605+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
606+
}
607+
608+
CelExpr localVar = exprFactory.newIdentifier(OPTIONAL_MAP_VAR);
609+
CelExpr localVarCopy = exprFactory.copy(localVar);
610+
CelExpr conditionalExpr =
562611
exprFactory.newGlobalCall(
563612
Operator.CONDITIONAL.getFunction(),
564-
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), target),
613+
exprFactory.newReceiverCall(HAS_VALUE.getFunction(), localVar),
565614
exprFactory.fold(
566615
UNUSED_ITER_VAR,
567616
exprFactory.newList(),
568617
varName,
569-
exprFactory.newReceiverCall(VALUE.getFunction(), exprFactory.copy(target)),
618+
exprFactory.newReceiverCall(VALUE.getFunction(), localVarCopy),
570619
exprFactory.newBoolLiteral(true),
571620
exprFactory.newIdentifier(varName),
572621
mapExpr),
573-
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction())));
622+
exprFactory.newGlobalCall(OPTIONAL_NONE.getFunction()));
623+
624+
return Optional.of(
625+
exprFactory.fold(
626+
UNUSED_ITER_VAR,
627+
exprFactory.newList(),
628+
OPTIONAL_MAP_VAR,
629+
target,
630+
exprFactory.newBoolLiteral(false),
631+
exprFactory.newIdentifier(OPTIONAL_MAP_VAR),
632+
conditionalExpr));
574633
}
575634

576635
private static Object indexOptionalMap(

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ java_library(
1717
"//common:compiler_common",
1818
"//common:container",
1919
"//common:options",
20+
"//common/ast",
2021
"//common/exceptions:attribute_not_found",
2122
"//common/exceptions:divide_by_zero",
2223
"//common/exceptions:index_out_of_bounds",

extensions/src/test/java/dev/cel/extensions/CelOptionalLibraryTest.java

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import dev.cel.common.CelOverloadDecl;
3535
import dev.cel.common.CelValidationException;
3636
import dev.cel.common.CelVarDecl;
37+
import dev.cel.common.ast.CelExpr;
3738
import dev.cel.common.types.CelType;
3839
import dev.cel.common.types.ListType;
3940
import dev.cel.common.types.MapType;
@@ -1571,6 +1572,68 @@ public void optionalFlatMapMacro_receiverHasValue_returnsOptionalValue() throws
15711572
assertThat(result).hasValue(43L);
15721573
}
15731574

1575+
@Test
1576+
public void optionalMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
1577+
Cel cel =
1578+
newCelBuilder()
1579+
.addVar("x", OptionalType.create(SimpleType.INT))
1580+
.setResultType(OptionalType.create(SimpleType.INT))
1581+
.build();
1582+
CelAbstractSyntaxTree ast = compile(cel, "x.optMap(y, y + 1)");
1583+
1584+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
1585+
}
1586+
1587+
@Test
1588+
public void optionalMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
1589+
Cel cel =
1590+
newCelBuilder()
1591+
.setResultType(OptionalType.create(SimpleType.INT))
1592+
.addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
1593+
.build();
1594+
CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optMap(y, y + 1)");
1595+
1596+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
1597+
assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
1598+
1599+
Optional<Long> result =
1600+
(Optional<Long>)
1601+
cel.createProgram(ast)
1602+
.eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
1603+
assertThat(result).hasValue(43L);
1604+
}
1605+
1606+
@Test
1607+
public void optionalFlatMapMacro_simpleTarget_notWrappedInComprehension() throws Exception {
1608+
Cel cel =
1609+
newCelBuilder()
1610+
.addVar("x", OptionalType.create(SimpleType.INT))
1611+
.setResultType(OptionalType.create(SimpleType.INT))
1612+
.build();
1613+
CelAbstractSyntaxTree ast = compile(cel, "x.optFlatMap(y, optional.of(y + 1))");
1614+
1615+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.CALL);
1616+
}
1617+
1618+
@Test
1619+
public void optionalFlatMapMacro_complexTarget_astWrappedInComprehension() throws Exception {
1620+
Cel cel =
1621+
newCelBuilder()
1622+
.setResultType(OptionalType.create(SimpleType.INT))
1623+
.addVar("msg", StructTypeReference.create(TestAllTypes.getDescriptor().getFullName()))
1624+
.build();
1625+
CelAbstractSyntaxTree ast = compile(cel, "msg.?single_int32.optFlatMap(y, optional.of(y + 1))");
1626+
1627+
assertThat(ast.getExpr().exprKind().getKind()).isEqualTo(CelExpr.ExprKind.Kind.COMPREHENSION);
1628+
assertThat(ast.getExpr().comprehension().accuVar()).isEqualTo("@target");
1629+
1630+
Optional<Long> result =
1631+
(Optional<Long>)
1632+
cel.createProgram(ast)
1633+
.eval(ImmutableMap.of("msg", TestAllTypes.newBuilder().setSingleInt32(42).build()));
1634+
assertThat(result).hasValue(43L);
1635+
}
1636+
15741637
@Test
15751638
public void optionalFlatMapMacro_withOptionalOfNonZeroValue_optionalEmptyWhenValueIsZero()
15761639
throws Exception {

0 commit comments

Comments
 (0)