cue/ast: make package clause a declaration

This makes handling comments during generation
and analysis considerably easier.

Also fixes comments printing of import specs.

Change-Id: I0a77405524b40eff47e1e1e0dd3183ed3563b263
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/2865
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
diff --git a/cmd/cue/cmd/add.go b/cmd/cue/cmd/add.go
index 57b1fbc..68f101b 100644
--- a/cmd/cue/cmd/add.go
+++ b/cmd/cue/cmd/add.go
@@ -30,6 +30,7 @@
 	"cuelang.org/go/cue/format"
 	"cuelang.org/go/cue/load"
 	"cuelang.org/go/cue/parser"
+	"cuelang.org/go/internal"
 	"github.com/spf13/cobra"
 )
 
@@ -252,9 +253,9 @@
 		if err != nil {
 			return nil, err
 		}
-		if f.Name != nil {
-			if pkg := flagPackage.String(cmd); pkg != "" && f.Name.Name != pkg {
-				return nil, fmt.Errorf("package mismatch (%s vs %s) for file %s", f.Name.Name, pkg, file)
+		if _, pkgName, _ := internal.PackageInfo(f); pkgName != "" {
+			if pkg := flagPackage.String(cmd); pkg != "" && pkgName != pkg {
+				return nil, fmt.Errorf("package mismatch (%s vs %s) for file %s", pkgName, pkg, file)
 			}
 			todo.build = getBuild(dir)
 		} else {
diff --git a/cmd/cue/cmd/get_go.go b/cmd/cue/cmd/get_go.go
index 34523d8..7a05f68 100644
--- a/cmd/cue/cmd/get_go.go
+++ b/cmd/cue/cmd/get_go.go
@@ -33,6 +33,7 @@
 
 	"cuelang.org/go/cue/format"
 	"cuelang.org/go/cue/parser"
+	"cuelang.org/go/internal"
 	"github.com/spf13/cobra"
 	"golang.org/x/tools/go/packages"
 )
@@ -497,7 +498,7 @@
 				return err
 			}
 
-			if f.Name != nil && f.Name.Name == p.Name {
+			if _, pkg, _ := internal.PackageInfo(f); pkg != "" && pkg == p.Name {
 				file := filepath.Base(path)
 				file = file[:len(file)-len(".cue")]
 				file += "_gen.cue"
diff --git a/cmd/cue/cmd/import.go b/cmd/cue/cmd/import.go
index a7c5b26..ce53316 100644
--- a/cmd/cue/cmd/import.go
+++ b/cmd/cue/cmd/import.go
@@ -403,9 +403,6 @@
 	}
 
 	f := &ast.File{}
-	if pkg != "" {
-		f.Name = ast.NewIdent(pkg)
-	}
 
 	h := hoister{
 		fields:   map[string]bool{},
@@ -504,6 +501,11 @@
 		f.Decls = append([]ast.Decl{imports}, f.Decls...)
 	}
 
+	if pkg != "" {
+		p := &ast.Package{Name: ast.NewIdent(pkg)}
+		f.Decls = append([]ast.Decl{p}, f.Decls...)
+	}
+
 	if flagList.Bool(cmd) {
 		switch x := index.field.Value.(type) {
 		case *ast.StructLit:
diff --git a/cue/ast.go b/cue/ast.go
index 8e7f5bc..3a5db12 100644
--- a/cue/ast.go
+++ b/cue/ast.go
@@ -204,6 +204,9 @@
 		}
 		ret = obj
 
+	case *ast.Package:
+		v.walk(n.Name)
+
 	case *ast.ImportDecl:
 		for _, s := range n.Specs {
 			v.walk(s)
diff --git a/cue/ast/ast.go b/cue/ast/ast.go
index a36c1d1..4755321 100644
--- a/cue/ast/ast.go
+++ b/cue/ast/ast.go
@@ -86,7 +86,10 @@
 func (*BadDecl) declNode()           {}
 func (*EmbedDecl) declNode()         {}
 func (*Alias) declNode()             {}
-func (*CommentGroup) declNode()      {}
+
+// Not technically declarations, but appearing at the same level.
+func (*Package) declNode()      {}
+func (*CommentGroup) declNode() {}
 
 // A Label is any prduction that can be used as a LHS label.
 type Label interface {
@@ -599,39 +602,42 @@
 	return s.Path.End()
 }
 
-// specNode() ensures that only spec nodes can be assigned to a Spec.
+// A BadDecl node is a placeholder for declarations containing
+// syntax errors for which no correct declaration nodes can be
+// created.
+type BadDecl struct {
+	comments
+	From, To token.Pos // position range of bad declaration
+}
+
+// A ImportDecl node represents a series of import declarations. A valid
+// Lparen position (Lparen.Line > 0) indicates a parenthesized declaration.
+type ImportDecl struct {
+	comments
+	Import token.Pos
+	Lparen token.Pos // position of '(', if any
+	Specs  []*ImportSpec
+	Rparen token.Pos // position of ')', if any
+}
+
+type Spec interface {
+	Node
+	specNode()
+}
+
 func (*ImportSpec) specNode() {}
 
-// A declaration is represented by one of the following declaration nodes.
-type (
-	// A BadDecl node is a placeholder for declarations containing
-	// syntax errors for which no correct declaration nodes can be
-	// created.
-	BadDecl struct {
-		comments
-		From, To token.Pos // position range of bad declaration
-	}
+func (*Package) specNode() {}
 
-	// A ImportDecl node represents a series of import declarations. A valid
-	// Lparen position (Lparen.Line > 0) indicates a parenthesized declaration.
-	ImportDecl struct {
-		comments
-		Import token.Pos
-		Lparen token.Pos // position of '(', if any
-		Specs  []*ImportSpec
-		Rparen token.Pos // position of ')', if any
-	}
-
-	// An EmbedDecl node represents a single expression used as a declaration.
-	// The expressions in this declaration is what will be emitted as
-	// configuration output.
-	//
-	// An EmbedDecl may only appear at the top level.
-	EmbedDecl struct {
-		comments
-		Expr Expr
-	}
-)
+// An EmbedDecl node represents a single expression used as a declaration.
+// The expressions in this declaration is what will be emitted as
+// configuration output.
+//
+// An EmbedDecl may only appear at the top level.
+type EmbedDecl struct {
+	comments
+	Expr Expr
+}
 
 // Pos and End implementations for declaration nodes.
 
@@ -662,19 +668,13 @@
 type File struct {
 	Filename string
 	comments
-	Package token.Pos // position of "package" pseudo-keyword
-	Name    *Ident    // package names
-	Decls   []Decl    // top-level declarations; or nil
+	Decls []Decl // top-level declarations; or nil
 
-	// TODO: Change Expr to Decl?
 	Imports    []*ImportSpec // imports in this file
 	Unresolved []*Ident      // unresolved identifiers in this file
 }
 
 func (f *File) Pos() token.Pos {
-	if f.Package != token.NoPos {
-		return f.Package
-	}
 	if len(f.Decls) > 0 {
 		return f.Decls[0].Pos()
 	}
@@ -689,8 +689,29 @@
 	if n := len(f.Decls); n > 0 {
 		return f.Decls[n-1].End()
 	}
-	if f.Name != nil {
-		return f.Name.End()
+	return token.NoPos
+}
+
+// A Package represents a package clause.
+type Package struct {
+	comments
+	PackagePos token.Pos // position of "package" pseudo-keyword
+	Name       *Ident    // package name
+}
+
+func (p *Package) Pos() token.Pos {
+	if p.PackagePos != token.NoPos {
+		return p.PackagePos
+	}
+	if p.Name != nil {
+		return p.Name.Pos()
+	}
+	return token.NoPos
+}
+
+func (p *Package) End() token.Pos {
+	if p.Name != nil {
+		return p.Name.End()
 	}
 	return token.NoPos
 }
diff --git a/cue/ast/walk.go b/cue/ast/walk.go
index 89e93a6..8fb1ee9 100644
--- a/cue/ast/walk.go
+++ b/cue/ast/walk.go
@@ -177,11 +177,11 @@
 
 	// Files and packages
 	case *File:
-		if n.Name != nil {
-			walk(v, n.Name)
-		}
 		walkDeclList(v, n.Decls)
 
+	case *Package:
+		walk(v, n.Name)
+
 	case *ListComprehension:
 		walk(v, n.Expr)
 		for _, c := range n.Clauses {
diff --git a/cue/build.go b/cue/build.go
index 8487a2b..013e029 100644
--- a/cue/build.go
+++ b/cue/build.go
@@ -102,9 +102,7 @@
 	if err != nil {
 		return nil, err
 	}
-	if file.Name != nil {
-		p.PkgName = file.Name.Name
-	}
+	_, p.PkgName, _ = internal.PackageInfo(file)
 	return r.complete(p)
 }
 
diff --git a/cue/build/instance.go b/cue/build/instance.go
index f0c9ce9..dcde72b 100644
--- a/cue/build/instance.go
+++ b/cue/build/instance.go
@@ -24,6 +24,7 @@
 	"cuelang.org/go/cue/errors"
 	"cuelang.org/go/cue/parser"
 	"cuelang.org/go/cue/token"
+	"cuelang.org/go/internal"
 )
 
 // An Instance describes the collection of files, and its imports, necessary
@@ -177,12 +178,7 @@
 // AddSyntax adds the given file to list of files for this instance. The package
 // name of the file must match the package name of the instance.
 func (inst *Instance) AddSyntax(file *ast.File) errors.Error {
-	pkg := ""
-	pos := file.Pos()
-	if file.Name != nil {
-		pkg = file.Name.Name
-		pos = file.Name.Pos()
-	}
+	_, pkg, pos := internal.PackageInfo(file)
 	if !inst.setPkg(pkg) && pkg != inst.PkgName {
 		err := errors.Newf(pos,
 			"package name %q conflicts with previous package name %q",
diff --git a/cue/format/format.go b/cue/format/format.go
index 517ea7e..5458c01 100644
--- a/cue/format/format.go
+++ b/cue/format/format.go
@@ -50,6 +50,12 @@
 	return func(c *config) { c.TabIndent = indent }
 }
 
+// TODO: make public
+// sortImportsOption causes import declarations to be sorted.
+func sortImportsOption() Option {
+	return func(c *config) { c.sortImports = true }
+}
+
 // TODO: other options:
 //
 // const (
@@ -105,7 +111,8 @@
 	Tabwidth  int // default: 4
 	Indent    int // default: 0 (all code is indented at least by this much)
 
-	simplify bool
+	simplify    bool
+	sortImports bool
 }
 
 func newConfig(opt []Option) *config {
diff --git a/cue/format/format_test.go b/cue/format/format_test.go
index 3ac03e0..7fb3902 100644
--- a/cue/format/format_test.go
+++ b/cue/format/format_test.go
@@ -50,6 +50,7 @@
 	rawFormat
 	idempotent
 	simplify
+	sortImps
 )
 
 // format parses src, prints the corresponding AST, verifies the resulting
@@ -61,6 +62,9 @@
 	if mode&simplify != 0 {
 		opts = append(opts, Simplify())
 	}
+	if mode&sortImps != 0 {
+		opts = append(opts, sortImportsOption())
+	}
 
 	res, err := Source(src, opts...)
 	if err != nil {
@@ -189,6 +193,7 @@
 	{"comments.input", "comments.golden", 0},
 	{"simplify.input", "simplify.golden", simplify},
 	{"expressions.input", "expressions.golden", 0},
+	{"imports.input", "imports.golden", sortImps},
 }
 
 func TestFiles(t *testing.T) {
@@ -235,8 +240,8 @@
 }
 func TestPackage(t *testing.T) {
 	f := &ast.File{
-		Name: ast.NewIdent("foo"),
 		Decls: []ast.Decl{
+			&ast.Package{Name: ast.NewIdent("foo")},
 			&ast.EmbedDecl{
 				Expr: &ast.BasicLit{
 					ValuePos: token.NoSpace.Pos(),
@@ -364,13 +369,13 @@
 }
 
 var decls = []string{
-	`import "fmt"`,
-	"pi = 3.1415\ne = 2.71828\n\nx = pi",
+	"package p\n\n" + `import "fmt"`,
+	"package p\n\n" + "pi = 3.1415\ne = 2.71828\n\nx = pi",
 }
 
 func TestDeclLists(t *testing.T) {
 	for _, src := range decls {
-		file, err := parser.ParseFile("", "package p\n"+src, parser.ParseComments)
+		file, err := parser.ParseFile("", src, parser.ParseComments)
 		if err != nil {
 			panic(err) // error in test
 		}
diff --git a/cue/parser/import.go b/cue/format/import.go
similarity index 72%
rename from cue/parser/import.go
rename to cue/format/import.go
index 9b5afe6..873de2c 100644
--- a/cue/parser/import.go
+++ b/cue/format/import.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package parser
+package format
 
 import (
 	"sort"
@@ -23,45 +23,50 @@
 )
 
 // sortImports sorts runs of consecutive import lines in import blocks in f.
-// It also removes duplicate imports when it is possible to do so without data loss.
-func sortImports(f *ast.File) {
-	for _, d := range f.Decls {
-		d, ok := d.(*ast.ImportDecl)
-		if !ok {
-			// Not an import declaration, so we're done.
-			// Imports are always first.
-			break
-		}
+// It also removes duplicate imports when it is possible to do so without data
+// loss.
+func sortImports(d *ast.ImportDecl) {
+	if !d.Lparen.IsValid() || len(d.Specs) == 0 {
+		// Not a block: sorted by default.
+		return
+	}
 
-		if !d.Lparen.IsValid() {
-			// Not a block: sorted by default.
-			continue
-		}
-
-		// Identify and sort runs of specs on successive lines.
-		i := 0
-		specs := d.Specs[:0]
-		for j, s := range d.Specs {
-			if j > i && s.Pos().Line() > 1+d.Specs[j-1].End().Line() {
-				// j begins a new run. End this one.
-				specs = append(specs, sortSpecs(f, d.Specs[i:j])...)
-				i = j
-			}
-		}
-		specs = append(specs, sortSpecs(f, d.Specs[i:])...)
-		d.Specs = specs
-
-		// Deduping can leave a blank line before the rparen; clean that up.
-		if len(d.Specs) > 0 {
-			lastSpec := d.Specs[len(d.Specs)-1]
-			lastLine := lastSpec.Pos().Line()
-			rParenLine := d.Rparen.Line()
-			for rParenLine > lastLine+1 {
-				rParenLine--
-				d.Rparen.File().MergeLine(rParenLine)
-			}
+	// Identify and sort runs of specs on successive lines.
+	i := 0
+	specs := d.Specs[:0]
+	for j, s := range d.Specs {
+		if j > i && (s.Pos().RelPos() >= token.NewSection || hasDoc(s)) {
+			setRelativePos(s, token.Newline)
+			// j begins a new run. End this one.
+			block := sortSpecs(d.Specs[i:j])
+			specs = append(specs, block...)
+			i = j
 		}
 	}
+	specs = append(specs, sortSpecs(d.Specs[i:])...)
+	setRelativePos(specs[0], token.Newline)
+	d.Specs = specs
+}
+
+func setRelativePos(s *ast.ImportSpec, r token.RelPos) {
+	if hasDoc(s) {
+		return
+	}
+	pos := s.Pos().WithRel(r)
+	if s.Name != nil {
+		s.Name.NamePos = pos
+	} else {
+		s.Path.ValuePos = pos
+	}
+}
+
+func hasDoc(s *ast.ImportSpec) bool {
+	for _, doc := range s.Comments() {
+		if doc.Doc {
+			return true
+		}
+	}
+	return false
 }
 
 func importPath(s *ast.ImportSpec) string {
@@ -107,11 +112,12 @@
 	End   token.Pos
 }
 
-func sortSpecs(f *ast.File, specs []*ast.ImportSpec) []*ast.ImportSpec {
+func sortSpecs(specs []*ast.ImportSpec) []*ast.ImportSpec {
 	// Can't short-circuit here even if specs are already sorted,
 	// since they might yet need deduplication.
 	// A lone import, however, may be safely ignored.
 	if len(specs) <= 1 {
+		setRelativePos(specs[0], token.NewSection)
 		return specs
 	}
 
@@ -134,13 +140,11 @@
 	for i, s := range specs {
 		if i == len(specs)-1 || !collapse(s, specs[i+1]) {
 			deduped = append(deduped, s)
-		} else {
-			p := s.Pos()
-			p.File().MergeLine(p.Line())
 		}
 	}
 	specs = deduped
 
+	setRelativePos(specs[0], token.NewSection)
 	return specs
 }
 
diff --git a/cue/format/node.go b/cue/format/node.go
index 32076ad..34b45f9 100644
--- a/cue/format/node.go
+++ b/cue/format/node.go
@@ -85,7 +85,9 @@
 func (f *formatter) walkSpecList(list []*ast.ImportSpec) {
 	f.before(nil)
 	for _, x := range list {
+		f.before(x)
 		f.importSpec(x)
+		f.after(x)
 	}
 	f.after(nil)
 }
@@ -93,7 +95,9 @@
 func (f *formatter) walkClauseList(list []ast.Clause) {
 	f.before(nil)
 	for _, x := range list {
+		f.before(x)
 		f.clause(x)
+		f.after(x)
 	}
 	f.after(nil)
 }
@@ -111,12 +115,6 @@
 
 func (f *formatter) file(file *ast.File) {
 	f.before(file)
-	if file.Name != nil {
-		f.print(file.Package, "package")
-		f.print(blank, file.Name, newsection, nooverride)
-	}
-	f.current.pos = 3
-	f.visitComments(3)
 	f.walkDeclList(file.Decls)
 	f.after(file)
 	f.print(token.EOF)
@@ -223,6 +221,10 @@
 	case *ast.BadDecl:
 		f.print(n.From, "*bad decl*", declcomma)
 
+	case *ast.Package:
+		f.print(n.PackagePos, "package")
+		f.print(blank, n.Name, newsection, nooverride)
+
 	case *ast.ImportDecl:
 		f.print(n.Import, "import")
 		if len(n.Specs) == 0 {
diff --git a/cue/format/testdata/imports.golden b/cue/format/testdata/imports.golden
new file mode 100644
index 0000000..abd15e4
--- /dev/null
+++ b/cue/format/testdata/imports.golden
@@ -0,0 +1,26 @@
+package foo
+
+import (
+	"cuelang.org/go/foo"
+	"cuelang.org/go/bar"
+	"time"
+)
+
+import (
+	"time"
+
+	// comment f2
+	f2 "cuelang.org/go/foo"
+	f1 "cuelang.org/go/foo"
+)
+
+import (
+	"time"
+
+	same "cuelang.org/go/foo" // comment 2
+	same "cuelang.org/go/foo" // comment 1
+)
+
+a: time.time
+b: foo.foo
+c: bar.Bar
diff --git a/cue/format/testdata/imports.input b/cue/format/testdata/imports.input
new file mode 100644
index 0000000..3d526ca
--- /dev/null
+++ b/cue/format/testdata/imports.input
@@ -0,0 +1,27 @@
+package foo
+
+import (
+    "cuelang.org/go/foo"
+    "cuelang.org/go/bar"
+    "time"
+)
+
+import (
+    "time"
+
+    // comment f2
+    f2 "cuelang.org/go/foo"
+    f1 "cuelang.org/go/foo"
+)
+
+import (
+    "time"
+
+    same "cuelang.org/go/foo" // comment 2
+    same "cuelang.org/go/foo" // comment 1
+)
+
+
+a: time.time
+b: foo.foo
+c: bar.Bar
\ No newline at end of file
diff --git a/cue/instance.go b/cue/instance.go
index 2b1d551..d79a0cd 100644
--- a/cue/instance.go
+++ b/cue/instance.go
@@ -15,8 +15,6 @@
 package cue
 
 import (
-	"strings"
-
 	"cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/build"
 	"cuelang.org/go/cue/errors"
@@ -147,11 +145,12 @@
 		return nil
 	}
 	for _, f := range inst.inst.Files {
-		if strings.HasPrefix(f.Filename, inst.Dir) {
+		pkg, _, _ := internal.PackageInfo(f)
+		if pkg == nil {
 			continue
 		}
 		var cg *ast.CommentGroup
-		for _, c := range f.Comments() {
+		for _, c := range pkg.Comments() {
 			if c.Position == 0 {
 				cg = c
 			}
diff --git a/cue/load/import.go b/cue/load/import.go
index cc7c9b8..46c44f6 100644
--- a/cue/load/import.go
+++ b/cue/load/import.go
@@ -31,6 +31,7 @@
 	"cuelang.org/go/cue/errors"
 	"cuelang.org/go/cue/parser"
 	"cuelang.org/go/cue/token"
+	"cuelang.org/go/internal"
 )
 
 // An importMode controls the behavior of the Import method.
@@ -357,10 +358,7 @@
 		return true
 	}
 
-	pkg := ""
-	if pf.Name != nil {
-		pkg = pf.Name.Name
-	}
+	_, pkg, _ := internal.PackageInfo(pf)
 	if pkg == "" && mode&allowAnonymous == 0 {
 		p.IgnoredCUEFiles = append(p.IgnoredCUEFiles, fullPath)
 		return false // don't mark as added
diff --git a/cue/marshal.go b/cue/marshal.go
index 43174d0..0feb4dd 100644
--- a/cue/marshal.go
+++ b/cue/marshal.go
@@ -149,7 +149,8 @@
 			}
 		}
 		if i.Name != "" {
-			file.Name = ast.NewIdent(i.Name)
+			pkg := &ast.Package{Name: ast.NewIdent(i.Name)}
+			file.Decls = append([]ast.Decl{pkg}, file.Decls...)
 		}
 
 		b, err := format.Node(file)
diff --git a/cue/parser/interface.go b/cue/parser/interface.go
index 0711b20..8090883 100644
--- a/cue/parser/interface.go
+++ b/cue/parser/interface.go
@@ -123,7 +123,6 @@
 			// ParseFile API and return a valid (but) empty
 			// *File
 			f = &ast.File{
-				Name: new(ast.Ident),
 				// Scope: NewScope(nil),
 			}
 		}
diff --git a/cue/parser/parser.go b/cue/parser/parser.go
index c331dc5..32e104e 100644
--- a/cue/parser/parser.go
+++ b/cue/parser/parser.go
@@ -1319,25 +1319,32 @@
 	if p.errors != nil {
 		return nil
 	}
+	p.openList()
+
+	var decls []ast.Decl
 
 	// The package clause is not a declaration: it does not appear in any
 	// scope.
-	pos := p.pos
-	var name *ast.Ident
 	if p.tok == token.IDENT && p.lit == "package" {
+		c := p.openComments()
+
+		pos := p.pos
+		var name *ast.Ident
 		p.expect(token.IDENT)
 		name = p.parseIdent()
 		if name.Name == "_" && p.mode&declarationErrorsMode != 0 {
 			p.errf(p.pos, "invalid package name _")
 		}
-		p.expectComma()
-	} else {
-		pos = token.NoPos
-	}
-	c.pos = 3
 
-	p.openList()
-	var decls []ast.Decl
+		pkg := &ast.Package{
+			PackagePos: pos,
+			Name:       name,
+		}
+		decls = append(decls, pkg)
+		p.expectComma()
+		c.closeNode(p, pkg)
+	}
+
 	if p.mode&packageClauseOnlyMode == 0 {
 		// import decls
 		for p.tok == token.IDENT && p.lit == "import" {
@@ -1354,8 +1361,6 @@
 	p.closeList()
 
 	f := &ast.File{
-		Package: pos,
-		Name:    name,
 		Imports: p.imports,
 		Decls:   decls,
 	}
diff --git a/cue/parser/parser_test.go b/cue/parser/parser_test.go
index 22f5728..557233e 100644
--- a/cue/parser/parser_test.go
+++ b/cue/parser/parser_test.go
@@ -255,7 +255,7 @@
 		// file.2
 
 		`,
-		"<[0// foo] [d0// uni] [l3// uniline] [3// file.1 // file.2] package foo, >",
+		"<[0// foo] <[d0// uni] [l3// uniline] [3// file.1 // file.2] package foo>>",
 	}, {
 		"line comments",
 		`// doc
@@ -470,11 +470,6 @@
 	}
 }
 
-func labelName(l ast.Label) string {
-	name, _ := ast.LabelName(l)
-	return name
-}
-
 // TestIncompleteSelection ensures that an incomplete selector
 // expression is parsed as a (blank) *SelectorExpr, not a
 // *BadExpr.
diff --git a/cue/parser/print.go b/cue/parser/print.go
index 2a91d64..9849e11 100644
--- a/cue/parser/print.go
+++ b/cue/parser/print.go
@@ -41,14 +41,14 @@
 	switch v := x.(type) {
 	case *ast.File:
 		out := ""
-		if v.Name != nil {
-			out += "package "
-			out += debugStr(v.Name)
-			out += ", "
-		}
 		out += debugStr(v.Decls)
 		return out
 
+	case *ast.Package:
+		out := "package "
+		out += debugStr(v.Name)
+		return out
+
 	case *ast.Alias:
 		out := debugStr(v.Ident)
 		out += " = "
diff --git a/cue/parser/walk.go b/cue/parser/walk.go
index eebe315..58b3085 100644
--- a/cue/parser/walk.go
+++ b/cue/parser/walk.go
@@ -172,10 +172,12 @@
 
 	// Files and packages
 	case *ast.File:
+		walkDeclList(v, n.Decls)
+
+	case *ast.Package:
 		if n.Name != nil {
 			walk(v, n.Name)
 		}
-		walkDeclList(v, n.Decls)
 
 	case *ast.ListComprehension:
 		walk(v, n.Expr)
diff --git a/cue/types_test.go b/cue/types_test.go
index e31e6b1..ea8ef72 100644
--- a/cue/types_test.go
+++ b/cue/types_test.go
@@ -1617,9 +1617,19 @@
 	// Another Foo.
 	Foo: {}
 	`
-	inst := getInstance(t, config)
+	var r Runtime
+	getInst := func(name, body string) *Instance {
+		inst, err := r.Compile("dir/file1.cue", body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return inst
+	}
+
+	inst := getInst("config", config)
+
 	v1 := inst.Value()
-	v2 := getInstance(t, config2).Value()
+	v2 := getInst("config2", config2).Value()
 	both := v1.Unify(v2)
 
 	testCases := []struct {
@@ -1668,6 +1678,11 @@
 		path: "baz field2",
 		doc:  "comment from bar on field 2\n",
 	}, {
+		val:  v2,
+		path: "Foo",
+		doc: `Another Foo.
+`,
+	}, {
 		val:  both,
 		path: "Foo",
 		doc: `Another Foo.
diff --git a/encoding/protobuf/parse.go b/encoding/protobuf/parse.go
index 4df5977..6432fce 100644
--- a/encoding/protobuf/parse.go
+++ b/encoding/protobuf/parse.go
@@ -112,11 +112,9 @@
 					default:
 						failf(x.Position, "unexpected ';' in %q", str)
 					}
-					p.file.Name = ast.NewIdent(p.shortPkgName)
 
 				case len(split) == 1:
 					p.shortPkgName = split[0]
-					p.file.Name = ast.NewIdent(p.shortPkgName)
 
 				default:
 					failf(x.Position, "malformed go_package clause %s", str)
@@ -127,6 +125,10 @@
 		}
 	}
 
+	if name := p.shortName(); name != "" {
+		p.file.Decls = append(p.file.Decls, &ast.Package{Name: ast.NewIdent(name)})
+	}
+
 	for _, e := range d.Elements {
 		switch x := e.(type) {
 		case *proto.Import:
@@ -137,6 +139,7 @@
 	}
 
 	imports := &ast.ImportDecl{}
+	importIdx := len(p.file.Decls)
 	p.file.Decls = append(p.file.Decls, imports)
 
 	for _, e := range d.Elements {
@@ -159,7 +162,9 @@
 	}
 
 	if len(imports.Specs) == 0 {
-		p.file.Decls = p.file.Decls[1:]
+		a := p.file.Decls
+		copy(a[importIdx:], a[importIdx+1:])
+		p.file.Decls = a[:len(a)-1]
 	}
 
 	return p, nil
@@ -209,7 +214,6 @@
 	if p.shortPkgName == "" && p.protoPkg != "" {
 		split := strings.Split(p.protoPkg, ".")
 		p.shortPkgName = split[len(split)-1]
-		p.file.Name = ast.NewIdent(p.shortPkgName)
 	}
 	return p.shortPkgName
 }
diff --git a/internal/internal.go b/internal/internal.go
index fa417ab..a7cdc42 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -21,6 +21,7 @@
 
 import (
 	"cuelang.org/go/cue/ast"
+	"cuelang.org/go/cue/token"
 	"github.com/cockroachdb/apd/v2"
 )
 
@@ -82,3 +83,17 @@
 	}
 	return elts, e
 }
+
+func PackageInfo(f *ast.File) (p *ast.Package, name string, tok token.Pos) {
+	for _, d := range f.Decls {
+		switch x := d.(type) {
+		case *ast.CommentGroup:
+		case *ast.Package:
+			if x.Name == nil {
+				break
+			}
+			return x, x.Name.Name, x.Name.Pos()
+		}
+	}
+	return nil, "", f.Pos()
+}