partialevaluation: Eliminate type-check failures

Previously partial evaluation could generate rules that could
fail the type-check and therefore fail to load.
However, type-failures simply indicate that the rules can never
succeed and therefore can safely be removed.

This change removes all rules that fail type-checking that are
generated during partial evaluation.

Fixes: #3012
Signed-off-by: Tim Hinrichs <tim@styra.com>

ast/compile.go

@@ -606,6 +606,14 @@ func (c *Compiler) RuleIndex(path Ref) RuleIndex {
 	return r.(RuleIndex)
 }
 
+// PassesTypeCheck determines whether the given body passes type checking
+func (c *Compiler) PassesTypeCheck(body Body) bool {
+	checker := newTypeChecker()
+	env := c.TypeEnv
+	_, errs := checker.CheckBody(env, body)
+	return len(errs) == 0
+}
+
 // ModuleLoader defines the interface that callers can implement to enable lazy
 // loading of modules during compilation.
 type ModuleLoader func(resolved map[string]*Module) (parsed map[string]*Module, err error)

ast/compile_test.go

@@ -4232,3 +4232,23 @@ deny {
 		t.Fatalf("Expected error for unsafe built-in but got %v", err)
 	}
 }
+
+func TestCompilerPassesTypeCheck(t *testing.T) {
+	c := NewCompiler().
+		WithCapabilities(&Capabilities{Builtins: []*Builtin{Split}})
+	// Must compile to initialize type environment after WithCapabilities
+	c.Compile(nil)
+	if c.PassesTypeCheck(MustParseBody(`a = input.a; split(a, ":", x); a0 = x[0]; a0 = null`)) {
+		t.Fatal("Did not successfully detect a type-checking violation")
+	}
+}
+
+func TestCompilerPassesTypeCheckNegative(t *testing.T) {
+	c := NewCompiler().
+		WithCapabilities(&Capabilities{Builtins: []*Builtin{Split, StartsWith}})
+	// Must compile to initialize type environment after WithCapabilities
+	c.Compile(nil)
+	if !c.PassesTypeCheck(MustParseBody(`a = input.a; split(a, ":", x); a0 = x[0]; startswith(a0, "foo", true)`)) {
+		t.Fatal("Incorrectly detected a type-checking violation")
+	}
+}

compile/compile_test.go

