Fix unset of input/data in REPL

Fixes #259

ast/policy.go

@@ -11,8 +11,6 @@ import (
 	"github.com/open-policy-agent/opa/util"
 )
 
-// TODO(tsandall): rename DefaultRootDocument/Ref accordingly
-
 // DefaultRootDocument is the default root document.
 //
 // All package directives inside source files are implicitly prefixed with the
@@ -22,6 +20,13 @@ var DefaultRootDocument = VarTerm("data")
 // InputRootDocument names the document containing query arguments.
 var InputRootDocument = VarTerm("input")
 
+// RootDocumentNames contains the names of top-level documents that can be
+// referred to in modules and queries.
+var RootDocumentNames = &Set{
+	DefaultRootDocument,
+	InputRootDocument,
+}
+
 // DefaultRootRef is a reference to the root of the default document.
 //
 // All refs to data in the policy engine's storage layer are prefixed with this ref.
@@ -32,12 +37,26 @@ var DefaultRootRef = Ref{DefaultRootDocument}
 // All refs to query arguments are prefixed with this ref.
 var InputRootRef = Ref{InputRootDocument}
 
+// RootDocumentRefs contains the prefixes of top-level documents that all
+// non-local references start with.
+var RootDocumentRefs = &Set{
+	NewTerm(DefaultRootRef),
+	NewTerm(InputRootRef),
+}
+
 // ReservedVars is the set of names that refer to implicitly ground vars.
-var ReservedVars = NewVarSet(DefaultRootDocument.Value.(Var), InputRootDocument.Value.(Var))
+var ReservedVars = NewVarSet(
+	DefaultRootDocument.Value.(Var),
+	InputRootDocument.Value.(Var),
+)
 
 // Wildcard represents the wildcard variable as defined in the language.
 var Wildcard = &Term{Value: Var("_")}
 
+// WildcardPrefix is the special character that all wildcard variables are
+// prefixed with when the statement they are contained in is parsed.
+var WildcardPrefix = "$"
+
 // Keywords contains strings that map to language keywords.
 var Keywords = [...]string{
 	"not",
@@ -61,10 +80,6 @@ func IsKeyword(s string) bool {
 	return false
 }
 
-// WildcardPrefix is the special character that all wildcard variables are
-// prefixed with when the statement they are contained in is parsed.
-var WildcardPrefix = "$"
-
 type (
 	// Statement represents a single statement in a policy module.
 	Statement interface {

repl/repl.go

@@ -390,13 +390,18 @@ func (r *REPL) cmdUnset(args []string) error {
 	}
 
 	term, err := ast.ParseTerm(args[0])
+
 	if err != nil {
 		return newBadArgsErr("argument must identify a rule")
 	}
 
 	v, ok := term.Value.(ast.Var)
+
 	if !ok {
-		return newBadArgsErr("argument must identify a rule")
+		if !ast.RootDocumentRefs.Contains(term) {
+			return newBadArgsErr("argument must identify a rule")
+		}
+		v = term.Value.(ast.Ref)[0].Value.(ast.Var)
 	}
 
 	mod := r.modules[r.currentModuleID]

repl/repl_test.go

@@ -238,11 +238,14 @@ func TestUnset(t *testing.T) {
 	var buffer bytes.Buffer
 	repl := newRepl(store, &buffer)
 
+	var err error
+
 	repl.OneShot(ctx, "magic = 23")
 	repl.OneShot(ctx, "p = 3.14")
 	repl.OneShot(ctx, "unset p")
 
-	err := repl.OneShot(ctx, "p")
+	err = repl.OneShot(ctx, "p")
+
 	if _, ok := err.(ast.Errors); !ok {
 		t.Fatalf("Expected AST error but got: %v", err)
 	}
@@ -283,10 +286,22 @@ func TestUnset(t *testing.T) {
 
 	buffer.Reset()
 	repl.OneShot(ctx, `package data.other`)
-	repl.OneShot(ctx, `unset magic`)
+	err = repl.OneShot(ctx, `unset magic`)
 	if buffer.String() != "warning: no matching rules in current module\n" {
 		t.Fatalf("Expected unset error for bad syntax but got: %v", buffer.String())
 	}
+
+	repl.OneShot(ctx, `input = {}`)
+
+	if err := repl.OneShot(ctx, `unset input`); err != nil {
+		t.Fatalf("Expected unset to succeed for input: %v", err)
+	}
+
+	err = repl.OneShot(ctx, `true = input`)
+
+	if !strings.Contains(err.Error(), "undefined") {
+		t.Fatalf("Expected undefined error but got: %v", err)
+	}
 }
 
 func TestOneShotEmptyBufferOneExpr(t *testing.T) {