cue/load: add support for build tags

Fixes #511

Change-Id: I012286cffe357ab7d835ef35e0f5e2ece00b9b89
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/7064
Reviewed-by: CUE cueckoo <cueckoo@gmail.com>
Reviewed-by: Marcel van Lohuizen <mpvl@golang.org>
diff --git a/cmd/cue/cmd/help.go b/cmd/cue/cmd/help.go
index 51ab0c0..91f3ac3 100644
--- a/cmd/cue/cmd/help.go
+++ b/cmd/cue/cmd/help.go
@@ -262,9 +262,32 @@
 
 var injectHelp = &cobra.Command{
 	Use:   "injection",
-	Short: "inject values from the command line",
-	Long: `Many of the cue commands allow injecting values
-from the command line using the --inject/-t flag.
+	Short: "inject files or values into specific fields for a build",
+	Long: `Many of the cue commands allow injecting values or
+selecting files from the command line using the --inject/-t flag.
+
+
+Injecting files
+
+A "build" attribute defines a boolean expression that causes a file
+to only be included in a build if its expression evaluates to true.
+There may only be a single @if attribute per file and it must
+appear before a package clause.
+
+The expression is a subset of CUE consisting only of identifiers
+and the operators &&, ||, !, where identifiers refer to tags
+defined by the user on the command line.
+
+For example, the following file will only be included in a build
+if the user includes the flag "-t prod" on the command line.
+
+   // File prod.cue
+   @if(prod)
+
+   package foo
+
+
+Injecting values
 
 The injection mechanism allows values to be injected into fields
 that are marked with a "tag" attribute. For any field of the form
diff --git a/cue/load/config.go b/cue/load/config.go
index d32f080..88b6a75 100644
--- a/cue/load/config.go
+++ b/cue/load/config.go
@@ -163,6 +163,16 @@
 	// build and to inject values into the AST.
 	//
 	//
+	// File selection
+	//
+	// Files with an attribute of the form @if(expr) before a package clause
+	// are conditionally included if expr resolves to true, where expr refers to
+	// boolean values in Tags.
+	//
+	// It is an error for a file to have more than one @if attribute or to
+	// have a @if attribute without or after a package clause.
+	//
+	//
 	// Value injection
 	//
 	// The Tags values are also used to inject values into fields with a
@@ -490,7 +500,10 @@
 		}
 	}
 
-	c.loader = &loader{cfg: &c}
+	c.loader = &loader{
+		cfg:       &c,
+		buildTags: make(map[string]bool),
+	}
 
 	// TODO: also make this work if run from outside the module?
 	switch {
diff --git a/cue/load/import.go b/cue/load/import.go
index 6ee6084..bc286a3 100644
--- a/cue/load/import.go
+++ b/cue/load/import.go
@@ -375,6 +375,15 @@
 		return false // don't mark as added
 	}
 
