<commit_before>package analyze

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/bazelbuild/buildtools/edit"
	"github.com/bazelbuild/rules_typescript/ts_auto_deps/workspace"
	"github.com/golang/protobuf/proto"

	appb "github.com/bazelbuild/buildtools/build_proto"
)

// QueryBasedTargetLoader uses Bazel query to load targets from BUILD files.
type QueryBasedTargetLoader struct {
	workdir     string
	bazelBinary string
}

// NewQueryBasedTargetLoader constructs a new QueryBasedTargetLoader rooted
// in workdir.
func NewQueryBasedTargetLoader(workdir, bazelBinary string) *QueryBasedTargetLoader {
	return &QueryBasedTargetLoader{
		workdir:     workdir,
		bazelBinary: bazelBinary,
	}
}

// LoadLabels uses Bazel query to load targets associated with labels from BUILD
// files.
func (q *QueryBasedTargetLoader) LoadLabels(labels []string) (map[string]*appb.Rule, error) {
	// Ensure the labels are unique to minimize the total number of targets that
	// need to be loaded.
	r, err := q.batchQuery(dedupeLabels(labels))
	if err != nil {
		return nil, err
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, target := range r.GetTarget() {
		label, err := q.ruleLabel(target)
		if err != nil {
			return nil, err
		}
		labelToRule[label] = target.GetRule()
	}
	return labelToRule, nil
}

// LoadImportPaths uses Bazel Query to load targets associated with import
// paths from BUILD files.
func (q *QueryBasedTargetLoader) LoadImportPaths(ctx context.Context, workspaceRoot string, paths []string) (map[string]*appb.Rule, error) {
	var remainingImportPaths []string
	results := make(map[string]*appb.Rule)
	for _, path := range paths {
		if trim := strings.TrimPrefix(path, workspace.Name()+"/"); trim != path {
			// TODO(jdhamlik): Optimize by grouping the queries into one larger query.
			// TODO(jdhamlik): Handle .d.ts and .tsx files.
			r, err := q.query(trim + ".ts")
			if err != nil {
				return nil, err
			}
			targets := r.GetTarget()
			// Expecting to get one, and only one, target per query.
			n := len(targets)
			if n < 1 {
				return nil, fmt.Errorf("failed to resolved a target for file %q", trim+".ts")
			}
			if n > 1 {
				return nil, fmt.Errorf("got %d targets when only one was expected", n)
			}
			t, err := q.loadRuleIncludingFile(targets[0].GetSourceFile().GetName())
			if err != nil {
				return nil, err
			}
			results[path] = t
		} else if trim := strings.TrimPrefix(path, "goog:"); trim != path {
			// There are no current heuristics in this implementation for
			// 'goog:' imports. That is handled by the ts_auto_deps binary proper.
			results[path] = nil
		} else {
			// The import is not explicitly under google3 or a closure-style
			// import.
			remainingImportPaths = append(remainingImportPaths, path)
		}
	}

	// Attempt to locate the file rooted in the workspace even though it isn't
	// prefixed by 'google3/'.
	for _, imp := range remainingImportPaths {
		// If the path has a suffix of ".ngfactory" or ".ngsummary", it might
		// be an Angular AOT generated file. We can infer the target as we
		// infer its corresponding ngmodule target by simply stripping the
		// ".ngfactory" / ".ngsummary" suffix
		path := strings.TrimSuffix(strings.TrimSuffix(imp, ".ngsummary"), ".ngfactory")
		res, err := q.batchQuery(pathWithExtensions(path))
		if err != nil {
			return nil, err
		}
		if len(res.GetTarget()) > 0 {
			target := res.GetTarget()[0]
			label, err := q.sourceFileLabel(target)
			if err != nil {
				return nil, err
			}
			rule, err := q.loadRuleIncludingFile(label)
			if err != nil {
				return nil, err
			}
			results[imp] = rule
			continue
		}

	}

	return results, nil
}

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)
	}
	return target.GetRule().GetName(), nil
}

func (q *QueryBasedTargetLoader) sourceFileLabel(target *appb.Target) (string, error) {
	if t := target.GetType(); t != appb.Target_SOURCE_FILE {
		return "", fmt.Errorf("target contains object of type %q instead of type %q", t, appb.Target_SOURCE_FILE)
	}
	return target.GetSourceFile().GetName(), nil
}

