Support generated sources in query-based taze.

1. Don't report generated sources as missing sources to be removed.
2. Correctly report that ts_libraries with generated sources provide the generated sources as exports.

PiperOrigin-RevId: 230936754
diff --git a/ts_auto_deps/analyze/analyze.go b/ts_auto_deps/analyze/analyze.go
index 04a8f57..851509a 100644
--- a/ts_auto_deps/analyze/analyze.go
+++ b/ts_auto_deps/analyze/analyze.go
@@ -44,19 +44,29 @@
 
 // TargetLoader provides methods for loading targets from BUILD files.
 type TargetLoader interface {
-	// LoadLabels loads targets from BUILD files associated with labels.
-	// It returns a mapping from labels to targets or an error, if any
+	// LoadTargets loads targets from BUILD files associated with labels. A target
+	// is a rule, source file, generated file, package group or environment group.
+	// It returns a mapping from labels to targets or an error, if any occurred.
+	//
+	// A label must be the absolute label associated with a target. For example,
+	// '//foo/bar:baz' is acceptable whereas 'bar:baz' or '//foo/bar' will result
+	// in undefined behavior. TODO(lucassloan): make this an error
+	//
+	// Only returns targets visible to currentPkg. If currentPkg is an empty
+	// string returns all targets regardless of visibility.
+	LoadTargets(currentPkg string, labels []string) (map[string]*appb.Target, error)
+	// LoadRules loads rules from BUILD files associated with labels.
+	// It returns a mapping from labels to rules or an error, if any
 	// occurred.
 	//
-	// A label must be the absolute label associated with a target. For
+	// A label must be the absolute label associated with a rule. For
 	// example, '//foo/bar:baz' is acceptable whereas 'bar:baz' or '//foo/bar'
-	// will result in undefined behavior. If no target is found associated
-	// with a provided label, the label should be excluded from the returned
-	// mapping but an error should not be returned.
+	// will result in undefined behavior.
+	// TODO(lucassloan): make this an error.
 	//
 	// Only returns rules visible to currentPkg. If currentPkg is an empty string
-	// returns all targets regardless of visibility.
-	LoadLabels(currentPkg string, labels []string) (map[string]*appb.Rule, error)
+	// returns all rules regardless of visibility.
+	LoadRules(currentPkg string, labels []string) (map[string]*appb.Rule, error)
 	// LoadImportPaths loads targets from BUILD files associated with import
 	// paths relative to a root directory. It returns a mapping from import
 	// paths to targets or an error, if any occurred.
@@ -102,11 +112,11 @@
 	if err != nil {
 		return nil, err
 	}
-	targets, err := a.loader.LoadLabels(currentPkg, labels)
+	rules, err := a.loader.LoadRules(currentPkg, labels)
 	if err != nil {
 		return nil, err
 	}
-	resolved, err := a.resolveImportsForTargets(ctx, currentPkg, root, targets)
+	resolved, err := a.resolveImportsForTargets(ctx, currentPkg, root, rules)
 	if err != nil {
 		return nil, err
 	}
@@ -126,30 +136,91 @@
 	// missingSources are source files which could not be opened on disk.
 	// These are added to the dependency reports and MissingSources.
 	missingSources []string
+	// A map from the labels in the target's srcs to the Targets those
+	// labels refer.
+	sources map[string]*appb.Target
 }
 
