Ensure logical operator fold expressions are bool

cel/folding.go

@@ -95,7 +95,7 @@ func (opt *constantFoldingOptimizer) Optimize(ctx *OptimizerContext, a *ast.AST)
 		for _, fold := range foldableExprs {
 			// If the expression could be folded because it's a non-strict call, and the
 			// branches are pruned, continue to the next fold.
-			if fold.Kind() == ast.CallKind && maybePruneBranches(ctx, fold) {
+			if fold.Kind() == ast.CallKind && maybePruneBranches(ctx, a, fold) {
 				continue
 			}
 			// Late-bound function calls cannot be folded.
@@ -202,12 +202,12 @@ func isLateBoundFunctionCall(ctx *OptimizerContext, expr ast.Expr) bool {
 // a branch can be removed. Evaluation will naturally prune logical and / or calls,
 // but conditional will not be pruned cleanly, so this is one small area where the
 // constant folding step reimplements a portion of the evaluator.
-func maybePruneBranches(ctx *OptimizerContext, expr ast.NavigableExpr) bool {
+func maybePruneBranches(ctx *OptimizerContext, a *ast.AST, expr ast.NavigableExpr) bool {
 	call := expr.AsCall()
 	args := call.Args()
 	switch call.FunctionName() {
 	case operators.LogicalAnd, operators.LogicalOr:
-		return maybeShortcircuitLogic(ctx, call.FunctionName(), args, expr)
+		return maybeShortcircuitLogic(ctx, a, call.FunctionName(), args, expr)
 	case operators.Conditional:
 		cond := args[0]
 		truthy := args[1]
@@ -248,7 +248,7 @@ func maybePruneBranches(ctx *OptimizerContext, expr ast.NavigableExpr) bool {
 	return false
 }
 
-func maybeShortcircuitLogic(ctx *OptimizerContext, function string, args []ast.Expr, expr ast.NavigableExpr) bool {
+func maybeShortcircuitLogic(ctx *OptimizerContext, a *ast.AST, function string, args []ast.Expr, expr ast.NavigableExpr) bool {
 	shortcircuit := types.False
 	skip := types.True
 	if function == operators.LogicalOr {
@@ -271,17 +271,31 @@ func maybeShortcircuitLogic(ctx *OptimizerContext, function string, args []ast.E
 	}
 	if len(newArgs) == 0 {
 		newArgs = append(newArgs, args[0])
-		ctx.UpdateExpr(expr, newArgs[0])
-		return true
+	}
+	if len(newArgs) == len(args) {
+		return false
 	}
 	if len(newArgs) == 1 {
+		if !isBoolType(a, newArgs[0]) {
+			return false
+		}
 		ctx.UpdateExpr(expr, newArgs[0])
 		return true
 	}
 	ctx.UpdateExpr(expr, ctx.NewCall(function, newArgs...))
 	return true
 }
 
+func isBoolType(a *ast.AST, e ast.Expr) bool {
+	if a != nil && a.GetType(e.ID()) == types.BoolType {
+		return true
+	}
+	if e.Kind() == ast.LiteralKind && e.AsLiteral().Type() == types.BoolType {
+		return true
+	}
+	return false
+}
+
 // pruneOptionalElements works from the bottom up to resolve optional elements within
 // aggregate literals.
 //

cel/folding_test.go

@@ -214,11 +214,11 @@ func TestConstantFoldingOptimizer(t *testing.T) {
 		},
 		{
 			expr:   `true && x`,
-			folded: `x`,
+			folded: `true && x`,
 		},
 		{
 			expr:   `x && true`,
-			folded: `x`,
+			folded: `x && true`,
 		},
 		{
 			expr:   `false && x`,
@@ -238,19 +238,59 @@ func TestConstantFoldingOptimizer(t *testing.T) {
 		},
 		{
 			expr:   `false || x`,
-			folded: `x`,
+			folded: `false || x`,
 		},
 		{
 			expr:   `x || false`,
-			folded: `x`,
+			folded: `x || false`,
+		},
+		{
+			expr:   `true && b`,
+			folded: `b`,
+		},
+		{
+			expr:   `b && true`,
+			folded: `b`,
+		},
+		{
+			expr:   `false || b`,
+			folded: `b`,
+		},
+		{
+			expr:   `b || false`,
+			folded: `b`,
+		},
+		{
+			expr:   `false || x`,
+			folded: `false || x`,
+		},
+		{
+			expr:   `x || false`,
+			folded: `x || false`,
+		},
+		{
+			expr:   `true && x`,
+			folded: `true && x`,
+		},
+		{
+			expr:   `x && true`,
+			folded: `x && true`,
 		},
 		{
 			expr:   `true && x && true && x`,
-			folded: `x && x`,
+			folded: `true && x && true && x`,
 		},
 		{
 			expr:   `false || x || false || x`,
-			folded: `x || x`,
+			folded: `false || x || false || x`,
+		},
+		{
+			expr:   `true && b && true && b`,
+			folded: `b && b`,
+		},
+		{
+			expr:   `false || b || false || b`,
+			folded: `b || b`,
 		},
 		{
 			expr:   `true && true`,
@@ -363,7 +403,11 @@ func TestConstantFoldingOptimizer(t *testing.T) {
 		},
 		{
 			expr:   `false || x || false || y`,
-			folded: `x || y`,
+			folded: `false || x || false || y`,
+		},
+		{
+			expr:   `false || b || false || b`,
+			folded: `b || b`,
 		},
 		{
 			expr:   `true ? (false ? x + 1 : x + 2) : x`,
@@ -497,6 +541,7 @@ func TestConstantFoldingOptimizer(t *testing.T) {
 		Types(&proto3pb.TestAllTypes{}),
 		Variable("x", DynType),
 		Variable("y", DynType),
+		Variable("b", BoolType),
 		// work around different package convention in piper vs github.
 		// google.expr.proto3.test.ImportedGlobalEnum.IMPORT_BAZ
 		Constant("c", IntType, types.Int(2)),