// loadRuleIncludingFile loads the target associated with a file label.
func (q *QueryBasedTargetLoader) loadRuleIncludingFile(fileLabel string) (*appb.Rule, error) {
	_, pkg, file := edit.ParseLabel(fileLabel)
	// Filter the targets in the file label's package to only targets which
	// include the file in their 'srcs' attribute.
	r, err := q.query(fmt.Sprintf("attr('srcs', %s, //%s:*)", file, pkg))
	if err != nil {
		return nil, err
	}
	for _, target := range r.GetTarget() {
		rule := target.GetRule()
		for _, src := range listAttribute(rule, "srcs") {
			_, _, path := edit.ParseLabel(src)
			// Return the first rule which has a source file exactly matching
			// the requested file path.
			if path == file {
				return rule, nil
			}
		}
	}
	return nil, fmt.Errorf("failed to resolved a target for file label %q", fileLabel)
}

// batchQuery runs a set of queries with a single call to Bazel query and the
// '--keep_going' flag.
func (q *QueryBasedTargetLoader) batchQuery(queries []string) (*appb.QueryResult, error) {
	// Join all of the queries with a '+' character according to Bazel's
	// syntax for running multiple queries.
	return q.query("--keep_going", strings.Join(queries, "+"))
}

func (q *QueryBasedTargetLoader) query(args ...string) (*appb.QueryResult, error) {
	n := len(args)
	if n < 1 {
		return nil, fmt.Errorf("expected at least one argument")
	}
	if query := args[n-1]; query == "" {
		// An empty query was provided so return an empty result without
		// making a call to Bazel.
		return &appb.QueryResult{}, nil
	}
	var stdout, stderr bytes.Buffer
	args = append([]string{"query", "--output=proto"}, args...)
	cmd := exec.Command(q.bazelBinary, args...)
	cmd.Dir = q.workdir
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr<commit_msg>Refactor LoadImportPaths to batch queries<commit_after>package analyze

import (
	"bytes"
	"context"
	"fmt"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/bazelbuild/buildtools/edit"
	"github.com/bazelbuild/rules_typescript/ts_auto_deps/workspace"
	"github.com/golang/protobuf/proto"

	appb "github.com/bazelbuild/buildtools/build_proto"
)

// QueryBasedTargetLoader uses Bazel query to load targets from BUILD files.
type QueryBasedTargetLoader struct {
	workdir     string
	bazelBinary string
}

// NewQueryBasedTargetLoader constructs a new QueryBasedTargetLoader rooted
// in workdir.
func NewQueryBasedTargetLoader(workdir, bazelBinary string) *QueryBasedTargetLoader {
	return &QueryBasedTargetLoader{
		workdir:     workdir,
		bazelBinary: bazelBinary,
	}
}

// LoadLabels uses Bazel query to load targets associated with labels from BUILD
// files.
func (q *QueryBasedTargetLoader) LoadLabels(labels []string) (map[string]*appb.Rule, error) {
	// Ensure the labels are unique to minimize the total number of targets that
	// need to be loaded.
	r, err := q.batchQuery(dedupeLabels(labels))
	if err != nil {
		return nil, err
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, target := range r.GetTarget() {
		label, err := q.ruleLabel(target)
		if err != nil {
			return nil, err
		}
		labelToRule[label] = target.GetRule()
	}
	return labelToRule, nil
}

// LoadImportPaths uses Bazel Query to load targets associated with import
// paths from BUILD files.
func (q *QueryBasedTargetLoader) LoadImportPaths(ctx context.Context, workspaceRoot string, paths []string) (map[string]*appb.Rule, error) {
	results := make(map[string]*appb.Rule)

	addedPaths := make(map[string]bool)
	var possiblePaths []string
	for _, path := range paths {
		if strings.HasPrefix(path, "goog:") {
			// 'goog:' imports are resolved using an sstable.
			results[path] = nil
			continue
		}
		if !strings.HasPrefix(path, "@") {
			if _, ok := addedPaths[path]; !ok {
				addedPaths[path] = true

				// If the path has a suffix of ".ngfactory" or ".ngsummary", it might
				// be an Angular AOT generated file. We can infer the target as we
				// infer its corresponding ngmodule target by simply stripping the
				// ".ngfactory" / ".ngsummary" suffix
				path = strings.TrimSuffix(strings.TrimSuffix(path, ".ngsummary"), ".ngfactory")
				path = strings.TrimPrefix(path, workspace.Name()+"/")

				possiblePaths = append(possiblePaths, pathWithExtensions(path)...)
				possiblePaths = append(possiblePaths, pathWithExtensions(filepath.Join(path, "index"))...)
			}
		}
	}

	r, err := q.batchQuery(possiblePaths)
	if err != nil {
		return nil, err
	}
	var sourceFileLabels []string
	for _, target := range r.GetTarget() {
		label, err := q.sourceFileLabel(target)
		if err != nil {
			return nil, err
		}
		sourceFileLabels = append(sourceFileLabels, label)
	}
	sourceLabelToRule, err := q.loadRulesIncludingSourceFiles(workspaceRoot, sourceFileLabels)
	if err != nil {
		return nil, err
	}
	for sourceLabel, rule := range sourceLabelToRule {
		_, pkg, file := edit.ParseLabel(sourceLabel)
		pathWithoutExtension := strings.TrimSuffix(filepath.Join(pkg, stripTSExtension(file)), string(filepath.Separator)+"index")
		for _, path := range paths {
			if pathWithoutExtension == path {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngsummary") {
				results[path] = rule
			} else if pathWithoutExtension == strings.TrimSuffix(path, ".ngfactory") {
				results[path] = rule
			}
		}
	}

	return results, nil
}

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)
	}
	return target.GetRule().GetName(), nil
}