-func (t *resolvedTarget) srcs() ([]string, error) {
-	srcs, err := sources(t.rule)
-	if err != nil {
-		// Targets without sources are considered errors.
-		return nil, err
+// setSources sets the sources on t.  It returns an error if one of the srcs of
+// t's rule isn't in loadedSrcs.
+func (t *resolvedTarget) setSources(loadedSrcs map[string]*appb.Target) error {
+	for _, label := range listAttribute(t.rule, "srcs") {
+		src := loadedSrcs[label]
+		if src == nil {
+			return fmt.Errorf("no source found for label %s", label)
+		}
+		t.sources[label] = src
 	}
+	return nil
+}
+
+// srcs returns the labels of the sources of t.
+func (t *resolvedTarget) srcs() ([]string, error) {
+	srcs := listAttribute(t.rule, "srcs")
+	if srcs == nil {
+		return nil, fmt.Errorf("target %q missing \"srcs\" attribute", t.label)
+	}
+
 	return srcs, nil
 }
 
+// literalSrcPaths returns the file paths of the non-generated sources of t.
+func (t *resolvedTarget) literalSrcPaths() ([]string, error) {
+	srcs := listAttribute(t.rule, "srcs")
+	if srcs == nil {
+		return nil, fmt.Errorf("target %q missing \"srcs\" attribute", t.label)
+	}
+	var literalFilePaths []string
+	for _, label := range listAttribute(t.rule, "srcs") {
+		src := t.sources[label]
+		if src == nil {
+			return nil, fmt.Errorf("src %q has no associated target", label)
+		}
+		// There's no syntactic way to determine if a label is a source file
+		// so check against the type of the relevant target
+		if src.GetType() == appb.Target_SOURCE_FILE {
+			literalFilePaths = append(literalFilePaths, labelToPath(label))
+		}
+	}
+	return literalFilePaths, nil
+}
+
+// getAllLiteralSrcPaths returns the file paths of all the non-generated sources
+// of the targets.
+func getAllLiteralSrcPaths(targets map[string]*resolvedTarget) ([]string, error) {
+	var allLiteralSrcPaths []string
+	for _, t := range targets {
+		literalSrcPaths, err := t.literalSrcPaths()
+		if err != nil {
+			return nil, err
+		}
+		allLiteralSrcPaths = append(allLiteralSrcPaths, literalSrcPaths...)
+	}
+
+	return allLiteralSrcPaths, nil
+}
+
 func (t *resolvedTarget) deps() []string {
 	return listAttribute(t.rule, "deps")
 }
 
 // provides returns whether the resolved target can provide the path provided.
 func (t *resolvedTarget) provides(path string) bool {
-	srcs, err := t.srcs()
-	if err != nil {
-		return false
-	}
-	for _, src := range srcs {
-		if src == path {
-			return true
+	for _, label := range listAttribute(t.rule, "srcs") {
+		src := t.sources[label]
+		if src.GetType() == appb.Target_SOURCE_FILE {
+			// For literal sources, check the path of the source
+			if labelToPath(label) == path {
+				return true
+			}
+		} else if src.GetType() == appb.Target_RULE {
+			// For generated souces, check against the paths of rule's
+			// outputs
+			for _, genSrc := range src.GetRule().GetRuleOutput() {
+				if labelToPath(genSrc) == path {
+					return true
+				}
+			}
 		}
 	}
 	return false
@@ -162,6 +233,7 @@
 		dependencies: make(map[string]*appb.Rule),
 		imports:      make(map[string][]*ts_auto_depsImport),
 		rule:         r,
+		sources:      make(map[string]*appb.Target),
 	}
 }
 
@@ -180,7 +252,7 @@
 		allDeps = append(allDeps, target.deps()...)
 		allSrcs = append(allSrcs, srcs...)
 	}
-	deps, err := a.loader.LoadLabels(currentPkg, allDeps)
+	deps, err := a.loader.LoadRules(currentPkg, allDeps)
 	if err != nil {
 		return nil, err
 	}
@@ -191,7 +263,25 @@
 			t.dependencies[dep] = deps[dep]
 		}
 	}
-	imports, errs := extractAllImports(root, allSrcs)
+	// load all the sources in the targets, so that literal and generated
+	// targets can be distinguished
+	srcs, err := a.loader.LoadTargets(currentPkg, allSrcs)
+	if err != nil {
+		return nil, err
+	}
+	for _, t := range targets {
+		err := t.setSources(srcs)
+		if err != nil {
+			return nil, err
+		}
+	}
+	// only extract the imports out of the literal sources, since ts_auto_deps can't
+	// see the contents of generated files
+	allLiteralSrcPaths, err := getAllLiteralSrcPaths(targets)
+	if err != nil {
+		return nil, err
+	}
+	imports, errs := extractAllImports(root, allLiteralSrcPaths)
 	for _, err := range errs {
 		// NotExist errors are caught and added to the generated dependency
 		// reports as missing source files. Only errors which are not NotExist
@@ -201,7 +291,7 @@
 		}
 	}
 	for _, t := range targets {
-		srcs, err := t.srcs()
+		srcs, err := t.literalSrcPaths()
 		if err != nil {
 			return nil, err
 		}
@@ -328,19 +418,9 @@
 	return target.GetName()
 }
 
-// sources creates an array of all sources listed in the 'srcs' attribute
-// on each target in targets.
-func sources(target *appb.Rule) ([]string, error) {
-	srcs := listAttribute(target, "srcs")
-	if srcs == nil {
-		return nil, fmt.Errorf("target %q missing \"srcs\" attribute", target.GetName())
-	}
-	for i, src := range srcs {
-		_, pkg, file := edit.ParseLabel(src)
-		// TODO(jdhamlik): Handle generated files.
-		srcs[i] = platform.Normalize(filepath.Clean(filepath.Join(pkg, file)))
-	}
-	return srcs, nil
+func labelToPath(label string) string {
+	_, pkg, file := edit.ParseLabel(label)
+	return platform.Normalize(filepath.Clean(filepath.Join(pkg, file)))
 }
 
 // generateReports generates reports for each label in labels.
@@ -419,7 +499,7 @@
 			unusedDeps = append(unusedDeps, dep)
 		}
 	}
-	labelToRule, err := a.loader.LoadLabels("", unusedDeps)
+	labelToRule, err := a.loader.LoadRules("", unusedDeps)
 	if err != nil {
 		return nil, err
 	}
diff --git a/ts_auto_deps/analyze/analyze_test.go b/ts_auto_deps/analyze/analyze_test.go
index 26aac4d..700ea84 100644
--- a/ts_auto_deps/analyze/analyze_test.go
+++ b/ts_auto_deps/analyze/analyze_test.go
@@ -6,6 +6,8 @@
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -44,8 +46,20 @@
 	}
 }
 
-func (bl *fakeTargetLoader) LoadLabels(_ string, labels []string) (map[string]*appb.Rule, error) {
-	return bl.loadTargets(bl.targetsByLabels, labels)
+func (bl *fakeTargetLoader) LoadRules(_ string, labels []string) (map[string]*appb.Rule, error) {
+	return bl.loadRules(bl.targetsByLabels, labels)
+}
+
+func (bl *fakeTargetLoader) LoadTargets(_ string, labels []string) (map[string]*appb.Target, error) {
+	targets := make(map[string]*appb.Target)
+	for _, l := range labels {
+		if strings.Contains(l, ".") {
+			targets[l] = &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()}
+		} else {
+			targets[l] = &appb.Target{Type: appb.Target_RULE.Enum()}
+		}
+	}
+	return targets, nil
 }
 
 func (bl *fakeTargetLoader) byLabel(label, value string) {
@@ -53,14 +67,14 @@
 }
 
 func (bl *fakeTargetLoader) LoadImportPaths(_ context.Context, _, _ string, paths []string) (map[string]*appb.Rule, error) {
-	return bl.loadTargets(bl.targetsByImportPaths, paths)
+	return bl.loadRules(bl.targetsByImportPaths, paths)
 }
 
 func (bl *fakeTargetLoader) byImportPath(importPath, value string) {
 	bl.targetsByImportPaths[importPath] = value
 }
 
