cue/load: rewrite pre-resolved files after injection

This also moves tag discover to load.Instances to
prevent finding tags in imported packages. This
"feature" slipped in when moving this functionaltiy
to cue/load.

Change-Id: Ib04d826d95ffebd152310aaad26d0a896991f076
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/7121
Reviewed-by: Paul Jolly <paul@myitcv.org.uk>
Reviewed-by: CUE cueckoo <cueckoo@gmail.com>
diff --git a/cmd/cue/cmd/testdata/script/issue520.txt b/cmd/cue/cmd/testdata/script/issue520.txt
new file mode 100644
index 0000000..16fb107
--- /dev/null
+++ b/cmd/cue/cmd/testdata/script/issue520.txt
@@ -0,0 +1,30 @@
+cue cmd -t greeting=hello prefix
+
+-- cue.mod/module.cue --
+module: "mod.com"
+-- my.cue --
+package tools
+
+msg: string @tag(greeting)
+
+-- my_tool.cue --
+package tools
+
+import (
+	"tool/cli"
+)
+
+greeting: string @tag(greeting)
+
+command: prefix: {
+	p1: cli.Print & {
+		text: greeting
+	}
+    salutation: string @tag(greeting)
+	p2: cli.Print & {
+		text: salutation
+	}
+	p3: cli.Print & {
+		text: msg @tag(greeting)
+	}
+}
diff --git a/cue/load/loader.go b/cue/load/loader.go
index 61fa06c..4d33065 100644
--- a/cue/load/loader.go
+++ b/cue/load/loader.go
@@ -27,6 +27,7 @@
 
 	"golang.org/x/xerrors"
 
+	"cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/build"
 	"cuelang.org/go/cue/errors"
 	"cuelang.org/go/cue/token"
@@ -84,6 +85,14 @@
 		a = append(a, l.cueFilesPackage(files))
 	}
 
+	for _, p := range a {
+		tags, err := findTags(p)
+		if err != nil {
+			p.ReportError(err)
+		}
+		l.tags = append(l.tags, tags...)
+	}
+
 	// TODO(api): have API call that returns an error which is the aggregate
 	// of all build errors. Certain errors, like these, hold across builds.
 	if err := injectTags(c.Tags, l); err != nil {
@@ -92,6 +101,22 @@
 		}
 	}
 
+	if l.replacements == nil {
+		return a
+	}
+
+	for _, p := range a {
+		for _, f := range p.Files {
+			ast.Walk(f, nil, func(n ast.Node) {
+				if ident, ok := n.(*ast.Ident); ok {
+					if v, ok := l.replacements[ident.Node]; ok {
+						ident.Node = v
+					}
+				}
+			})
+		}
+	}
+
 	return a
 }
 
@@ -110,10 +135,11 @@
 )
 
 type loader struct {
-	cfg       *Config
-	stk       importStack
-	tags      []tag // tags found in files
-	buildTags map[string]bool
+	cfg          *Config
+	stk          importStack
+	tags         []tag // tags found in files
+	buildTags    map[string]bool
+	replacements map[ast.Node]ast.Node
 }
 
 func (l *loader) abs(filename string) string {
@@ -205,11 +231,6 @@
 		}
 		d.Close()
 	}
-	tags, err := findTags(p)
-	if err != nil {
-		p.ReportError(err)
-	}
-	l.tags = append(l.tags, tags...)
 }
 
 func cleanImport(path string) string {
diff --git a/cue/load/tags.go b/cue/load/tags.go
index 1417c0b..d09f97d 100644
--- a/cue/load/tags.go
+++ b/cue/load/tags.go
@@ -88,12 +88,17 @@
 	return t, nil
 }
 
-func (t *tag) inject(value string) errors.Error {
+func (t *tag) inject(value string, l *loader) errors.Error {
 	e, err := cli.ParseValue(token.NoPos, t.key, value, t.kind)
 	if err != nil {
 		return err
 	}
-	t.field.Value = ast.NewBinExpr(token.AND, t.field.Value, e)
+	injected := ast.NewBinExpr(token.AND, t.field.Value, e)
+	if l.replacements == nil {
+		l.replacements = map[ast.Node]ast.Node{}
+	}
+	l.replacements[t.field.Value] = injected
+	t.field.Value = injected
 	return nil
 }
 
@@ -161,7 +166,7 @@
 			for _, t := range l.tags {
 				if t.key == s[:p] {
 					found = true
-					if err := t.inject(s[p+1:]); err != nil {
+					if err := t.inject(s[p+1:], l); err != nil {
 						return err
 					}
 				}
@@ -174,7 +179,7 @@
 				for _, sh := range t.shorthands {
 					if sh == s {
 						found = true
-						if err := t.inject(s); err != nil {
+						if err := t.inject(s, l); err != nil {
 							return err
 						}
 					}