Update analysis to mark input as base document

Refs against the input root document were not being considered "base
document" refs. These changes update the analysis to treat all of the
root documents (i.e., input and data) as "base documents" by default.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

dependencies/deps.go

@@ -286,7 +286,7 @@ func extractEq(exprs ast.Body) (vars map[ast.Var][]ast.Ref, others []*ast.Expr)
 
 func expandRef(r ast.Ref, vars map[ast.Var]*util.HashMap) []ast.Ref {
 	head, rest := r[0], r[1:]
-	if ast.DefaultRootDocument.Equal(head) {
+	if ast.RootDocumentNames.Contains(head) {
 		return []ast.Ref{r}
 	}
 
@@ -317,7 +317,7 @@ func joinVarRefs(vars map[ast.Var][]ast.Ref) map[ast.Var]*util.HashMap {
 		for v, rs := range vars {
 			for _, r := range rs {
 				head, rest := r[0], r[1:]
-				if ast.DefaultRootDocument.Equal(head) {
+				if ast.RootDocumentNames.Contains(head) {
 					if _, ok := joined[v].Get(r); !ok {
 						joined[v].Put(r, struct{}{})
 						done = false

dependencies/deps_test.go

@@ -440,6 +440,53 @@ func TestBaseAndVirtual(t *testing.T) {
 	}
 }
 
+func TestBase(t *testing.T) {
+	modules := map[string]*ast.Module{
+		"test": ast.MustParseModule(`
+			package test
+
+			p {
+				input = x
+				x = y
+				y.z = "foo"
+			}
+
+			q {
+				input.a = "bar"
+			}
+		`),
+	}
+
+	compiler := ast.NewCompiler()
+	compiler.Compile(modules)
+	if compiler.Failed() {
+		t.Fatal(compiler.Errors)
+	}
+
+	body := ast.MustParseBody("data.test.p")
+
+	refs, err := Base(compiler, body)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// TODO(tsandall): dependency analysis should be able to identify that full
+	// extent of input is not required here (only input.z and input.a are
+	// needed)
+	exp := []ast.Ref{ast.MustParseRef("input")}
+
+	if len(exp) != 1 {
+		t.Fatalf("Expected %v but got %v", exp, refs)
+	}
+
+	for i := range refs {
+		if refs[i].Compare(exp[i]) != 0 {
+			t.Fatalf("Expected %v but got: %v", exp, refs)
+		}
+	}
+
+}
+
 func runDeps(t *testing.T, x interface{}) (min, full []ast.Ref) {
 	min, err := Minimal(x)
 	if err != nil {