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>
diff --git a/cmd/cue/cmd/orphans.go b/cmd/cue/cmd/orphans.go
index de54d92..7032e1a 100644
--- a/cmd/cue/cmd/orphans.go
+++ b/cmd/cue/cmd/orphans.go
@@ -151,19 +151,7 @@
}
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
}
diff --git a/cue/ast/ast.go b/cue/ast/ast.go
index 78117b6..be6b8a2 100644
--- a/cue/ast/ast.go
+++ b/cue/ast/ast.go
@@ -967,6 +967,40 @@
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 {
@@ -974,7 +1008,7 @@
switch x := d.(type) {
case *Package:
return x.Name.Name
- case *CommentGroup:
+ case *CommentGroup, *Attribute:
default:
return ""
}
diff --git a/cue/ast/astutil/sanitize.go b/cue/ast/astutil/sanitize.go
index 193dae8..061a46b 100644
--- a/cue/ast/astutil/sanitize.go
+++ b/cue/ast/astutil/sanitize.go
@@ -169,23 +169,16 @@
}
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 {
diff --git a/cue/ast/astutil/util.go b/cue/ast/astutil/util.go
index e3439b8..2a8ab55 100644
--- a/cue/ast/astutil/util.go
+++ b/cue/ast/astutil/util.go
@@ -109,17 +109,23 @@
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 @@
// 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
}
diff --git a/cue/build.go b/cue/build.go
index 65f0932..5587650 100644
--- a/cue/build.go
+++ b/cue/build.go
@@ -306,18 +306,11 @@
}
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) {
diff --git a/cue/marshal.go b/cue/marshal.go
index c48d136..98a3511 100644
--- a/cue/marshal.go
+++ b/cue/marshal.go
@@ -144,12 +144,12 @@
// 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)
diff --git a/encoding/openapi/decode.go b/encoding/openapi/decode.go
index a8a61b8..a90efd3 100644
--- a/encoding/openapi/decode.go
+++ b/encoding/openapi/decode.go
@@ -65,17 +65,16 @@
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 @@
}
}
- 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
diff --git a/internal/core/runtime/runtime.go b/internal/core/runtime/runtime.go
index 3a3c6cc..fd53f0c 100644
--- a/internal/core/runtime/runtime.go
+++ b/internal/core/runtime/runtime.go
@@ -69,18 +69,11 @@
// 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 {
diff --git a/internal/internal.go b/internal/internal.go
index 066d26f..0aaff75 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -112,21 +112,6 @@
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 @@
if cgs = ast.Comments(d); cgs != nil {
break
}
+ // TODO: what to do here?
if _, ok := d.(*ast.Attribute); !ok {
break
}