+	if include, err := shouldBuildFile(pf, fp); !include {
+		if err != nil {
+			fp.err = errors.Append(fp.err, err)
+		}
+		p.IgnoredCUEFiles = append(p.InvalidCUEFiles, fullPath)
+		p.IgnoredFiles = append(p.InvalidFiles, file)
+		return false
+	}
+
 	if pkg != "" && pkg != "_" {
 		if p.PkgName == "" {
 			p.PkgName = pkg
diff --git a/cue/load/loader.go b/cue/load/loader.go
index da5a590..61fa06c 100644
--- a/cue/load/loader.go
+++ b/cue/load/loader.go
@@ -86,7 +86,7 @@
 
 	// 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.tags); err != nil {
+	if err := injectTags(c.Tags, l); err != nil {
 		for _, p := range a {
 			p.ReportError(err)
 		}
@@ -110,9 +110,10 @@
 )
 
 type loader struct {
-	cfg  *Config
-	stk  importStack
-	tags []tag // tags found in files
+	cfg       *Config
+	stk       importStack
+	tags      []tag // tags found in files
+	buildTags map[string]bool
 }
 
 func (l *loader) abs(filename string) string {
diff --git a/cue/load/loader_test.go b/cue/load/loader_test.go
index 073b45c..17177fe 100644
--- a/cue/load/loader_test.go
+++ b/cue/load/loader_test.go
@@ -253,6 +253,47 @@
 root:   $CWD/testdata
 dir:    $CWD/testdata/toolonly
 display:./toolonly`,
+	}, {
+		cfg: &Config{
+			Dir:  testdataDir,
+			Tags: []string{"prod"},
+		},
+		args: args("./tags"),
+		want: `
+path:   example.org/test/tags
+module: example.org/test
+root:   $CWD/testdata
+dir:    $CWD/testdata/tags
+display:./tags
+files:
+	$CWD/testdata/tags/prod.cue`,
+	}, {
+		cfg: &Config{
+			Dir:  testdataDir,
+			Tags: []string{"prod", "foo=bar"},
+		},
+		args: args("./tags"),
+		want: `
+path:   example.org/test/tags
+module: example.org/test
+root:   $CWD/testdata
+dir:    $CWD/testdata/tags
+display:./tags
+files:
+	$CWD/testdata/tags/prod.cue`,
+	}, {
+		cfg: &Config{
+			Dir:  testdataDir,
+			Tags: []string{"prod"},
+		},
+		args: args("./tagsbad"),
+		want: `
+err:    multiple @if attributes (and 2 more errors)
+path:   example.org/test/tagsbad
+module: example.org/test
+root:   $CWD/testdata
+dir:    $CWD/testdata/tagsbad
+display:./tagsbad`,
 	}}
 	for i, tc := range testCases {
 		t.Run(strconv.Itoa(i)+"/"+strings.Join(tc.args, ":"), func(t *testing.T) {
diff --git a/cue/load/tags.go b/cue/load/tags.go
index fad4d48..d30b502 100644
--- a/cue/load/tags.go
+++ b/cue/load/tags.go
@@ -21,6 +21,7 @@
 	"cuelang.org/go/cue/ast"
 	"cuelang.org/go/cue/build"
 	"cuelang.org/go/cue/errors"
+	"cuelang.org/go/cue/parser"
 	"cuelang.org/go/cue/token"
 	"cuelang.org/go/internal"
 	"cuelang.org/go/internal/cli"
@@ -138,13 +139,13 @@
 	return tags, errs
 }
 
-func injectTags(tags []string, a []tag) errors.Error {
+func injectTags(tags []string, l *loader) errors.Error {
 	// Parses command line args
 	for _, s := range tags {
 		p := strings.Index(s, "=")
-		found := false
+		found := l.buildTags[s]
 		if p > 0 { // key-value
-			for _, t := range a {
+			for _, t := range l.tags {
 				if t.key == s[:p] {
 					found = true
 					if err := t.inject(s[p+1:]); err != nil {
@@ -156,7 +157,7 @@
 				return errors.Newf(token.NoPos, "no tag for %q", s[:p])
 			}
 		} else { // shorthand
-			for _, t := range a {
+			for _, t := range l.tags {
 				for _, sh := range t.shorthands {
 					if sh == s {
 						found = true
@@ -167,9 +168,104 @@
 				}
 			}
 			if !found {
-				return errors.Newf(token.NoPos, "no shorthand for %q", s)
+				return errors.Newf(token.NoPos, "tag %q not used in any file", s)
 			}
 		}
 	}
 	return nil
 }
+
+func shouldBuildFile(f *ast.File, fp *fileProcessor) (bool, errors.Error) {
+	tags := fp.c.Tags
+
+	a, errs := getBuildAttr(f)
+	if errs != nil {
+		return false, errs
+	}
+	if a == nil {
+		return true, nil
+	}
+
+	_, body := a.Split()
+
+	expr, err := parser.ParseExpr("", body)
+	if err != nil {
+		return false, errors.Promote(err, "")
+	}
+
+	tagMap := map[string]bool{}
+	for _, t := range tags {
+		tagMap[t] = !strings.ContainsRune(t, '=')
+	}
+
+	c := checker{tags: tagMap, loader: fp.c.loader}
+	include := c.shouldInclude(expr)
+	if c.err != nil {
+		return false, c.err
+	}
+	return include, nil
+}
+
+func getBuildAttr(f *ast.File) (*ast.Attribute, errors.Error) {
+	var a *ast.Attribute
+	for _, d := range f.Decls {
+		switch x := d.(type) {
+		case *ast.Attribute:
+			key, _ := x.Split()
+			if key != "if" {
+				continue
+			}
+			if a != nil {
+				err := errors.Newf(d.Pos(), "multiple @if attributes")
+				err = errors.Append(err,
+					errors.Newf(a.Pos(), "previous declaration here"))
+				return nil, err
+			}
+			a = x
+
+		case *ast.Package:
+			break
+		}
+	}
+	return a, nil
+}
+
+type checker struct {
+	loader *loader
+	tags   map[string]bool
+	err    errors.Error
+}
+
+func (c *checker) shouldInclude(expr ast.Expr) bool {
+	switch x := expr.(type) {
+	case *ast.Ident:
+		c.loader.buildTags[x.Name] = true
+		return c.tags[x.Name]
+
+	case *ast.BinaryExpr:
+		switch x.Op {
+		case token.LAND:
+			return c.shouldInclude(x.X) && c.shouldInclude(x.Y)
+
+		case token.LOR:
+			return c.shouldInclude(x.X) || c.shouldInclude(x.Y)
+
+		default:
+			c.err = errors.Append(c.err, errors.Newf(token.NoPos,
+				"invalid operator %v", x.Op))
+			return false
+		}
+
+	case *ast.UnaryExpr:
+		if x.Op != token.NOT {
+			c.err = errors.Append(c.err, errors.Newf(token.NoPos,
+				"invalid operator %v", x.Op))
+		}
+		return !c.shouldInclude(x.X)
+
+	default:
+		c.err = errors.Append(c.err, errors.Newf(token.NoPos,
+			"invalid type %T in build attribute", expr))
+		return false
+	}
+}
diff --git a/cue/load/testdata/tags/prod.cue b/cue/load/testdata/tags/prod.cue
new file mode 100644
index 0000000..ce61d47
--- /dev/null
+++ b/cue/load/testdata/tags/prod.cue
@@ -0,0 +1,5 @@
+@if(prod)
+
+package tags
+
+foo: string @tag(foo)
diff --git a/cue/load/testdata/tags/stage.cue b/cue/load/testdata/tags/stage.cue
new file mode 100644
index 0000000..2d72a68
--- /dev/null
+++ b/cue/load/testdata/tags/stage.cue
@@ -0,0 +1,3 @@
+@if(stage)
+
+package tags
diff --git a/cue/load/testdata/tagsbad/prod.cue b/cue/load/testdata/tagsbad/prod.cue
new file mode 100644
index 0000000..f0598d9
--- /dev/null
+++ b/cue/load/testdata/tagsbad/prod.cue
@@ -0,0 +1,4 @@
+@if(foo)
+@if(bar)
+
+package tagsbad
diff --git a/cue/load/testdata/tagsbad/stage.cue b/cue/load/testdata/tagsbad/stage.cue
new file mode 100644
index 0000000..45af3ad
--- /dev/null
+++ b/cue/load/testdata/tagsbad/stage.cue
@@ -0,0 +1,3 @@
+package tagsbad
+
+@if(prod)