@@ -1020,6 +1020,137 @@ func TestOptimizerOutput(t *testing.T) {
 				`,
 			},
 		},
+		{
+			note:        "generate rules with type violations: complete doc",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p {
+						x := split(input.a, ":")
+						f(x[0])
+					}
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+			roots: []string{"test"},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p = __result__ { split(input.a, ":", __local3__1); startswith(__local3__1[0], "foo"); __result__ = true }
+				`,
+				"test.rego": `
+					package test
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+		},
+		{
+			note:        "generate rules with type violations: partial set",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p[msg] {
+						x := split(input.a, ":")
+						f(x[0])
+						msg := "test string"
+					}
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+			roots: []string{"test"},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p["test string"] { split(input.a, ":", __local4__1); startswith(__local4__1[0], "foo") }
+				`,
+				"test.rego": `
+					package test
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+		},
+		{
+			note:        "generate rules with type violations: partial object",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p[k] = value  {
+						x := split(input.a, ":")
+						f(x[0])
+						k := "a"
+						value := 1
+					}
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+			roots: []string{"test"},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p["a"] = 1 { split(input.a, ":", __local5__1); startswith(__local5__1[0], "foo") }
+				`,
+				"test.rego": `
+					package test
+
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+		},
+		{
+			note:        "generate rules with type violations: negation",
+			entrypoints: []string{"data.test.p"},
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p  { not q }
+					q {
+						x := split(input.a, ":")
+						f(x[0])
+					}
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+			},
+			roots: []string{"test"},
+			wantModules: map[string]string{
+				"optimized/test.rego": `
+					package test
+
+					p = __result__ { not data.partial.__not1_0_2__; __result__ = true }
+				`,
+				"test.rego": `
+					package test
+					q = true { assign(x, split(input.a, ":")); f(x[0]) }
+					f(x) { x == null }
+					f(x) { startswith(x, "foo") }
+				`,
+				"optimized/partial.rego": `
+					package partial
+            		__not1_0_2__ = true { split(input.a, ":", __local3__3); startswith(__local3__3[0], "foo") }
+				`,
+			},
+		},
 	}
 
 	for _, tc := range tests {

topdown/eval.go

@@ -503,6 +503,11 @@ func (e *eval) evalNotPartial(iter evalIterator) error {
 	child.eval(func(*eval) error {
 		query := e.saveStack.Peek()
 		plugged := query.Plug(e.caller.bindings)
+		// Skip this rule body if it fails to type-check.
+		// Type-checking failure means the rule body will never succeed.
+		if !e.compiler.PassesTypeCheck(plugged) {
+			return nil
+		}
 		if cp != nil {
 			plugged = applyCopyPropagation(cp, e.instr, plugged)
 		}
@@ -2095,32 +2100,34 @@ func (e evalVirtualPartial) partialEvalSupportRule(iter unifyIterator, rule *ast
 
 		current := e.e.saveStack.PopQuery()
 		plugged := current.Plug(e.e.caller.bindings)
+		// Skip this rule body if it fails to type-check.
+		// Type-checking failure means the rule body will never succeed.
+		if e.e.compiler.PassesTypeCheck(plugged) {
+			var key, value *ast.Term
 
-		var key, value *ast.Term
+			if rule.Head.Key != nil {
+				key = child.bindings.PlugNamespaced(rule.Head.Key, e.e.caller.bindings)
+			}
 
-		if rule.Head.Key != nil {
-			key = child.bindings.PlugNamespaced(rule.Head.Key, e.e.caller.bindings)
+			if rule.Head.Value != nil {
+				value = child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings)
+			}
+
+			head := ast.NewHead(rule.Head.Name, key, value)
+
+			if !e.e.inliningControl.shallow {
+				cp := copypropagation.New(head.Vars()).
+					WithEnsureNonEmptyBody(true).
+					WithCompiler(e.e.compiler)
+				plugged = applyCopyPropagation(cp, e.e.instr, plugged)
+			}
+
+			e.e.saveSupport.Insert(path, &ast.Rule{
+				Head:    head,
+				Body:    plugged,
+				Default: rule.Default,
+			})
 		}
-
-		if rule.Head.Value != nil {
-			value = child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings)
-		}
-
-		head := ast.NewHead(rule.Head.Name, key, value)
-
-		if !e.e.inliningControl.shallow {
-			cp := copypropagation.New(head.Vars()).
-				WithEnsureNonEmptyBody(true).
-				WithCompiler(e.e.compiler)
-			plugged = applyCopyPropagation(cp, e.e.instr, plugged)
-		}
-
-		e.e.saveSupport.Insert(path, &ast.Rule{
-			Head:    head,
-			Body:    plugged,
-			Default: rule.Default,
-		})
-
 		child.traceRedo(rule)
 		e.e.saveStack.PushQuery(current)
 		return nil
@@ -2373,22 +2380,24 @@ func (e evalVirtualComplete) partialEvalSupportRule(iter unifyIterator, rule *as
 
 		current := e.e.saveStack.PopQuery()
 		plugged := current.Plug(e.e.caller.bindings)
+		// Skip this rule body if it fails to type-check.
+		// Type-checking failure means the rule body will never succeed.
+		if e.e.compiler.PassesTypeCheck(plugged) {
+			head := ast.NewHead(rule.Head.Name, nil, child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings))
 
-		head := ast.NewHead(rule.Head.Name, nil, child.bindings.PlugNamespaced(rule.Head.Value, e.e.caller.bindings))
+			if !e.e.inliningControl.shallow {
+				cp := copypropagation.New(head.Vars()).
+					WithEnsureNonEmptyBody(true).
+					WithCompiler(e.e.compiler)
+				plugged = applyCopyPropagation(cp, e.e.instr, plugged)
+			}
 
-		if !e.e.inliningControl.shallow {
-			cp := copypropagation.New(head.Vars()).
-				WithEnsureNonEmptyBody(true).
-				WithCompiler(e.e.compiler)
-			plugged = applyCopyPropagation(cp, e.e.instr, plugged)
+			e.e.saveSupport.Insert(path, &ast.Rule{
+				Head:    head,
+				Body:    plugged,
+				Default: rule.Default,
+			})
 		}
-
-		e.e.saveSupport.Insert(path, &ast.Rule{
-			Head:    head,
-			Body:    plugged,
-			Default: rule.Default,
-		})
-
 		child.traceRedo(rule)
 		e.e.saveStack.PushQuery(current)
 		return nil

topdown/query.go

@@ -350,6 +350,12 @@ func (q *Query) PartialRun(ctx context.Context) (partials []ast.Body, support []
 			body.Append(bindingExprs[i])
 		}
 
+		// Skip this rule body if it fails to type-check.
+		// Type-checking failure means the rule body will never succeed.
+		if !e.compiler.PassesTypeCheck(body) {
+			return nil
+		}
+
 		if !q.shallowInlining {
 			body = applyCopyPropagation(p, e.instr, body)
 		}

topdown/uuid_test.go

@@ -69,8 +69,12 @@ func TestUUIDRFC4122SeedError(t *testing.T) {
 func TestUUIDRFC4122SavingDuringPartialEval(t *testing.T) {
 
 	query := `foo = "x"; uuid.rfc4122(foo,x)`
+	c := ast.NewCompiler().
+		WithCapabilities(&ast.Capabilities{Builtins: []*ast.Builtin{ast.UUIDRFC4122}})
+	// Must compile to initialize type environment after WithCapabilities
+	c.Compile(nil)
 
-	q := NewQuery(ast.MustParseBody(query)).WithSeed(rand.New(rand.NewSource(0))).WithCompiler(ast.NewCompiler())
+	q := NewQuery(ast.MustParseBody(query)).WithSeed(rand.New(rand.NewSource(0))).WithCompiler(c)
 
 	queries, modules, err := q.PartialRun(context.Background())
 	if err != nil {