Fix == rewriting on embedded terms

In a660d8f we added a rewriting stage that converts == to unification.
This simplifies evaluation and allows the rule index to be used,
however, the rewriting was incorrectly applied to == expressions that
are intended to yield a true/false value (and not undefined).

Fixes open-policy-agent#995

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

ast/compile.go

@@ -2113,7 +2113,7 @@ func rewriteEquals(x interface{}) {
 	WalkExprs(x, func(x *Expr) bool {
 		if x.IsCall() {
 			operator := x.Operator()
-			if operator.Equal(doubleEq) {
+			if operator.Equal(doubleEq) && len(x.Operands()) == 2 {
 				x.SetOperator(NewTerm(unifyOp))
 			}
 		}

ast/compile_test.go

@@ -1320,6 +1320,16 @@ func TestCompilerRewriteDoubleEq(t *testing.T) {
 			input: "p { count([1,2]) == 2 }",
 			exp:   `count([1,2], __local0__); __local0__ = 2`,
 		},
+		{
+			note:  "embedded",
+			input: "p { x = 1; y = [x == 0] }",
+			exp:   `x = 1; equal(x, 0, __local0__); y = [__local0__]`,
+		},
+		{
+			note:  "embedded in call",
+			input: `p { x = 0; neq(true, x == 1) }`,
+			exp:   `x = 0; equal(x, 1, __local0__); neq(true, __local0__)`,
+		},
 	}
 	for _, tc := range tests {
 		test.Subtest(t, tc.note, func(t *testing.T) {