interpreter.PruneAst does not prune IterRange expression

[dbuduev]

Describe the bug
interpreter.PruneAst function ignores the IterRange field of a comprehension expression.

To Reproduce
Check which components this affects:

• parser
• checker
• interpreter

Sample expression and input that reproduces the issue:

users is a list of const strings, and R is an unknown attribute pattern.

users.filter(u, u.startsWith(R.attr.prefix))

Test setup:

package main

import (
	"fmt"
	"log"

	"github.com/google/cel-go/cel"
	"github.com/google/cel-go/checker/decls"
	"github.com/google/cel-go/interpreter"
	expr "google.golang.org/genproto/googleapis/api/expr/v1alpha1"
)

func main() {
	env, err := cel.NewEnv(cel.Declarations(
		decls.NewVar("R", decls.NewMapType(decls.String, decls.Dyn)),
		decls.NewVar("users", decls.NewListType(decls.String)),
	))
	exitOnErr(err)

	ast, iss := env.Compile(`users.filter(u, u.startsWith(R.attr.prefix))`)
	exitOnErr(iss.Err())

	prg, err := env.Program(ast, cel.EvalOptions(cel.OptTrackState, cel.OptPartialEval))
	exitOnErr(err)

	in := map[string]any{
		"users": []string{"alice", "bob"},
	}
	vars, err := cel.PartialVars(in, cel.AttributePattern("R").QualString("attr").Wildcard())
	exitOnErr(err)

	_, details, err := prg.Eval(vars)
	exitOnErr(err)

	pruned := interpreter.PruneAst(ast.Expr(), ast.SourceInfo().GetMacroCalls(), details.State()).Expr
	ce := pruned.ExprKind.(*expr.Expr_ComprehensionExpr).ComprehensionExpr
	printExpr("IterRange:", ce.IterRange)
}

func printExpr(name string, e *expr.Expr) {
	prunedAST := cel.ParsedExprToAst(&expr.ParsedExpr{Expr: e})
	s, err := cel.AstToString(prunedAST)
	exitOnErr(err)
	fmt.Println(name, s)
}

func exitOnErr(err error) {
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
}

Expected behavior
The output should be:

IterRange: ["alice", "bob"]

instead of

IterRange: users

Additional context
This works as expected in v0.15.2, but in v0.15.3, the pruner was modified to skip the IterRange expression. Is this a breaking change?

[TristonianJones]

@dbuduev Thanks for the report, you will need to enable macro call tracking to get the pruning to work reliably with all macros: https://github.com/google/cel-go/blob/v0.16.0/cel/options.go#L533

[dbuduev]

@TristonianJones Thanks for your prompt reply. I have enabled macro call tracking, but IterRange remains unevaluated. I passed the pruned expression into the parser.Unparse function, which indeed replaced users with a literal.

A modified example

package main

import (
	"fmt"
	"log"

	"github.com/google/cel-go/cel"
	"github.com/google/cel-go/checker/decls"
	"github.com/google/cel-go/interpreter"
	"github.com/google/cel-go/parser"
	expr "google.golang.org/genproto/googleapis/api/expr/v1alpha1"
)

func main() {
	env, err := cel.NewEnv(cel.EnableMacroCallTracking(), cel.Declarations(
		decls.NewVar("R", decls.NewMapType(decls.String, decls.Dyn)),
		decls.NewVar("users", decls.NewListType(decls.String)),
	))
	exitOnErr(err)

	ast, iss := env.Compile(`users.filter(u, u.startsWith(R.attr.prefix))`)
	exitOnErr(iss.Err())

	prg, err := env.Program(ast, cel.EvalOptions(cel.OptTrackState, cel.OptPartialEval))
	exitOnErr(err)

	in := map[string]any{
		"users": []string{"alice", "bob"},
	}
	vars, err := cel.PartialVars(in, cel.AttributePattern("R").QualString("attr").Wildcard())
	exitOnErr(err)

	_, details, err := prg.Eval(vars)
	exitOnErr(err)

	pruned := interpreter.PruneAst(ast.Expr(), ast.SourceInfo().GetMacroCalls(), details.State())
	ce := pruned.GetExpr().GetComprehensionExpr()
	printExpr("IterRange:", ce.IterRange)
	printExpr("LoopStep:", ce.LoopStep)
	fmt.Println("IterVar:", ce.IterVar)

	s, err := parser.Unparse(pruned.Expr, pruned.SourceInfo)
	exitOnErr(err)
	fmt.Println("Unparsed expr:", s)
}