-func (bl *fakeTargetLoader) loadTargets(source map[string]string, keys []string) (map[string]*appb.Rule, error) {
+func (bl *fakeTargetLoader) loadRules(source map[string]string, keys []string) (map[string]*appb.Rule, error) {
 	targets := make(map[string]*appb.Rule)
 	for _, key := range keys {
 		value, ok := source[key]
@@ -479,6 +493,208 @@
 	}
 }
 
+func createResolvedTarget(srcs []string) *resolvedTarget {
+	return &resolvedTarget{
+		rule: &appb.Rule{
+			Attribute: []*appb.Attribute{
+				&appb.Attribute{
+					Name:            proto.String("srcs"),
+					Type:            appb.Attribute_STRING_LIST.Enum(),
+					StringListValue: srcs,
+				},
+			},
+		},
+		sources: map[string]*appb.Target{
+			"//a:file.ts":   &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+			"//b:file.ts":   &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+			"//b:generator": &appb.Target{Type: appb.Target_RULE.Enum()},
+			"//b:wiz":       &appb.Target{Type: appb.Target_RULE.Enum()},
+		},
+	}
+}
+
+func TestLiteralSrcPaths(t *testing.T) {
+	tests := []struct {
+		name     string
+		srcs     []string
+		err      error
+		expected []string
+	}{
+		{
+			"OneLiteralSource",
+			[]string{"//a:file.ts"},
+			nil,
+			[]string{"a/file.ts"},
+		},
+		{
+			"MultipleLiteralSources",
+			[]string{"//a:file.ts", "//b:file.ts"},
+			nil,
+			[]string{"a/file.ts", "b/file.ts"},
+		},
+		{
+			"MultipleGeneratedSources",
+			[]string{"//b:generator", "//b:wiz"},
+			nil,
+			nil,
+		},
+		{
+			"MixedSources",
+			[]string{"//a:file.ts", "//b:file.ts", "//b:generator", "//b:wiz"},
+			nil,
+			[]string{"a/file.ts", "b/file.ts"},
+		},
+		{
+			"MissingSource",
+			[]string{"//not/in/the/set/of/resolved:sources"},
+			fmt.Errorf("src %q has no associated target", "//not/in/the/set/of/resolved:sources"),
+			nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			rt := createResolvedTarget(test.srcs)
+			literalSrcPaths, err := rt.literalSrcPaths()
+			if !reflect.DeepEqual(err, test.err) {
+				t.Errorf("got err %q, expected %q", err, test.err)
+			}
+
+			if diff := pretty.Compare(literalSrcPaths, test.expected); diff != "" {
+				t.Errorf("failed to get correct literal source paths: (-got, +want)\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestGetAllLiteralSrcPaths(t *testing.T) {
+	tests := []struct {
+		name      string
+		srcsLists [][]string
+		err       error
+		expected  []string
+	}{
+		{
+			"OneTarget",
+			[][]string{
+				[]string{"//a:file.ts", "//b:file.ts"},
+			},
+			nil,
+			[]string{"a/file.ts", "b/file.ts"},
+		},
+		{
+			"MultipleTargets",
+			[][]string{
+				[]string{"//a:file.ts"},
+				[]string{"//b:file.ts"},
+			},
+			nil,
+			[]string{"a/file.ts", "b/file.ts"},
+		},
+		{
+			"MissingSource",
+			[][]string{
+				[]string{"//not/in/the/set/of/resolved:sources"},
+			},
+			fmt.Errorf("src %q has no associated target", "//not/in/the/set/of/resolved:sources"),
+			nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			rts := make(map[string]*resolvedTarget)
+			for i, srcs := range test.srcsLists {
+				rts[strconv.Itoa(i)] = createResolvedTarget(srcs)
+			}
+			literalSrcPaths, err := getAllLiteralSrcPaths(rts)
+			if !reflect.DeepEqual(err, test.err) {
+				t.Errorf("got err %q, expected %q", err, test.err)
+			}
+
+			if diff := pretty.Compare(literalSrcPaths, test.expected); diff != "" {
+				t.Errorf("failed to get correct literal source paths: (-got, +want)\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestSetSources(t *testing.T) {
+	tests := []struct {
+		name       string
+		srcs       []string
+		loadedSrcs map[string]*appb.Target
+		err        error
+		expected   map[string]*appb.Target
+	}{
+		{
+			"NoSources",
+			nil,
+			nil,
+			nil,
+			nil,
+		},
+		{
+			"OneSource",
+			[]string{"//a:file.ts"},
+			map[string]*appb.Target{
+				"//a:file.ts": &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+			},
+			nil,
+			map[string]*appb.Target{
+				"//a:file.ts": &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+			},
+		},
+		{
+			"ExtraSources",
+			[]string{"//a:file.ts"},
+			map[string]*appb.Target{
+				"//a:file.ts":   &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+				"//b:file.ts":   &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+				"//b:generator": &appb.Target{Type: appb.Target_RULE.Enum()},
+				"//b:wiz":       &appb.Target{Type: appb.Target_RULE.Enum()},
+			},
+			nil,
+			map[string]*appb.Target{
+				"//a:file.ts": &appb.Target{Type: appb.Target_SOURCE_FILE.Enum()},
+			},
+		},
+		{
+			"MissingSources",
+			[]string{"//a:file.ts"},
+			nil,
+			fmt.Errorf("no source found for label %s", "//a:file.ts"),
+			nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			rt := &resolvedTarget{
+				rule: &appb.Rule{
+					Attribute: []*appb.Attribute{
+						&appb.Attribute{
+							Name:            proto.String("srcs"),
+							Type:            appb.Attribute_STRING_LIST.Enum(),
+							StringListValue: test.srcs,
+						},
+					},
+				},
+				sources: make(map[string]*appb.Target),
+			}
+
+			err := rt.setSources(test.loadedSrcs)
+			if !reflect.DeepEqual(err, test.err) {
+				t.Errorf("got err %q, expected %q", err, test.err)
+			}
+
+			if diff := pretty.Compare(rt.sources, test.expected); diff != "" {
+				t.Errorf("failed to set correct sources: (-got, +want)\n%s", diff)
+			}
+		})
+	}
+}
+
 func missingDeps(report *arpb.DependencyReport) []string {
 	var deps []string
 	for _, group := range report.GetMissingDependencyGroup() {
diff --git a/ts_auto_deps/analyze/loader.go b/ts_auto_deps/analyze/loader.go
index 0308267..51cfd44 100644
--- a/ts_auto_deps/analyze/loader.go
+++ b/ts_auto_deps/analyze/loader.go
@@ -42,8 +42,8 @@
 	// analyzed in the "-recursive" case, these caches will be garbage
 	// collected between directories.
 	pkgCache map[string]*pkgCacheEntry
-	// labelCache is a mapping from a label to its loaded rule.
-	labelCache map[string]*appb.Rule
+	// labelCache is a mapping from a label to its loaded target.
+	labelCache map[string]*appb.Target
 
 	// queryCount is the total number of queries executed by the target loader.
 	queryCount int
@@ -57,13 +57,33 @@
 		bazelBinary: bazelBinary,
 
 		pkgCache:   make(map[string]*pkgCacheEntry),
-		labelCache: make(map[string]*appb.Rule),
+		labelCache: make(map[string]*appb.Target),
 	}
 }
 
-// LoadLabels uses Bazel query to load targets associated with labels from BUILD
+// LoadRules uses Bazel query to load rules associated with labels from BUILD
 // files.
-func (q *QueryBasedTargetLoader) LoadLabels(pkg string, labels []string) (map[string]*appb.Rule, error) {
+func (q *QueryBasedTargetLoader) LoadRules(pkg string, labels []string) (map[string]*appb.Rule, error) {
+	labelToTarget, err := q.LoadTargets(pkg, labels)
+	if err != nil {
+		return nil, err
+	}
+
+	labelToRule := make(map[string]*appb.Rule)
+	for _, label := range labels {
+		target := labelToTarget[label]
+		if target.GetType() == appb.Target_RULE {
+			labelToRule[label] = target.GetRule()
+		} else {
+			return nil, fmt.Errorf("target contains object of type %q instead of type %q", target.GetType(), appb.Target_RULE)
+		}
+	}
+	return labelToRule, nil
+}
+
+// LoadTargets uses Bazel query to load targets associated with labels from BUILD
+// files.
+func (q *QueryBasedTargetLoader) LoadTargets(pkg string, labels []string) (map[string]*appb.Target, error) {
 	var labelCacheMisses []string
 	for _, label := range labels {
 		if _, ok := q.labelCache[labelCacheKey(pkg, label)]; !ok {
@@ -84,11 +104,11 @@
 			return nil, err
 		}
 		for _, target := range r.GetTarget() {
-			label, err := q.ruleLabel(target)
+			label, err := q.targetLabel(target)
 			if err != nil {
 				return nil, err
 			}
-			q.labelCache[labelCacheKey(pkg, label)] = target.GetRule()
+			q.labelCache[labelCacheKey(pkg, label)] = target
 		}
 		for _, label := range labelCacheMisses {
 			key := labelCacheKey(pkg, label)
@@ -101,11 +121,11 @@
 			}
 		}
 	}
-	labelToRule := make(map[string]*appb.Rule)
+	labelToTarget := make(map[string]*appb.Target)
 	for _, label := range labels {
-		labelToRule[label] = q.labelCache[labelCacheKey(pkg, label)]
+		labelToTarget[label] = q.labelCache[labelCacheKey(pkg, label)]
 	}
-	return labelToRule, nil
+	return labelToTarget, nil
 }
 
 func labelCacheKey(currentPkg, label string) string {
@@ -168,7 +188,7 @@
 
 	labelToRule := make(map[string]*appb.Rule)
 	for len(generators) > 0 {
-		generatorToRule, err := q.LoadLabels(currentPkg, generators)
+		generatorToRule, err := q.LoadRules(currentPkg, generators)
 		if err != nil {
 			return nil, err
 		}
@@ -218,6 +238,8 @@
 	return results, nil
 }
 
+// ruleLabel returns the label for a target which is a rule.  Returns an error if
+// target is not a rule.
 func (q *QueryBasedTargetLoader) ruleLabel(target *appb.Target) (string, error) {
 	if t := target.GetType(); t != appb.Target_RULE {
 		return "", fmt.Errorf("target contains object of type %q instead of type %q", t, appb.Target_RULE)
@@ -225,6 +247,8 @@
 	return target.GetRule().GetName(), nil
 }
 
+// fileLabel returns the label for a target which is a file.  Returns an error if
+// target is not a source file or a generated file.
 func (q *QueryBasedTargetLoader) fileLabel(target *appb.Target) (string, error) {
 	switch t := target.GetType(); t {
 	case appb.Target_GENERATED_FILE:
@@ -236,6 +260,25 @@
 	}
 }
 
+// targetLabel returns the label for a target.  Returns an error if target is an
+// unknown type.
+func (q *QueryBasedTargetLoader) targetLabel(target *appb.Target) (string, error) {
+	switch t := target.GetType(); t {
+	case appb.Target_GENERATED_FILE:
+		return target.GetGeneratedFile().GetName(), nil
+	case appb.Target_SOURCE_FILE:
+		return target.GetSourceFile().GetName(), nil
+	case appb.Target_RULE:
+		return target.GetRule().GetName(), nil
+	case appb.Target_PACKAGE_GROUP:
+		return target.GetPackageGroup().GetName(), nil
+	case appb.Target_ENVIRONMENT_GROUP:
+		return target.GetEnvironmentGroup().GetName(), nil
+	default:
+		return "", fmt.Errorf("target contains object of unknown type %q", t)
+	}
+}
+
 // loadRuleIncludingSourceFiles loads all rules which include labels in
 // sourceFileLabels, Returns a map from source file label to the rule which
 // includes it.