topdown: Fix evaluator to include with statements on saved negated exprs

The evaluator was incorrectly assuming that if negated exprs were
being saved during partial eval that any with statements would have
been applied. This was not the case.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

topdown/eval.go

@@ -1175,10 +1175,16 @@ func (e *eval) saveCall(declArgsLen int, terms []*ast.Term, iter unifyIterator)
 
 func (e *eval) saveInlinedNegatedExprs(exprs []*ast.Expr, iter unifyIterator) error {
 
-	// This function does not include with statements on the exprs because
-	// they will have already been saved and therefore had their any relevant
-	// with statements set.
+	with := make([]*ast.With, len(e.query[e.index].With))
+
+	for i := range e.query[e.index].With {
+		cpy := e.query[e.index].With[i].Copy()
+		cpy.Value = e.bindings.PlugNamespaced(cpy.Value, e.caller.bindings)
+		with[i] = cpy
+	}
+
 	for _, expr := range exprs {
+		expr.With = with
 		e.saveStack.Push(expr, nil, nil)
 		e.traceSave(expr)
 	}

topdown/topdown_partial_test.go

@@ -2058,7 +2058,34 @@ func TestTopDownPartialEval(t *testing.T) {
 			note:        "negation: save inline negated with",
 			query:       `not input with data.x as 2; data.x = 1`,
 			data:        `{"x": 1}`,
-			wantQueries: []string{"not input"},
+			wantQueries: []string{"not input with data.x as 2"},
+		},
+		{
+			note:  "negation: save negated expr using plugged with value",
+			query: "data.test.p = true",
+			modules: []string{`
+				package test
+
+				p {
+					x = 1
+					not q with input.x as x
+				}
+
+				q {
+					r[input.x]
+				}
+
+				r[1]
+				r[2]
+			`},
+			disableInlining: []string{"data.test.q"},
+			wantQueries:     []string{"not data.partial.test.q with input.x as 1"},
+			wantSupport: []string{`
+				package partial.test
+
+				q { 1 = input.x }
+				q { 2 = input.x }
+			`},
 		},
 		{
 			note:        "negation: save inline negated with (undefined)",