func printExpr(name string, e *expr.Expr) {
	prunedAST := cel.ParsedExprToAst(&expr.ParsedExpr{Expr: e})
	s, err := cel.AstToString(prunedAST)
	exitOnErr(err)
	fmt.Println(name, s)
}

func exitOnErr(err error) {
	if err != nil {
		log.Fatalf("Error: %v", err)
	}
}

The output is:

IterRange: users
LoopStep: u.startsWith(R.attr.prefix) ? (__result__ + [u]) : __result__
IterVar: u
Unparsed expr: ["alice", "bob"].filter(u, "bob".startsWith(R.attr.prefix))

⚠️ parser.Unparse replaced an iteration variable with "bob" in the loop step expression.

[TristonianJones]

That's definitely problematic. I'll give this a look tomorrow. Sorry about the trouble with the residuals. They're a bit tricky to get right.

[dbuduev]

I have run the aforementioned example with this latest cel-go version in the master branch (github.com/google/cel-go v0.16.1-0.20230616212823-338b3c80e688).

1. IterRange is not pruned.
2. The issue with the parser.Unparse() function is fixed.

The output:

IterRange: users
LoopStep: u.startsWith(R.attr.prefix) ? (__result__ + [u]) : __result__
IterVar: u
Unparsed expr: ["alice", "bob"].filter(u, u.startsWith(R.attr.prefix))

[TristonianJones]

@dbuduev the test case in the report is included in the prune tests. Are you confirming the fix or asking for different behavior?

[dbuduev]

I see the test case in the pruner tests:

	{
		in: partialActivation(map[string]any{
			"users": []string{"alice", "bob"},
		}, NewAttributePattern("r").QualString("attr").Wildcard()),
		expr: `users.filter(u, r.attr.prefix.endsWith(u))`,
		out:  `["alice", "bob"].filter(u, r.attr.prefix.endsWith(u))`,
	},

Then the assert:
https://github.com/google/cel-go/blob/master/interpreter/prune_test.go#L468-L469

newExpr := PruneAst(ast.GetExpr(), ast.GetSourceInfo().GetMacroCalls(), state)
actual, err := parser.Unparse(newExpr.GetExpr(), newExpr.GetSourceInfo())
...
t.Errorf("prune[%d], diff: %s", i, test.DiffMessage("structure", actual, tst.out))

My point is that it is the parser.Unparse() function that substitutes the variable, not the PruneAst.

[TristonianJones]

I see, that's because it's the macro call information that's pruned, which is why the unparse renders it differently than you see in the main AST

[dbuduev]

In our use case, we are not using the unparse, but only the PruneAst, which has stopped pruning the comprehensions' IterRange since v0.15.3.
The unparse is smart and thus compensates for the change in the pruner tests.

Do you recommend always calling the unparse (and then parse again) after the pruning AST?

[TristonianJones]

Oh, I see. It just depends on the use case, I guess. Can you tell me more about what you're trying to do?

If you really need the comprehension pruning, I can add it back. I thought the issue was with combo of prune / unparse since the output was not semantically correct.

[dbuduev]

We have a use case when we don't have all information to evaluate a CEL expression, so we use PruneAst to partially evaluate it, and then convert the residual expression to a custom form, which then can be evaluated later. Here is the full description.

If you really need the comprehension pruning, I can add it back.

That would be great, but we likely can evaluate it separately if not.

EDIT: I wasn't sure if the change in PruneAst was deliberate. Thank you for the clarification.