repl: Allow interactive ref head rule definitions

v1/ast/parser_ext.go

@@ -346,10 +346,11 @@ func ParsePartialObjectDocRuleFromEqExpr(module *Module, lhs, rhs *Term) (*Rule,
 	body := NewBody(NewExpr(BooleanTerm(true).SetLocation(rhs.Location)).SetLocation(rhs.Location))
 
 	rule := &Rule{
-		Location: rhs.Location,
-		Head:     head,
-		Body:     body,
-		Module:   module,
+		Location:      rhs.Location,
+		Head:          head,
+		Body:          body,
+		Module:        module,
+		generatedBody: true,
 	}
 
 	return rule, nil

v1/ast/parser_test.go

@@ -5773,6 +5773,36 @@ data = {"bar": 2}`
 	}
 }
 
+func TestRuleFromExprGeneratedBody(t *testing.T) {
+	// Rules parsed from a single expression have a generated body, so no `if`
+	// keyword is required of them under rego-v1. This is what allows rules to
+	// be defined interactively in the REPL.
+	tests := []string{
+		`x := 1`,
+		`x = 1`,
+		`a[0] := 1`,
+		`a["foo"] = "bar"`,
+		`p.q.r := 1`,
+	}
+
+	mod := MustParseModule("package test")
+
+	for _, tc := range tests {
+		t.Run(tc, func(t *testing.T) {
+			rule, err := ParseRuleFromExpr(mod, MustParseBody(tc)[0])
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+			if !rule.generatedBody {
+				t.Fatal("Expected rule to have generated body")
+			}
+			if errs := CheckRegoV1(rule); len(errs) > 0 {
+				t.Fatalf("Unexpected errors: %v", errs)
+			}
+		})
+	}
+}
+
 func TestWildcards(t *testing.T) {
 
 	assertParseOneTerm(t, "ref", "a.b[_].c[_]", RefTerm(

v1/repl/repl.go

@@ -731,25 +731,20 @@ func (r *REPL) cmdUnknown(s []string) error {
 
 func (r *REPL) cmdUnset(ctx context.Context, args []string) error {
 	if len(args) != 1 {
-		return newBadArgsErr("unset <var>: expects exactly one argument")
+		return newBadArgsErr("unset <ref>: expects exactly one argument")
 	}
 
 	term, err := ast.ParseTerm(args[0])
 	if err != nil {
 		return newBadArgsErr("argument must identify a rule")
 	}
 
-	v, ok := term.Value.(ast.Var)
-
+	ref, ok := ruleRefFromTerm(term)
 	if !ok {
-		ref, ok := term.Value.(ast.Ref)
-		if !ok || !ast.RootDocumentNames.Contains(ref[0]) {
-			return newBadArgsErr("arguments must identify a rule")
-		}
-		v = ref[0].Value.(ast.Var)
+		return newBadArgsErr("arguments must identify a rule")
 	}
 
-	unset, err := r.unsetRule(ctx, v)
+	unset, err := r.unsetRule(ctx, ref)
 	if err != nil {
 		return err
 	} else if !unset {
@@ -779,17 +774,17 @@ func (r *REPL) cmdUnsetPackage(ctx context.Context, args []string) error {
 	return nil
 }
 
-func (r *REPL) unsetRule(ctx context.Context, name ast.Var) (bool, error) {
+func (r *REPL) unsetRule(ctx context.Context, ref ast.Ref) (bool, error) {
 	if r.currentModuleID == "" {
 		return false, nil
 	}
 
 	mod := r.modules[r.currentModuleID]
 	rules := []*ast.Rule{}
 
-	for _, r := range mod.Rules {
-		if r.Head.Name != name {
-			rules = append(rules, r)
+	for _, rule := range mod.Rules {
+		if !refsOverlap(rule.Head.Ref(), ref) {
+			rules = append(rules, rule)
 		}
 	}
 
@@ -807,6 +802,45 @@ func (r *REPL) unsetRule(ctx context.Context, name ast.Var) (bool, error) {
 	return true, nil
 }
 
+// ruleRefFromTerm returns the rule head ref identified by term: the terms "p",
+// "p.q.r" and `p["q"]` identify the rule head refs p, p.q.r and p.q. Refs
+// rooted at a root document are only identified by their first term, e.g.
+// "input" identifies the rule defined by `input = {...}`.
+func ruleRefFromTerm(term *ast.Term) (ast.Ref, bool) {
+	switch v := term.Value.(type) {
+	case ast.Var:
+		return ast.Ref{term}, true
+	case ast.Ref:
+		if _, ok := v[0].Value.(ast.Var); !ok {
+			return nil, false
+		}
+		if ast.RootDocumentNames.Contains(v[0]) {
+			return v[:1], true
+		}
+		// Ref.IsGround ignores the leading var, so this only rejects refs with
+		// non-ground terms after the head, e.g. "a[x]".
+		if !v.IsGround() {
+			return nil, false
+		}
+		return v, true
+	}
+	return nil, false
+}
+
+// refsOverlap returns true if one of the refs is a prefix of, or equal to, the
+// other. Rule head refs that overlap in this way define (parts of) the same
+// document: "user" overlaps "user.role", so unsetting "user" removes
+// `user.role := "admin"`, while defining `user.role := "admin"` only replaces
+// previous definitions of user.role, leaving user.email in place.
+func refsOverlap(a, b ast.Ref) bool {
+	for i := range min(len(a), len(b)) {
+		if a[i].Value.Compare(b[i].Value) != 0 {
+			return false
+		}
+	}
+	return true
+}
+
 func (r *REPL) unsetPackage(_ context.Context, pkg *ast.Package) (bool, error) {
 	path := pkg.Path.String()
 	_, ok := r.modules[path]
@@ -895,7 +929,7 @@ func (r *REPL) compileRule(ctx context.Context, rule *ast.Rule) error {
 
 	if rule.Head.Assign {
 		var err error
-		unset, err = r.unsetRule(ctx, rule.Head.Name)
+		unset, err = r.unsetRule(ctx, rule.Head.Ref())
 		if err != nil {
 			return err
 		}
@@ -1295,6 +1329,12 @@ func (r *REPL) evalPackage(p *ast.Package) error {
 //	 > a
 //	 1
 //
+// The left hand side may be a ref, defining part of a document:
+//
+//		> user["role"] := "admin"
+//	 > user
+//	 {"role": "admin"}
+//
 // If the expression is a = statement, then an additional check on the left
 // hand side occurs. For example:
 //
@@ -1324,8 +1364,11 @@ func (r *REPL) interpretAsRule(ctx context.Context, compiler *ast.Compiler, body
 	if rule == nil || err != nil {
 		return false, err
 	}
-	// TODO(sr): support interactive ref head rule definitions
-	if len(rule.Head.Ref()) > 1 {
+
+	// Statements about a root document are queries, never rule definitions:
+	// "data.foo.bar = 1" asks if data.foo.bar is 1. The single-term case is
+	// excluded so that `input = {...}` keeps defining a rule.
+	if ref := rule.Head.Ref(); len(ref) > 1 && ast.RootDocumentNames.Contains(ref[0]) {
 		return false, nil
 	}
 
@@ -1600,7 +1643,7 @@ var extra = [...]commandDesc{
 var builtin = [...]commandDesc{
 	{"show", []string{""}, "show active module definition"},
 	{"show debug", []string{""}, "show REPL settings"},
-	{"unset", []string{"<var>"}, "unset rules in currently active module"},
+	{"unset", []string{"<ref>"}, "unset rules in currently active module"},
 	{"unset-package", []string{"<var>"}, "unset packages in currently active module"},
 	{"json", []string{}, "set output format to JSON"},
 	{"pretty", []string{}, "set output format to pretty"},
@@ -1662,35 +1705,64 @@ func dumpStorage(ctx context.Context, store storage.Store, txn storage.Transacti
 	return e.Encode(data)
 }
 
+// isGlobalInModule returns true if term refers to an import, or to a document
+// that the module already defines rules for. Statements about such documents
+// are evaluated as queries rather than interpreted as rule definitions.
 func isGlobalInModule(compiler *ast.Compiler, module *ast.Module, term *ast.Term) bool {
 
-	var name ast.Var
+	var ref ast.Ref
 
-	if ast.RootDocumentRefs.Contains(term) {
-		name = term.Value.(ast.Ref)[0].Value.(ast.Var)
-	} else if v, ok := term.Value.(ast.Var); ok {
-		name = v
-	} else {
+	switch v := term.Value.(type) {
+	case ast.Var:
+		ref = ast.Ref{term}
+	case ast.Ref:
+		if _, ok := v[0].Value.(ast.Var); !ok {
+			return false
+		}
+		ref = v
+	default:
 		return false
 	}
 
+	name := ref[0].Value.(ast.Var)
+
 	for _, imp := range module.Imports {
 		if imp.Name().Compare(name) == 0 {
 			return true
 		}
 	}
 
-	path := module.Package.Path.Copy().Append(ast.StringTerm(string(name)))
+	// Only the ground prefix of the ref can be looked up in the rule tree, e.g.
+	// "a[i]" is looked up as "a". The rule tree keys the leading var of the ref
+	// as a string.
+	prefix := ref.GroundPrefix()
+	path := make(ast.Ref, 0, len(prefix))
+	path = append(path, ast.StringTerm(string(name)))
+	path = append(path, prefix[1:]...)
+
 	node := compiler.RuleTree
 
+	for _, elem := range module.Package.Path {
+		node = node.Child(elem.Value)
+		if node == nil {
+			return false
+		}
+	}
+
 	for _, elem := range path {
 		node = node.Child(elem.Value)
 		if node == nil {
 			return false
 		}
+		if len(node.Values) > 0 {
+			return true
+		}
 	}
 
-	return len(node.Values) > 0
+	// The whole prefix resolved to a node without rules of its own: the ref
+	// refers to an existing document only if the dropped suffix can be
+	// satisfied by rules below that node, e.g. "a[i]" when a[0] is defined.
+	return len(prefix) < len(ref)
 }
 
 func printHelp(output io.Writer, initPrompt string, report [][2]string) {