Update PruneAst to support constants of optional type (#1109)

dbuduev · web-flow · commit 91fb306ef06a · 2025-01-22T14:36:49.000-08:00
* Update PruneAst and unparser to support optional types constants.

Signed-off-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;

* recursive definitions

Signed-off-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;

* add prune tests for &lt;list&gt;.last()

Signed-off-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;

* clean-up

Signed-off-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;

---------

Signed-off-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;
Co-authored-by: Dennis Buduev &lt;dbuduev@users.noreply.github.com&gt;
diff --git a/interpreter/prune.go b/interpreter/prune.go
@@ -88,7 +88,7 @@ func PruneAst(expr ast.Expr, macroCalls map[int64]ast.Expr, state EvalState) *as
 
 func (p *astPruner) maybeCreateLiteral(id int64, val ref.Val) (ast.Expr, bool) {
 	switch v := val.(type) {
-	case types.Bool, types.Bytes, types.Double, types.Int, types.Null, types.String, types.Uint:
+	case types.Bool, types.Bytes, types.Double, types.Int, types.Null, types.String, types.Uint, *types.Optional:
 		p.state.SetValue(id, val)
 		return p.NewLiteral(id, val), true
 	case types.Duration:
diff --git a/interpreter/prune_test.go b/interpreter/prune_test.go
@@ -23,6 +23,7 @@ import (
 	"github.com/google/cel-go/common/decls"
 	"github.com/google/cel-go/common/types"
 	"github.com/google/cel-go/common/types/ref"
+	"github.com/google/cel-go/common/types/traits"
 	"github.com/google/cel-go/parser"
 	"github.com/google/cel-go/test"
 
@@ -216,6 +217,36 @@ var testCases = []testInfo{
 		expr: `a.?b`,
 		out:  `a.?b`,
 	},
+	{
+		in:   partialActivation(map[string]any{"a": map[string]any{"b": 10}}),
+		expr: `a.?b`,
+		out:  `optional.of(10)`,
+	},
+	{
+		in:   partialActivation(map[string]any{"a": map[string]any{"b": 10}}),
+		expr: `a[?"b"]`,
+		out:  `optional.of(10)`,
+	},
+	{
+		in:   unknownActivation(),
+		expr: `{'b': optional.of(10)}.?b`,
+		out:  `optional.of(optional.of(10))`,
+	},
+	{
+		in:   partialActivation(map[string]any{"a": map[string]any{}}),
+		expr: `a.?b`,
+		out:  `optional.none()`,
+	},
+	{
+		in:   unknownActivation(),
+		expr: `[10].last()`,
+		out:  "optional.of(10)",
+	},
+	{
+		in:   unknownActivation(),
+		expr: `[].last()`,
+		out:  "optional.none()",
+	},
 	{
 		in:   unknownActivation("a"),
 		expr: `a[?"b"]`,
@@ -561,5 +592,18 @@ func optionalDecls(t *testing.T) []*decls.FunctionDecl {
 				types.NewTypeParamType("K"),
 			}, optionalType),
 		),
+		funcDecl(t, "last", decls.Overload("list_last", []*types.Type{paramType}, optionalType,
+			decls.UnaryBinding(func(v ref.Val) ref.Val {
+				list := v.(traits.Lister)
+				sz := list.Size().Value().(int64)
+
+				if sz == 0 {
+					return types.OptionalNone
+				}
+
+				return types.OptionalOf(list.Get(types.Int(sz - 1)))
+			}),
+		),
+		),
 	}
 }
diff --git a/parser/unparser.go b/parser/unparser.go
@@ -24,6 +24,7 @@ import (
 	"github.com/google/cel-go/common/ast"
 	"github.com/google/cel-go/common/operators"
 	"github.com/google/cel-go/common/types"
+	"github.com/google/cel-go/common/types/ref"
 )
 
 // Unparse takes an input expression and source position information and generates a human-readable
@@ -273,8 +274,17 @@ func (un *unparser) visitCallUnary(expr ast.Expr) error {
 	return un.visitMaybeNested(args[0], nested)
 }
 
-func (un *unparser) visitConst(expr ast.Expr) error {
-	val := expr.AsLiteral()
+func (un *unparser) visitConstVal(val ref.Val) error {
+	optional := false
+	if optVal, ok := val.(*types.Optional); ok {
+		if !optVal.HasValue() {
+			un.str.WriteString("optional.none()")
+			return nil
+		}
+		optional = true
+		un.str.WriteString("optional.of(")
+		val = optVal.GetValue()
+	}
 	switch val := val.(type) {
 	case types.Bool:
 		un.str.WriteString(strconv.FormatBool(bool(val)))
@@ -303,7 +313,21 @@ func (un *unparser) visitConst(expr ast.Expr) error {
 		ui := strconv.FormatUint(uint64(val), 10)
 		un.str.WriteString(ui)
 		un.str.WriteString("u")
+	case *types.Optional:
+		if err := un.visitConstVal(val); err != nil {
+			return err
+		}
 	default:
+		return errors.New("unsupported constant")
+	}
+	if optional {
+		un.str.WriteString(")")
+	}
+	return nil
+}
+func (un *unparser) visitConst(expr ast.Expr) error {
+	val := expr.AsLiteral()
+	if err := un.visitConstVal(val); err != nil {
 		return fmt.Errorf("unsupported constant: %v", expr)
 	}
 	return nil
