topdown: Handling default functions in Partial Eval (#7499)

Making Partial Eval (PE) respect default functions.

Before this fix, Rego functions with declared default values weren't respected by PE, and the default declaration was omitted from generated support modules.

Fixes: #7220

Signed-off-by: Johan Fylling <[email protected]>

Author: johanfylling

Diff for: v1/topdown/eval.go

@@ -2019,15 +2019,36 @@ func (e evalFunc) eval(iter unifyIterator) error {
 		return e.e.saveCall(argCount, e.terms, iter)
 	}
 
-	if e.e.partial() && (e.e.inliningControl.shallow || e.e.inliningControl.Disabled(e.ref, false)) {
-		// check if the function definitions, or any of the arguments
-		// contain something unknown
-		unknown := e.e.unknown(e.ref, e.e.bindings)
-		for i := 1; !unknown && i <= argCount; i++ {
-			unknown = e.e.unknown(e.terms[i], e.e.bindings)
+	if e.e.partial() {
+		var mustGenerateSupport bool
+
+		if defRule := e.ir.Default; defRule != nil {
+			// The presence of a default func might force us to generate support
+			if len(defRule.Head.Args) == len(e.terms)-1 {
+				// The function is called without collecting the result in an output term,
+				// therefore any successful evaluation of the function is of interest, including the default value ...
+				if ret := defRule.Head.Value; ret == nil || !ret.Equal(ast.InternedBooleanTerm(false)) {
+					// ... unless the default value is false,
+					mustGenerateSupport = true
+				}
+			} else {
+				// The function is called with an output term, therefore any successful evaluation of the function is of interest.
+				// NOTE: Because of how the compiler rewrites function calls, we can't know if the result value is compared
+				// to a constant value, so we can't be as clever as we are for rules.
+				mustGenerateSupport = true
+			}
 		}
-		if unknown {
-			return e.partialEvalSupport(argCount, iter)
+
+		if mustGenerateSupport || e.e.inliningControl.shallow || e.e.inliningControl.Disabled(e.ref, false) {
+			// check if the function definitions, or any of the arguments
+			// contain something unknown
+			unknown := e.e.unknown(e.ref, e.e.bindings)
+			for i := 1; !unknown && i <= argCount; i++ {
+				unknown = e.e.unknown(e.terms[i], e.e.bindings)
+			}
+			if unknown {
+				return e.partialEvalSupport(argCount, iter)
+			}
 		}
 	}
 
@@ -2226,6 +2247,13 @@ func (e evalFunc) partialEvalSupport(declArgsLen int, iter unifyIterator) error
 				return err
 			}
 		}
+
+		if e.ir.Default != nil {
+			err := e.partialEvalSupportRule(e.ir.Default, path)
+			if err != nil {
+				return err
+			}
+		}
 	}
 
 	if !e.e.saveSupport.Exists(path) { // we haven't saved anything, nothing to call
@@ -2274,8 +2302,9 @@ func (e evalFunc) partialEvalSupportRule(rule *ast.Rule, path ast.Ref) error {
 			}
 
 			e.e.saveSupport.Insert(path, &ast.Rule{
-				Head: head,
-				Body: plugged,
+				Head:    head,
+				Body:    plugged,
+				Default: rule.Default,
 			})
 		}
 		child.traceRedo(rule)

Diff for: v1/topdown/topdown_partial_test.go

@@ -4068,6 +4068,138 @@ func TestTopDownPartialEval(t *testing.T) {
 			},
 			wantQueries: []string{""}, // unconditional true
 		},
+
+		{
+			note:  "default function, result not collected (non-false default value)",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					default f(x) := true # return true if x.size is undefined
+					f(x) if {
+						x.size < 100
+					}
+					p if {
+						f(input.x)
+					}
+				`},
+			wantQueries: []string{"data.partial.test.f(input.x)"},
+			wantSupport: []string{
+				`package partial.test
+				
+				default f(__local0__3) = true
+				f(__local1__2) = true if { __local2__2 = __local1__2.size; lt(__local2__2, 100) }`,
+			},
+		},
+		{
+			note:  "default function, result not collected (false default value)",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					default f(x) := false
+					f(x) if {
+						x.size < 100
+					}
+					p if {
+						f(input.x)
+					}
+				`},
+			wantQueries: []string{"lt(input.x.size, 100)"},
+		},
+		{
+			note:  "default function, result comparison (same as default)",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					default f(x) := true # return true if x.size is undefined
+					f(x) if {
+						x.size < 100
+					}
+					p if {
+						f(input.x) == true
+					}
+				`},
+			wantQueries: []string{"data.partial.test.f(input.x, true)"},
+			wantSupport: []string{
+				`package partial.test
+				
+				default f(__local0__3) = true
+				f(__local1__2) = true if { __local3__2 = __local1__2.size; lt(__local3__2, 100) }`,
+			},
+		},
+		{
+			note:  "default function, result comparison (not same as default)",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					default f(x) := true # return true if x.size is undefined
+					f(x) := y if {
+						y := x.size < 100
+					}
+					p if {
+						f(input.x) == false
+					}
+				`},
+			wantQueries: []string{"data.partial.test.f(input.x, false)"},
+			wantSupport: []string{
+				`package partial.test
+				
+				default f(__local0__3) = true
+				f(__local1__2) = __local2__2 if { __local5__2 = __local1__2.size; lt(__local5__2, 100, __local3__2); __local2__2 = __local3__2 }`,
+			},
+		},
+		{
+			note:  "default function, saved result",
+			query: "data.test.p = x",
+			modules: []string{`package test
+					default f(x) := true # return true if x.size is undefined
+					f(x) if {
+						x.size < 100
+					}
+					p := x if {
+						x := f(input.x)
+					}
+				`},
+			wantQueries: []string{"data.partial.test.f(input.x, x)"},
+			wantSupport: []string{
+				`package partial.test
+				
+				default f(__local0__3) = true
+				f(__local1__2) = true if { __local4__2 = __local1__2.size; lt(__local4__2, 100) }`,
+			},
+		},
+		{
+			// This test case is redundant, but serves as a counter example to the test above.
+			// Inlining can happen as there is no default function to consider
+			note:  "default function (no default)",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					f(x) if {
+						x.size < 100
+					}
+					p if {
+						f(input)
+					}
+				`},
+			wantQueries: []string{"lt(input.size, 100)"},
+		},
+		{
+			note:  "default function, shallow inlining",
+			query: "data.test.p = true",
+			modules: []string{`package test
+					default f(x) := true # return true if x.size is undefined
+					f(x) if {
+						x.size < 100
+					}
+					p if {
+						f(input)
+					}
+				`},
+			shallow:     true,
+			wantQueries: []string{"data.partial.test.p = true"},
+			wantSupport: []string{
+				`package partial.test
+					
+					p = true if { __local3__1 = input; data.partial.test.f(__local3__1) }
+					default f(__local0__3) = true
+					f(__local1__2) = true if { __local2__2 = __local1__2.size; lt(__local2__2, 100) }`,
+			},
+		},
 	}
 
 	ctx := context.Background()