diff --git a/dependencies/deps.go b/dependencies/deps.go
index 17b4c425fa..a132d68ce6 100644
--- a/dependencies/deps.go
+++ b/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
diff --git a/dependencies/deps_test.go b/dependencies/deps_test.go
index cc881a4a8e..26c77f15ed 100644
--- a/dependencies/deps_test.go
+++ b/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 {