cue: clean up handling of preamble declarations

There is a lot of code that needs to find the cutoff point
for preamble vs body in the declarations of an ast.File.

This cleans this up.

Change-Id: I322542f1d5bade6d67ea70bd1085175e7ff8f7c1
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/7086
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
Reviewed-by: CUE cueckoo <cueckoo@gmail.com>

cmd/cue/cmd/orphans.go

@@ -151,19 +151,7 @@ func (b *buildPlan) placeOrphans(i *build.Instance) (ok bool, err error) {
 }
 
 func toExpr(f *ast.File) (expr ast.Expr, pkg *ast.Package) {
-	var p int
-outer:
-	for i, d := range f.Decls {
-		switch x := d.(type) {
-		case *ast.Package:
-			pkg = x
-		case *ast.ImportDecl:
-			p = i + 1
-		case *ast.CommentGroup:
-		default:
-			break outer
-		}
-	}
+	p := len(f.Preamble())
 	return &ast.StructLit{Elts: f.Decls[p:]}, pkg
 }
 

cue/ast/ast.go

@@ -967,14 +967,48 @@ type File struct {
 	comments
 }
 
+// Preamble returns the declarations of the preamble.
+func (f *File) Preamble() []Decl {
+	p := 0
+outer:
+	for i, d := range f.Decls {
+		switch d.(type) {
+		default:
+			break outer
+
+		case *Package:
+			p = i + 1
+		case *CommentGroup:
+		case *Attribute:
+		case *ImportDecl:
+			p = i + 1
+		}
+	}
+	return f.Decls[:p]
+}
+
+func (f *File) VisitImports(fn func(d *ImportDecl)) {
+	for _, d := range f.Decls {
+		switch x := d.(type) {
+		case *CommentGroup:
+		case *Package:
+		case *Attribute:
+		case *ImportDecl:
+			fn(x)
+		default:
+			return
+		}
+	}
+}
+
 // PackageName returns the package name associated with this file or "" if no
 // package is associated.
 func (f *File) PackageName() string {
 	for _, d := range f.Decls {
 		switch x := d.(type) {
 		case *Package:
 			return x.Name.Name
-		case *CommentGroup:
+		case *CommentGroup, *Attribute:
 		default:
 			return ""
 		}

cue/ast/astutil/sanitize.go

@@ -169,23 +169,16 @@ func (z *sanitizer) markUsed(s *scope, n *ast.Ident) bool {
 }
 
 func (z *sanitizer) cleanImports() {
-	for _, d := range z.file.Decls {
-		switch id := d.(type) {
-		case *ast.Package, *ast.CommentGroup:
-		case *ast.ImportDecl:
-			k := 0
-			for _, s := range id.Specs {
-				if _, ok := z.referenced[s]; ok {
-					id.Specs[k] = s
-					k++
-				}
+	z.file.VisitImports(func(d *ast.ImportDecl) {
+		k := 0
+		for _, s := range d.Specs {
+			if _, ok := z.referenced[s]; ok {
+				d.Specs[k] = s
+				k++
 			}
-			id.Specs = id.Specs[:k]
-
-		default:
-			return
 		}
-	}
+		d.Specs = d.Specs[:k]
+	})
 }
 
 func (z *sanitizer) handleIdent(s *scope, n *ast.Ident) bool {

cue/ast/astutil/util.go

@@ -109,17 +109,23 @@ func insertImport(decls *[]ast.Decl, spec *ast.ImportSpec) *ast.ImportSpec {
 
 	var imports *ast.ImportDecl
 	var orig *ast.ImportSpec
-	i := 0
+
+	p := 0
 outer:
-	for ; i < len(a); i++ {
+	for i := 0; i < len(a); i++ {
 		d := a[i]
 		switch t := d.(type) {
 		default:
 			break outer
 
 		case *ast.Package:
+			p = i + 1
 		case *ast.CommentGroup:
+			p = i + 1
+		case *ast.Attribute:
+			continue
 		case *ast.ImportDecl:
+			p = i + 1
 			imports = t
 			for _, s := range t.Specs {
 				y, _ := ParseImportSpec(s)
@@ -137,8 +143,8 @@ outer:
 	// Import not found, add one.
 	if imports == nil {
 		imports = &ast.ImportDecl{}
-		preamble := append(a[:i:i], imports)
-		a = append(preamble, a[i:]...)
+		preamble := append(a[:p:p], imports)
+		a = append(preamble, a[p:]...)
 		*decls = a
 	}
 

cue/build.go

@@ -306,18 +306,11 @@ func (v *visitor) addErr(e errors.Error) {
 }
 
 func (v *visitor) file(file *ast.File) {
-	for _, d := range file.Decls {
-		switch x := d.(type) {
-		case *ast.Package:
-		case *ast.ImportDecl:
-			for _, s := range x.Specs {
-				v.spec(s)
-			}
-		case *ast.CommentGroup:
-		default:
-			return
+	file.VisitImports(func(x *ast.ImportDecl) {
+		for _, s := range x.Specs {
+			v.spec(s)
 		}
-	}
+	})
 }
 
 func (v *visitor) spec(spec *ast.ImportSpec) {

cue/marshal.go

@@ -144,12 +144,12 @@ func (r *Runtime) Marshal(instances ...*Instance) (b []byte, err error) {
 		// TODO: support exporting instance
 		file, _ := export.Def(r.idx.Runtime, i.root)
 		imports := []string{}
-		for _, i := range internal.Imports(file) {
-			for _, spec := range i.(*ast.ImportDecl).Specs {
+		file.VisitImports(func(i *ast.ImportDecl) {
+			for _, spec := range i.Specs {
 				info, _ := astutil.ParseImportSpec(spec)
 				imports = append(imports, info.ID)
 			}
-		}
+		})
 
 		if i.PkgName != "" {
 			p, name, _ := internal.PackageInfo(file)

encoding/openapi/decode.go

@@ -65,17 +65,16 @@ func Extract(data *cue.Instance, c *Config) (*ast.File, error) {
 		add(cg)
 	}
 
-	i := 0
-	for ; i < len(js.Decls); i++ {
-		switch x := js.Decls[i].(type) {
+	preamble := js.Preamble()
+	body := js.Decls[len(preamble):]
+	for _, d := range preamble {
+		switch x := d.(type) {
 		case *ast.Package:
 			return nil, errors.Newf(x.Pos(), "unexpected package %q", x.Name.Name)
 
-		case *ast.ImportDecl, *ast.CommentGroup:
+		default:
 			add(x)
-			continue
 		}
-		break
 	}
 
 	// TODO: allow attributes before imports? Would be easier.
@@ -112,9 +111,9 @@ func Extract(data *cue.Instance, c *Config) (*ast.File, error) {
 		}
 	}
 
-	if i < len(js.Decls) {
-		ast.SetRelPos(js.Decls[i], token.NewSection)
-		f.Decls = append(f.Decls, js.Decls[i:]...)
+	if len(body) > 0 {
+		ast.SetRelPos(body[0], token.NewSection)
+		f.Decls = append(f.Decls, body...)
 	}
 
 	return f, nil

internal/core/runtime/runtime.go

@@ -69,18 +69,11 @@ func (x *Runtime) Build(b *build.Instance) (v *adt.Vertex, errs errors.Error) {
 
 	// Build transitive dependencies.
 	for _, file := range b.Files {
-		for _, d := range file.Decls {
-			switch g := d.(type) {
-			case *ast.Package:
-			case *ast.ImportDecl:
-				for _, s := range g.Specs {
-					errs = errors.Append(errs, x.buildSpec(b, s))
-				}
-			case *ast.CommentGroup:
-			default:
-				break
+		file.VisitImports(func(d *ast.ImportDecl) {
+			for _, s := range d.Specs {
+				errs = errors.Append(errs, x.buildSpec(b, s))
 			}
-		}
+		})
 	}
 
 	if errs != nil {

internal/internal.go

@@ -112,21 +112,6 @@ func ListEllipsis(n *ast.ListLit) (elts []ast.Expr, e *ast.Ellipsis) {
 	return elts, e
 }
 
-func Imports(f *ast.File) (a []ast.Decl) {
-	for _, d := range f.Decls {
-		switch x := d.(type) {
-		case *ast.CommentGroup:
-		case *ast.Package:
-		case *ast.Attribute:
-		case *ast.ImportDecl:
-			a = append(a, x)
-		default:
-			return a
-		}
-	}
-	return a
-}
-
 func PackageInfo(f *ast.File) (p *ast.Package, name string, tok token.Pos) {
 	for _, d := range f.Decls {
 		switch x := d.(type) {
@@ -225,6 +210,7 @@ func FileComment(f *ast.File) *ast.CommentGroup {
 			if cgs = ast.Comments(d); cgs != nil {
 				break
 			}
+			// TODO: what to do here?
 			if _, ok := d.(*ast.Attribute); !ok {
 				break
 			}