cue: refactor marshal to simplify move

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

cue/export.go

@@ -33,6 +33,24 @@ func doEval(m options) bool {
 	return !m.raw
 }
 
+func exportFile(ctx *context, inst *Instance, v value, m options) (f *ast.File, imports []string) {
+	e := exporter{ctx, m, nil, map[label]bool{}, map[string]importInfo{}, false, nil}
+	top, ok := v.evalPartial(ctx).(*structLit)
+	if ok {
+		top, err := top.expandFields(ctx)
+		if err != nil {
+			v = err
+		} else {
+			for _, a := range top.Arcs {
+				e.top[a.Label] = true
+			}
+		}
+	}
+
+	value := e.expr(v)
+	return e.toFile(inst, value, m)
+}
+
 func export(ctx *context, inst *Instance, v value, m options) (n ast.Node, imports []string) {
 	e := exporter{ctx, m, nil, map[label]bool{}, map[string]importInfo{}, false, nil}
 	top, ok := v.evalPartial(ctx).(*structLit)
@@ -52,7 +70,10 @@ func export(ctx *context, inst *Instance, v value, m options) (n ast.Node, impor
 		// TODO: unwrap structs?
 		return value, nil
 	}
+	return e.toFile(inst, value, m)
+}
 
+func (e *exporter) toFile(inst *Instance, value ast.Expr, m options) (n *ast.File, imports []string) {
 	file := &ast.File{}
 	if inst != nil {
 		if inst.PkgName != "" {

cue/marshal.go

@@ -22,10 +22,12 @@ import (
 	"strings"
 
 	"cuelang.org/go/cue/ast"
+	"cuelang.org/go/cue/ast/astutil"
 	"cuelang.org/go/cue/build"
 	"cuelang.org/go/cue/errors"
 	"cuelang.org/go/cue/format"
 	"cuelang.org/go/cue/token"
+	"cuelang.org/go/internal"
 )
 
 // root.
@@ -139,17 +141,15 @@ func (r *Runtime) Marshal(instances ...*Instance) (b []byte, err error) {
 			return p
 		}
 		// TODO: support exporting instance
-		n, imports := export(ctx, nil, i.rootValue, options{raw: true})
-
-		file, ok := n.(*ast.File)
-		if !ok {
-			file = &ast.File{}
-			if obj, ok := n.(*ast.StructLit); ok {
-				file.Decls = append(file.Decls, obj.Elts...)
-			} else {
-				file.Decls = append(file.Decls, &ast.EmbedDecl{Expr: n.(ast.Expr)})
+		file, _ := exportFile(ctx, nil, i.rootValue, options{raw: true})
+		imports := []string{}
+		for _, i := range internal.Imports(file) {
+			for _, spec := range i.(*ast.ImportDecl).Specs {
+				info, _ := astutil.ParseImportSpec(spec)
+				imports = append(imports, info.ID)
 			}
 		}
+
 		if i.PkgName != "" {
 			pkg := &ast.Package{Name: ast.NewIdent(i.PkgName)}
 			file.Decls = append([]ast.Decl{pkg}, file.Decls...)