Update compiler to check for partially replaced virtual docs

Previously the compiler only checked whether the with keyword was
applied to non-input/data documents or functions. This mean that the
with keyword target could be used to partially replace virtual
documents. Unfortunately the evalutor does not support this kind of
replacement. For example, in v0.10.6:

> p = {"a": 1, "b": 2}
Rule 'p' defined in package repl. Type 'show' to see rules.
> p = x with p.a as 100
+---------------+
|       x       |
+---------------+
| {"a":1,"b":2} |
+---------------+

These changes update the compiler to walk the rule tree and check if
the with target would replace a virtual document and generates an
error in that case.

Allowing the with keyword to partially replace virtual documents comes
with its own set of problems. For example, suppose we defined a
virtual document 'p' and queried for 'p' using the with modifier
(e.g., 'p with p.foo as 2'):

p[x] = 1 {
     expr1
     x = "foo"
}

What should the value of 'p' be if an error occurs during evaluation
of 'expr1'? It's arguable that the error should be masked because the
value generated by the rule would be replaced. Unfortunately, the
evaluator may not know the value of 'x' when the error occurs so it
doesn't have enough information to decide.

Signed-off-by: Torin Sandall <[email protected]>

ast/compile.go

@@ -930,7 +930,7 @@ func (c *Compiler) rewriteWithModifiers() {
 			if !ok {
 				return x, nil
 			}
-			body, err := rewriteWithModifiersInBody(f, body, c.GetRules)
+			body, err := rewriteWithModifiersInBody(c, f, body)
 			if err != nil {
 				c.err(err)
 			}
@@ -1113,7 +1113,7 @@ func (qc *queryCompiler) checkTypes(qctx *QueryContext, body Body) (Body, error)
 
 func (qc *queryCompiler) rewriteWithModifiers(qctx *QueryContext, body Body) (Body, error) {
 	f := newEqualityFactory(newLocalVarGenerator(body))
-	body, err := rewriteWithModifiersInBody(f, body, qc.compiler.GetRules)
+	body, err := rewriteWithModifiersInBody(qc.compiler, f, body)
 	if err != nil {
 		return nil, Errors{err}
 	}
@@ -2692,10 +2692,10 @@ func rewriteDeclaredVar(g *localVarGenerator, stack *localDeclaredVars, v Var) (
 // rewriteWithModifiersInBody will rewrite the body so that with modifiers do
 // not contain terms that require evaluation as values. If this function
 // encounters an invalid with modifier target then it will raise an error.
-func rewriteWithModifiersInBody(f *equalityFactory, body Body, getRules func(ref Ref) []*Rule) (Body, *Error) {
+func rewriteWithModifiersInBody(c *Compiler, f *equalityFactory, body Body) (Body, *Error) {
 	var result Body
 	for i := range body {
-		exprs, err := rewriteWithModifier(f, body[i], getRules)
+		exprs, err := rewriteWithModifier(c, f, body[i])
 		if err != nil {
 			return nil, err
 		}
@@ -2710,11 +2710,11 @@ func rewriteWithModifiersInBody(f *equalityFactory, body Body, getRules func(ref
 	return result, nil
 }
 
-func rewriteWithModifier(f *equalityFactory, expr *Expr, getRules func(ref Ref) []*Rule) ([]*Expr, *Error) {
+func rewriteWithModifier(c *Compiler, f *equalityFactory, expr *Expr) ([]*Expr, *Error) {
 
 	var result []*Expr
 	for i := range expr.With {
-		err := validateTarget(expr.With[i].Target, getRules)
+		err := validateTarget(c, expr.With[i].Target)
 		if err != nil {
 			return nil, err
 		}
@@ -2735,20 +2735,34 @@ func rewriteWithModifier(f *equalityFactory, expr *Expr, getRules func(ref Ref)
 	return result, nil
 }
 
-func validateTarget(term *Term, getRules func(ref Ref) []*Rule) *Error {
+func validateTarget(c *Compiler, term *Term) *Error {
 	if !isInputRef(term) && !isDataRef(term) {
 		return NewError(TypeErr, term.Location, "with keyword target must start with %v or %v", InputRootDocument, DefaultRootDocument)
 	}
 
 	if isDataRef(term) {
-		if ref, ok := term.Value.(Ref); ok {
-			rules := getRules(ref)
-			for _, rule := range rules {
-				if len(rule.Head.Args) > 0 {
-					return NewError(CompileErr, term.Location, "with keyword cannot replace rules with arguments")
+		ref := term.Value.(Ref)
+		node := c.RuleTree
+		for i := 0; i < len(ref)-1; i++ {
+			child := node.Child(ref[i].Value)
+			if child == nil {
+				break
+			} else if len(child.Values) > 0 {
+				return NewError(CompileErr, term.Loc(), "with keyword cannot partially replace virtual document(s)")
+			}
+			node = child
+		}
+
+		if node != nil {
+			if child := node.Child(ref[len(ref)-1].Value); child != nil {
+				for _, value := range child.Values {
+					if len(value.(*Rule).Head.Args) > 0 {
+						return NewError(CompileErr, term.Loc(), "with keyword cannot replace functions")
+					}
 				}
 			}
 		}
+
 	}
 	return nil
 }

ast/compile_test.go

@@ -1591,9 +1591,20 @@ func TestCompilerMockFunction(t *testing.T) {
 	p {true with data.test.is_allowed as "blah" }
 	`)
 	compileStages(c, c.rewriteWithModifiers)
+	assertCompilerErrorStrings(t, c, []string{"rego_compile_error: with keyword cannot replace functions"})
+}
+
+func TestCompilerMockVirtualDocumentPartially(t *testing.T) {
+	c := NewCompiler()
+
+	c.Modules["test"] = MustParseModule(`
+	package test
+	p = {"a": 1}
+	q = x { p = x with p.a as 2 }
+	`)
 
-	expectedError := fmt.Errorf("rego_compile_error: with keyword cannot replace rules with arguments")
-	assertCompilerErrorStrings(t, c, []string{expectedError.Error()})
+	compileStages(c, c.rewriteWithModifiers)
+	assertCompilerErrorStrings(t, c, []string{"rego_compile_error: with keyword cannot partially replace virtual document(s)"})
 }
 
 func TestCompilerSetGraph(t *testing.T) {