ast+topdown+planner: replacement of non-built-in functions via 'with'

Follow-up to open-policy-agent#4540

We can now mock functions that are user-defined:

    package test

    f(_) = 1 {
        input.x = "x"
    }
    p = y {
        y := f(1) with f as 2
    }

...following the same scoping rules as laid out for built-in mocks.
The replacement can be a value (replacing all calls), or a built-in,
or another non-built-in function.

Also addresses bugs in the previous slice:
* topdown/evalCall: account for empty rules result from indexer
* topdown/eval: capture value replacement in PE could panic

Note: in PE, we now drop 'with' for function mocks of any kind:

These are always fully replaced in the saved support modules, so
this should be OK.

When keeping them, we'd also have to either copy the existing definitions
into the support module; or create a function stub in it.

Fixes open-policy-agent#4449.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

ast/compile.go

@@ -4798,6 +4798,14 @@ func rewriteWithModifier(c *Compiler, f *equalityFactory, expr *Expr) ([]*Expr,
 
 func validateWith(c *Compiler, expr *Expr, i int) (bool, *Error) {
 	target, value := expr.With[i].Target, expr.With[i].Value
+
+	// Ensure that values that are built-ins are rewritten to Ref (not Var)
+	if v, ok := value.Value.(Var); ok {
+		if _, ok := c.builtins[v.String()]; ok {
+			value.Value = Ref([]*Term{NewTerm(v)})
+		}
+	}
+
 	switch {
 	case isDataRef(target):
 		ref := target.Value.(Ref)
@@ -4813,11 +4821,16 @@ func validateWith(c *Compiler, expr *Expr, i int) (bool, *Error) {
 		}
 
 		if node != nil {
+			// NOTE(sr): at this point in the compiler stages, we don't have a fully-populated
+			// TypeEnv yet -- so we have to make do with this check to see if the replacement
+			// target is a function. It's probably wrong for arity-0 functions, but those are
+			// and edge case anyways.
 			if child := node.Child(ref[len(ref)-1].Value); child != nil {
-				for _, value := range child.Values {
-					if len(value.(*Rule).Head.Args) > 0 {
-						// TODO(sr): UDF
-						return false, NewError(CompileErr, target.Loc(), "with keyword used on non-built-in function")
+				for _, v := range child.Values {
+					if len(v.(*Rule).Head.Args) > 0 {
+						if validateWithFunctionValue(c.builtins, c.RuleTree, value) {
+							return false, nil
+						}
 					}
 				}
 			}
@@ -4830,29 +4843,18 @@ func validateWith(c *Compiler, expr *Expr, i int) (bool, *Error) {
 		if v, ok := target.Value.(Var); ok {
 			target.Value = Ref([]*Term{NewTerm(v)})
 		}
-		if v, ok := value.Value.(Var); ok {
-			if _, ok := c.builtins[v.String()]; ok {
-				value.Value = Ref([]*Term{NewTerm(v)})
-			}
-		}
 
 		targetRef := target.Value.(Ref)
 		bi := c.builtins[targetRef.String()] // safe because isBuiltinRefOrVar checked this
 		if err := validateWithBuiltinTarget(bi, targetRef, target.Loc()); err != nil {
 			return false, err
 		}
 
-		if v, ok := value.Value.(Ref); ok {
-			if c.RuleTree.Find(v) != nil { // ref exists in rule tree
-				return false, nil
-			}
-			if _, ok := c.builtins[v.String()]; ok { // built-in replaced by other built-in
-				return false, nil
-			}
+		if validateWithFunctionValue(c.builtins, c.RuleTree, value) {
+			return false, nil
 		}
-
 	default:
-		return false, NewError(TypeErr, target.Location, "with keyword target must reference existing %v, %v, or a built-in function", InputRootDocument, DefaultRootDocument)
+		return false, NewError(TypeErr, target.Location, "with keyword target must reference existing %v, %v, or a function", InputRootDocument, DefaultRootDocument)
 	}
 	return requiresEval(value), nil
 }
@@ -4878,6 +4880,15 @@ func validateWithBuiltinTarget(bi *Builtin, target Ref, loc *location.Location)
 	return nil
 }
 
+func validateWithFunctionValue(bs map[string]*Builtin, ruleTree *TreeNode, value *Term) bool {
+	if v, ok := value.Value.(Ref); ok {
+		if ruleTree.Find(v) != nil { // ref exists in rule tree
+			return true
+		}
+	}
+	return isBuiltinRefOrVar(bs, value)
+}
+
 func isInputRef(term *Term) bool {
 	if ref, ok := term.Value.(Ref); ok {
 		if ref.HasPrefix(InputRootRef) {

ast/compile_test.go

@@ -4040,7 +4040,7 @@ func TestCompilerRewriteWithValue(t *testing.T) {
 		{
 			note:    "invalid target",
 			input:   `p { true with foo.q as 1 }`,
-			wantErr: fmt.Errorf("rego_type_error: with keyword target must reference existing input, data, or a built-in function"),
+			wantErr: fmt.Errorf("rego_type_error: with keyword target must reference existing input, data, or a function"),
 		},
 		{
 			note:     "built-in function: replaced by (unknown) var",
@@ -4495,21 +4495,6 @@ func TestRewritePrintCallsWithElseImplicitArgs(t *testing.T) {
 }
 
 func TestCompilerMockFunction(t *testing.T) {
-	c := NewCompiler()
-	c.Modules["test"] = MustParseModule(`
-	package test
-
-	is_allowed(label) {
-	    label == "test_label"
-	}
-
-	p {true with data.test.is_allowed as "blah" }
-	`)
-	compileStages(c, c.rewriteWithModifiers)
-	assertCompilerErrorStrings(t, c, []string{"rego_compile_error: with keyword used on non-built-in function"})
-}
-
-func TestCompilerMockBuiltinFunction(t *testing.T) {
 	tests := []struct {
 		note          string
 		module, extra string
@@ -4645,6 +4630,54 @@ func TestCompilerMockBuiltinFunction(t *testing.T) {
 				p { bar(foo.bar("one")) with bar as mock with foo.bar as mock_mock }
 			`,
 		},
+		{
+			note: "non-built-in function replaced value",
+			module: `package test
+				original(_)
+				p { original(true) with original as 123 }
+			`,
+		},
+		{
+			note: "non-built-in function replaced by another, arity 0",
+			module: `package test
+				original() = 1
+				mock() = 2
+				p { original() with original as mock }
+			`,
+			err: "rego_type_error: undefined function data.test.original", // TODO(sr): file bug -- this doesn't depend on "with" used or not
+		},
+		{
+			note: "non-built-in function replaced by another, arity 1",
+			module: `package test
+				original(_)
+				mock(_)
+				p { original(true) with original as mock }
+			`,
+		},
+		{
+			note: "non-built-in function replaced by built-in",
+			module: `package test
+				original(_)
+				p { original([1]) with original as count }
+			`,
+		},
+		{
+			note: "non-built-in function replaced by another, arity mismatch",
+			module: `package test
+				original(_)
+				mock(_, _)
+				p { original([1]) with original as mock }
+			`,
+			err: "rego_type_error: data.test.original: arity mismatch\n\thave: (any, any)\n\twant: (any)",
+		},
+		{
+			note: "non-built-in function replaced by built-in, arity mismatch",
+			module: `package test
+				original(_)
+				p { original([1]) with original as concat }
+			`,
+			err: "rego_type_error: data.test.original: arity mismatch\n\thave: (string, any<array[string], set[string]>)\n\twant: (any)",
+		},
 	}
 
 	for _, tc := range tests {
@@ -6415,7 +6448,7 @@ func TestQueryCompiler(t *testing.T) {
 			q:        "x = 1 with foo.p as null",
 			pkg:      "",
 			imports:  nil,
-			expected: fmt.Errorf("1 error occurred: 1:12: rego_type_error: with keyword target must reference existing input, data, or a built-in function"),
+			expected: fmt.Errorf("1 error occurred: 1:12: rego_type_error: with keyword target must reference existing input, data, or a function"),
 		},
 		{
 			note:     "rewrite with value",

docs/content/policy-language.md

@@ -1487,7 +1487,7 @@ following syntax:
 ```
 
 The `<target>`s must be references to values in the input document (or the input
-document itself) or data document, or references to built-in functions.
+document itself) or data document, or references to functions (built-in or not).
 
 {{< info >}}
 When applied to the `data` document, the `<target>` must not attempt to
@@ -1516,7 +1516,7 @@ outer := result {
 }
 ```
 
-When `<target>` is a reference to a built-in function, like `http.send`, then
+When `<target>` is a reference to a function, like `http.send`, then
 its `<value>` can be any of the following:
 1. a value: `with http.send as {"body": {"success": true }}`
 2. a reference to another function: `with http.send as mock_http_send`
@@ -1533,10 +1533,10 @@ See the following example:
 package opa.examples
 import future.keywords.in
 
-f(x) = count(x)
+f(x) := count(x)
 
-mock_count(x) = 0 { "x" in x }
-mock_count(x) = count(x) { not "x" in x }
+mock_count(x) := 0 { "x" in x }
+mock_count(x) := count(x) { not "x" in x }
 ```
 
 ```live:with_builtins/1:query:merge_down
@@ -1558,12 +1558,12 @@ Each replacement function evaluation will start a new scope: it's valid to use
 package opa.examples
 import future.keywords.in
 
-f(x) = count(x) {
+f(x) := count(x) {
     rule_using_concat with concat as "foo,bar"
 }
 
-mock_count(x) = 0 { "x" in x }
-mock_count(x) = count(x) { not "x" in x }
+mock_count(x) := 0 { "x" in x }
+mock_count(x) := count(x) { not "x" in x }
 
 rule_using_concat {
     concat(",", input.x) == "foo,bar"

docs/content/policy-testing.md

@@ -233,15 +233,15 @@ opa test --format=json pass_fail_error_test.rego
 
 ## Data and Function Mocking
 
-OPA's `with` keyword can be used to replace the data document or built-in functions by mocks.
+OPA's `with` keyword can be used to replace the data document or called functions with mocks.
 Both base and virtual documents can be replaced.
 
-When replacing built-in functions, the following constraints are in place:
+When replacing functions, built-in or otherwise, the following constraints are in place:
 
 1. Replacing `internal.*` functions, or `rego.metadata.*`, or `eq`; or relations (`walk`) is not allowed.
 2. Replacement and replaced function need to have the same arity.
 3. Replaced functions can call the functions they're replacing, and those calls
-   will call out to the original built-in function, and not cause recursion.
+   will call out to the original function, and not cause recursion.
 
 Below is a simple policy that depends on the data document.
 
@@ -360,7 +360,7 @@ data.authz.test_allow: PASS (458.752µs)
 PASS: 1/1
 ```
 
-In simple cases, a built-in function can also be replaced with a value, as in
+In simple cases, a function can also be replaced with a value, as in
 
 ```live:with_keyword_builtins/tests/value:module:read_only
 test_allow_value {
@@ -374,22 +374,19 @@ test_allow_value {
 Every invocation of the function will then return the replacement value, regardless
 of the function's arguments.
 
-Note that it's also possible to replace one built-in function by another.
-
-
-**User-defined functions** cannot be replaced by the `with` keyword.
-For example, in the below policy the function `cannot_replace` cannot be replaced.
+Note that it's also possible to replace one built-in function by another; or a non-built-in
+function by a built-in function.
 
 **authz.rego**:
 
 ```live:with_keyword_funcs:module:read_only
 package authz
 
-invalid_replace {
-    cannot_replace(input.label)
+replace_rule {
+    replace(input.label)
 }
 
-cannot_replace(label) {
+replace(label) {
     label == "test_label"
 }
 ```
@@ -399,14 +396,16 @@ cannot_replace(label) {
 ```live:with_keyword_funcs/tests:module:read_only
 package authz
 
-test_invalid_replace {
-    invalid_replace with input as {"label": "test_label"} with cannot_replace as true
+test_replace_rule {
+    replace_rule with input.label as "does-not-matter" with replace as true
 }
 ```
 
 ```console
 $ opa test -v authz.rego authz_test.rego
-1 error occurred: authz_test.rego:4: rego_compile_error: with keyword cannot replace rego functions
+data.authz.test_replace_rule: PASS (648.314µs)
+--------------------------------------------------------------------------------
+PASS: 1/1
 ```
 
 

format/testfiles/test_with.rego

@@ -22,4 +22,11 @@ func_replacements {
     count(array.concat(input.x, [])) with input.x as "foo"
     with array.concat as true
     with count as mock_f
+}
+
+original(x) = x+1
+
+more_func_replacements {
+    original(1) with original as mock_f
+    original(1) with original as 1234
 }

format/testfiles/test_with.rego.formatted

@@ -22,3 +22,10 @@ func_replacements {
 		with array.concat as true
 		with count as mock_f
 }
+
+original(x) = x + 1
+
+more_func_replacements {
+	original(1) with original as mock_f
+	original(1) with original as 1234
+}