planner: fix plan for 'every' (open-policy-agent#4346)

Before, we had been planing `every x in xs { BODY }` as, roughly,

    NOT
      SCAN xs
        NOT
         BODY

This poses a problem: scans never fail -- they iterate _their_
bodies, and break out of their block when they are done iterating.
As far as the NOT is concerned, it always looks like the SCAN was
successful.

We had worked around that for the "outer" SCAN by breaking out of
the NOT's block at the end. However, the same trick can't be used
when BODY contains another SCAN, since we can't tell beforehand how
deeply we're nested, and doesn't know how far to break out again.

So in this commit, we replace the NOTs by some custom "condition
var" construct that works similar to how NOT gets compiled, but
sets the "inner" condition variable when the BODY's plan succeeds.
Likewise, the "outer" condition variable is tied to that result.

We're practically using the condition variables to encode

    every x, y in xs { p(x,y) }
    ~> p(x1, y1) AND p(x2, y2) AND ... AND p(xn, yn)
    ~> NOT (NOT p(x1, y1) OR NOT p(x2, y2) OR ... OR NOT p(xn, yn))

The conditon variables are now implemented in the IR. (When using
ir.NotStmt, they are an artifact of compiling the IR to wasm.) This
is achieved by resetting a new local, and assigning it a dummy value
(true) to signal the condition was met.

----

This path was taken because I could not come up with a way to deal
with the "inner SCAN" situation just by using NOT blocks. The "outer"
NOT/SCAN could perhaps be salvaged, but I found it easier to reason
about one mechanism applied twice than to mix and match.

Another abandoned path was using ir.ReturnLocalStmt in the query's
plan iterator: while it worked well for simple cases, the problems
considered were that we might do the wrong thing in nested sitations,
like comprehensions. Also, it's not something we'd done in any other
place.

Signed-off-by: Stephan Renatus <[email protected]>

internal/planner/planner.go

@@ -719,51 +719,68 @@ func (p *Planner) planExprTerm(e *ast.Expr, iter planiter) error {
 func (p *Planner) planExprEvery(e *ast.Expr, iter planiter) error {
 	every := e.Terms.(*ast.Every)
 
-	not0 := &ir.NotStmt{
-		Block: &ir.Block{},
-	}
-	prev := p.curr
-	p.curr = not0.Block
-
-	// NOT { # not0
-	//   c0 = 1
-	//   SCAN domain
-	//     NOT { # not1
-	//       c1 = 1
-	//       QUERY
-	//       c1 = 0
-	//     } // check c1, break out of scan
-	//   c0 = 0
-	// } // check c0, break to undefined
+	cond0 := p.newLocal() // outer not
+	cond1 := p.newLocal() // inner not
+
+	// We're using condition variables together with IsDefinedStmt to encode
+	// this:
+	// every x, y in xs { p(x,y) }
+	// ~> p(x1, y1) AND p(x2, y2) AND ... AND p(xn, yn)
+	// ~> NOT (NOT p(x1, y1) OR NOT p(x2, y2) OR ... OR NOT p(xn, yn))
+	//
+	// cond1 is initialized to 0, and set to TRUE if p(xi, yi) succeeds for
+	// a binding of (xi, yi). We then use IsUndefined to check that this has NOT
+	// happened (NOT p(xi, yi)).
+	// cond0 is initialized to 0, and set to TRUE if cond1 happens to not
+	// be set: it's encoding the NOT ( ... OR ... OR ... ) part of this.
+
+	p.appendStmt(&ir.ResetLocalStmt{
+		Target: cond0,
+	})
 
 	err := p.planTerm(every.Domain, func() error {
 		return p.planScan(every.Key, func(ir.Local) error {
-			scan := p.curr
+			p.appendStmt(&ir.ResetLocalStmt{
+				Target: cond1,
+			})
+			nested := &ir.BlockStmt{Blocks: []*ir.Block{{}}}
 
-			not1 := &ir.NotStmt{
-				Block: &ir.Block{},
-			}
-			p.curr = not1.Block
+			prev := p.curr
+			p.curr = nested.Blocks[0]
 
 			lval := p.ltarget
 			err := p.planUnifyLocal(lval, every.Value, func() error {
 				return p.planQuery(every.Body, 0, func() error {
-					p.appendStmtToBlock(not1, scan)
+					p.appendStmt(&ir.AssignVarStmt{
+						Source: op(ir.Bool(true)),
+						Target: cond1,
+					})
 					return nil
 				})
 			})
 			if err != nil {
 				return err
 			}
-			p.appendStmtToBlock(&ir.BreakStmt{Index: 2}, scan)
+
+			p.curr = prev
+			p.appendStmt(nested)
+			p.appendStmt(&ir.IsUndefinedStmt{
+				Source: cond1,
+			})
+			p.appendStmt(&ir.AssignVarStmt{
+				Source: op(ir.Bool(true)),
+				Target: cond0,
+			})
 			return nil
 		})
 	})
 	if err != nil {
 		return err
 	}
 
-	p.appendStmtToBlock(not0, prev)
+	p.appendStmt(&ir.IsUndefinedStmt{
+		Source: cond0,
+	})
 	return iter()
 }
 

internal/wasm/sdk/test/e2e/exceptions.yaml

@@ -1,3 +1,2 @@
 # Exception Format is <test name>: <reason>
 "functions/default": "not supported in topdown, https://github.com/open-policy-agent/opa/issues/2445"
-"every/example every/some": "WIP"

test/cases/testdata/every/every.yaml

@@ -113,4 +113,19 @@ cases:
   note: 'every/with: body'
   query: data.test.p = x
   want_result:
-  - x: true
+  - x: true
+- data:
+  modules:
+  - |
+    package test
+    import future.keywords.every
+
+    p[v] {
+      every v in [1, 2] { v < 3 }
+      v := 10
+      v > 3
+    }
+  note: 'every/followed by another query'
+  query: data.test.p = x
+  want_result:
+  - x: [10]

test/cases/testdata/every/textbook.yaml

@@ -100,9 +100,32 @@ cases:
         startswith(c.image, repo)
       }
     }
-  note: every/example every/some
+  note: every/example every/some, fail
   query: data.test.p = x
   want_result: []
+- data:
+  input:
+    containers:
+    - image: hooli.com/bitcoin-miner
+    - image: acmecorp.net/webapp
+    - image: hooli.com/nginx
+  modules:
+  - |
+    package test
+    import future.keywords.every
+
+    allowed_repos := {"hooli.com/", "acmecorp.net/"}
+
+    p {
+      every c in input.containers {
+        some repo in allowed_repos
+        startswith(c.image, repo)
+      }
+    }
+  note: every/example every/some, success
+  query: data.test.p = x
+  want_result:
+  - x: true
 - data:
   input:
     servers: