diff --git a/ast/compile.go b/ast/compile.go index 2089e9b580..74c3f80565 100644 --- a/ast/compile.go +++ b/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) diff --git a/ast/compile_test.go b/ast/compile_test.go index 2162ac3dee..c4f249fd4c 100644 --- a/ast/compile_test.go +++ b/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") + } +} diff --git a/compile/compile_test.go b/compile/compile_test.go index db455a6964..e09abd52c3 100644 --- a/compile/compile_test.go +++ b/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 { diff --git a/topdown/eval.go b/topdown/eval.go index b51f5a6572..878303aa37 100644 --- a/topdown/eval.go +++ b/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 diff --git a/topdown/query.go b/topdown/query.go index 2c23389dbb..b22fa9c570 100644 --- a/topdown/query.go +++ b/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) } diff --git a/topdown/uuid_test.go b/topdown/uuid_test.go index c2aa7540c5..b4cbe8d363 100644 --- a/topdown/uuid_test.go +++ b/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 {