func (q *QueryBasedTargetLoader) sourceFileLabel(target *appb.Target) (string, error) {
	if t := target.GetType(); t != appb.Target_SOURCE_FILE {
		return "", fmt.Errorf("target contains object of type %q instead of type %q", t, appb.Target_SOURCE_FILE)
	}
	return target.GetSourceFile().GetName(), nil
}

// loadRuleIncludingSourceFiles loads all rules which include labels in
// sourceFileLabels, Returns a map from source file label to the rule which
// includes it.
func (q *QueryBasedTargetLoader) loadRulesIncludingSourceFiles(workspaceRoot string, sourceFileLabels []string) (map[string]*appb.Rule, error) {
	pkgToLabels := make(map[string][]string)
	queries := make([]string, 0, len(sourceFileLabels))
	for _, label := range sourceFileLabels {
		_, pkg, file := edit.ParseLabel(label)
		pkgToLabels[pkg] = append(pkgToLabels[pkg], label)
		// Query for all targets in the package which use file.
		queries = append(queries, fmt.Sprintf("attr('srcs', %s, //%s:*)", file, pkg))
	}
	r, err := q.batchQuery(queries)
	if err != nil {
		return nil, err
	}
	labelToRule := make(map[string]*appb.Rule)
	for _, target := range r.GetTarget() {
		label, err := q.ruleLabel(target)
		if err != nil {
			return nil, err
		}
		rule := target.GetRule()
		_, pkg, _ := edit.ParseLabel(label)
		labels := pkgToLabels[pkg]
		for _, src := range listAttribute(rule, "srcs") {
			for _, l := range labels {
				if src == l {
					labelToRule[l] = rule
					break
				}
			}
		}
	}
	return labelToRule, nil
}

// batchQuery runs a set of queries with a single call to Bazel query and the
// '--keep_going' flag.
func (q *QueryBasedTargetLoader) batchQuery(queries []string) (*appb.QueryResult, error) {
	// Join all of the queries with a '+' character according to Bazel's
	// syntax for running multiple queries.
	return q.query("--keep_going", strings.Join(queries, "+"))
}

func (q *QueryBasedTargetLoader) query(args ...string) (*appb.QueryResult, error) {
	n := len(args)
	if n < 1 {
		return nil, fmt.Errorf("expected at least one argument")
	}
	if query := args[n-1]; query == "" {
		// An empty query was provided so return an empty result without
		// making a call to Bazel.
		return &appb.QueryResult{}, nil
	}
	var stdout, stderr bytes.Buffer
	args = append([]string{"query", "--output=proto"}, args...)
	cmd := exec.Command(q.bazelBinary, args...)
	cmd.Dir = q.workdir
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
{'fraction_non_alphanumeric': 0.12786057496538805, 'fraction_numerical': 0.0009772782799902272, 'mean_word_length': 3.898284802552852, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7462834', 'n_tokens_mistral': 4241, 'n_tokens_neox': 3915, 'n_words': 1205}