cmd/cue/cmd: hoist expression logic into iterator

Change-Id: I1edc0a43e4709b1146b2c4cef9a4728e84cf9f95
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/5095
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
diff --git a/cmd/cue/cmd/common.go b/cmd/cue/cmd/common.go
index 7e0a220..81d9550 100644
--- a/cmd/cue/cmd/common.go
+++ b/cmd/cue/cmd/common.go
@@ -143,10 +143,20 @@
 // data files. In the latter case, there must be either 0 or 1 other
 // instance, with which the data instance may be merged.
 func (b *buildPlan) instances() iterator {
+	var i iterator
 	if len(b.orphanedData) == 0 && len(b.orphanedSchema) == 0 {
-		return &instanceIterator{a: buildInstances(b.cmd, b.insts), i: -1}
+		i = &instanceIterator{a: buildInstances(b.cmd, b.insts), i: -1}
+	} else {
+		i = newStreamingIterator(b)
 	}
-	return newStreamingIterator(b)
+	if len(b.expressions) > 0 {
+		return &expressionIter{
+			iter: i,
+			expr: b.expressions,
+			i:    len(b.expressions),
+		}
+	}
+	return i
 }
 
 type iterator interface {
@@ -297,6 +307,41 @@
 	return i.e
 }
 
+type expressionIter struct {
+	iter iterator
+	expr []ast.Expr
+	i    int
+}
+
+func (i *expressionIter) err() error { return i.iter.err() }
+func (i *expressionIter) close()     { i.iter.close() }
+func (i *expressionIter) id() string { return i.iter.id() }
+
+func (i *expressionIter) scan() bool {
+	i.i++
+	if i.i < len(i.expr) {
+		return true
+	}
+	if !i.iter.scan() {
+		return false
+	}
+	i.i = 0
+	return true
+}
+
+func (i *expressionIter) file() *ast.File { return nil }
+
+func (i *expressionIter) instance() *cue.Instance {
+	if len(i.expr) == 0 {
+		return i.iter.instance()
+	}
+	inst := i.iter.instance()
+	v := i.iter.instance().Eval(i.expr[i.i])
+	ni := internal.MakeInstance(v).(*cue.Instance)
+	ni.DisplayName = fmt.Sprintf("%s|%s", inst.DisplayName, i.expr[i.i])
+	return ni
+}
+
 func parseArgs(cmd *Command, args []string, cfg *load.Config) (p *buildPlan, err error) {
 	if cfg == nil {
 		cfg = defaultConfig
diff --git a/cmd/cue/cmd/eval.go b/cmd/cue/cmd/eval.go
index 8b1ea7c..c2f47f4 100644
--- a/cmd/cue/cmd/eval.go
+++ b/cmd/cue/cmd/eval.go
@@ -23,7 +23,6 @@
 	"cuelang.org/go/cue"
 	"cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/format"
-	"cuelang.org/go/internal"
 )
 
 // newEvalCmd creates a new eval command
@@ -100,9 +99,7 @@
 
 	iter := b.instances()
 	defer iter.close()
-	for iter.scan() {
-		inst := iter.instance().Value()
-
+	for i := 0; iter.scan(); i++ {
 		// TODO: use ImportPath or some other sanitized path.
 		if len(b.insts) > 1 {
 			fmt.Fprintf(w, "\n// %s\n", iter.id())
@@ -129,33 +126,20 @@
 			opts = append(opts, format.Simplify())
 		}
 
-		if b.expressions == nil {
-			v := v
-			if flagConcrete.Bool(cmd) && !flagIgnore.Bool(cmd) {
-				if err := v.Validate(cue.Concrete(true)); err != nil {
-					exitOnErr(cmd, err, false)
-					continue
-				}
-			}
-			writeNode(format.Node(getSyntax(v, syn), opts...))
+		if len(b.expressions) > 1 {
+			fmt.Fprint(w, "// ")
+			writeNode(format.Node(b.expressions[i%len(b.expressions)]))
 		}
-		for _, e := range b.expressions {
-			if len(b.expressions) > 1 {
-				fmt.Fprint(w, "// ")
-				writeNode(format.Node(e))
-			}
-			v := internal.EvalExpr(inst, e).(cue.Value)
-			if err := v.Err(); err != nil {
-				return err
-			}
-			if flagConcrete.Bool(cmd) && !flagIgnore.Bool(cmd) {
-				if err := v.Validate(cue.Concrete(true)); err != nil {
-					exitOnErr(cmd, err, false)
-					continue
-				}
-			}
-			writeNode(format.Node(getSyntax(v, syn), opts...))
+		if err := v.Err(); err != nil {
+			return err
 		}
+		if flagConcrete.Bool(cmd) && !flagIgnore.Bool(cmd) {
+			if err := v.Validate(cue.Concrete(true)); err != nil {
+				exitOnErr(cmd, err, false)
+				continue
+			}
+		}
+		writeNode(format.Node(getSyntax(v, syn), opts...))
 	}
 	exitOnErr(cmd, iter.err(), true)
 	return nil
diff --git a/cmd/cue/cmd/export.go b/cmd/cue/cmd/export.go
index 194a2da..c7d6f9c 100644
--- a/cmd/cue/cmd/export.go
+++ b/cmd/cue/cmd/export.go
@@ -17,8 +17,6 @@
 import (
 	"github.com/spf13/cobra"
 
-	"cuelang.org/go/cue"
-	"cuelang.org/go/internal"
 	"cuelang.org/go/internal/encoding"
 	"cuelang.org/go/internal/filetypes"
 )
@@ -121,18 +119,8 @@
 	defer iter.close()
 	for iter.scan() {
 		inst := iter.instance()
-
-		if b.expressions == nil {
-			err = enc.Encode(inst)
-			exitOnErr(cmd, err, true)
-			continue
-		}
-		for _, e := range b.expressions {
-			v := internal.MakeInstance(inst.Eval(e)).(*cue.Instance)
-			exitOnErr(cmd, v.Err, true)
-			err = enc.Encode(v)
-			exitOnErr(cmd, err, true)
-		}
+		err = enc.Encode(inst)
+		exitOnErr(cmd, err, true)
 	}
 	exitOnErr(cmd, iter.err(), true)
 	return nil