Add simple rule indexing implementation

ast/compile.go

@@ -72,6 +72,7 @@ type Compiler struct {
 
 	generatedVars map[*Module]VarSet
 	moduleLoader  ModuleLoader
+	ruleIndices   *util.HashMap
 	stages        []func()
 }
 
@@ -161,6 +162,12 @@ func NewCompiler() *Compiler {
 		Modules:       map[string]*Module{},
 		TypeEnv:       NewTypeEnv(),
 		generatedVars: map[*Module]VarSet{},
+		ruleIndices: util.NewHashMap(func(a, b util.T) bool {
+			r1, r2 := a.(Ref), b.(Ref)
+			return r1.Equal(r2)
+		}, func(x util.T) int {
+			return x.(Ref).Hash()
+		}),
 	}
 
 	c.ModuleTree = NewModuleTree(nil)
@@ -178,6 +185,7 @@ func NewCompiler() *Compiler {
 		c.checkSafetyRuleBodies,
 		c.checkRecursion,
 		c.checkTypes,
+		c.buildRuleIndices,
 	}
 
 	return c
@@ -339,6 +347,18 @@ func (c *Compiler) GetRules(ref Ref) (rules []*Rule) {
 	return rules
 }
 
+// RuleIndex returns a RuleIndex built for the rule set referred to by path.
+// The path must refer to the rule set exactly, i.e., given a rule set at path
+// data.a.b.c.p, refs data.a.b.c.p.x and data.a.b.c would not return a
+// RuleIndex built for the rule.
+func (c *Compiler) RuleIndex(path Ref) RuleIndex {
+	r, ok := c.ruleIndices.Get(path)
+	if !ok {
+		return nil
+	}
+	return r.(RuleIndex)
+}
+
 // 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)
@@ -356,6 +376,23 @@ func (c *Compiler) WithModuleLoader(f ModuleLoader) *Compiler {
 	return c
 }
 
+// buildRuleIndices constructs indices for rules.
+func (c *Compiler) buildRuleIndices() {
+
+	c.RuleTree.DepthFirst(func(node *RuleTreeNode) bool {
+		if len(node.Rules) > 1 {
+			index := newBaseDocEqIndex(func(ref Ref) bool {
+				return len(c.GetRules(ref.GroundPrefix())) > 0
+			})
+			if index.Build(node.Rules) {
+				c.ruleIndices.Put(node.Rules[0].Path(), index)
+			}
+		}
+		return false
+	})
+
+}
+
 // checkRecursion ensures that there are no recursive rule definitions, i.e., there are
 // no cycles in the RuleGraph.
 func (c *Compiler) checkRecursion() {
@@ -690,7 +727,6 @@ func (qc *queryCompiler) Compile(query Body) (Body, error) {
 		qc.resolveRefs,
 		qc.checkWithModifiers,
 		qc.checkSafety,
-		qc.checkInput,
 		qc.checkTypes,
 	}
 
@@ -742,73 +778,6 @@ func (qc *queryCompiler) checkSafety(_ *QueryContext, body Body) (Body, error) {
 	return reordered, nil
 }
 
-func (qc *queryCompiler) checkInput(qctx *QueryContext, body Body) (Body, error) {
-	return body, qc.checkInputRec(qctx.InputDefined(), body)
-}
-
-func (qc *queryCompiler) checkInputRec(definedPrev bool, body Body) error {
-
-	// Perform DFS for conflicting or missing input document.
-	for _, expr := range body {
-
-		definedCurr := definesInput(expr)
-
-		if definedPrev && definedCurr {
-			return NewError(InputErr, expr.Location, "input document conflict")
-		} else if !definedCurr && !definedPrev && referencesInput(expr) {
-			return NewError(InputErr, expr.Location, "input document not defined")
-		}
-
-		var err error
-
-		// Check closures contained in this expression.
-		vis := NewGenericVisitor(func(x interface{}) bool {
-			if err != nil {
-				return true
-			}
-			switch x := x.(type) {
-			case *ArrayComprehension:
-				if err = qc.checkInputRec(definedPrev || definedCurr, x.Body); err != nil {
-					return true
-				}
-			}
-			return false
-		})
-
-		Walk(vis, expr)
-
-		if err != nil {
-			return err
-		}
-
-		// Check rule bodies referred to by this expression.
-		vis = NewGenericVisitor(func(x interface{}) bool {
-			if err != nil {
-				return true
-			}
-			switch x := x.(type) {
-			case Ref:
-				if x.HasPrefix(DefaultRootRef) {
-					for _, rule := range qc.compiler.GetRules(x.GroundPrefix()) {
-						if err = qc.checkInputRec(definedPrev || definedCurr, rule.Body); err != nil {
-							return true
-						}
-					}
-				}
-			}
-			return false
-		})
-
-		Walk(vis, expr)
-
-		if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
 // referencesInput returns true if expr refers to the input document. This
 // function will not visit closures.
 func referencesInput(expr *Expr) bool {

ast/compile_test.go

@@ -1081,12 +1081,8 @@ func TestQueryCompiler(t *testing.T) {
 		{"unsafe vars", "z", "", nil, "", fmt.Errorf("1 error occurred: 1:1: rego_unsafe_var_error: var z is unsafe")},
 		{"safe vars", `data; abc`, `package ex`, []string{"import input.xyz as abc"}, `{}`, `data; input.xyz`},
 		{"reorder", `x != 1; x = 0`, "", nil, "", `x = 0; x != 1`},
-		// {"bad builtin", "deadbeef(1,2,3)", "", nil, "", fmt.Errorf("1 error occurred: 1:1: rego_type_error: undefined built-in function")},
 		{"bad with target", "x = 1 with data.p as null", "", nil, "", fmt.Errorf("1 error occurred: 1:7: rego_type_error: with keyword target must be input")},
-		// wrapping refs in extra terms to cover error handling
-		{"undefined input", `[[true | [data.a.b.d.t, true]], true]`, "", nil, "", fmt.Errorf("5:12: rego_input_error: input document not defined")},
-		{"conflicting input", `[true | data.a.b.d.t with input as 1]`, "", nil, "2", fmt.Errorf("1:9: rego_input_error: input document conflict")},
-		{"conflicting input-2", `sum([1 | data.a.b.d.t with input as 2], x) with input as 3`, "", nil, "", fmt.Errorf("1:10: rego_input_error: input document conflict")},
+		{"check types", "x = data.a.b.c.z; y = null; x = y", "", nil, "", fmt.Errorf("match error\n\tleft  : number\n\tright : null")},
 	}
 
 	for _, tc := range tests {
@@ -1259,7 +1255,7 @@ func runQueryCompilerTest(t *testing.T, note, q, pkg string, imports []string, i
 			if err == nil {
 				t.Fatalf("Expected error from %v but got: %v", query, result)
 			}
-			if err.Error() != expected.Error() {
+			if !strings.Contains(err.Error(), expected.Error()) {
 				t.Fatalf("Expected error %v but got: %v", expected, err)
 			}
 		}