Export: enhance trimming using build statement metadata

src/BUILD.plz

@@ -10,8 +10,8 @@ go_binary(
     deps = [
         "///third_party/go/github.com_thought-machine_go-flags//:go-flags",
         "///third_party/go/go.uber.org_automaxprocs//maxprocs",
-        "//src/audit",
         "//src/assets",
+        "//src/audit",
         "//src/build",
         "//src/cache",
         "//src/clean",
@@ -29,7 +29,6 @@ go_binary(
         "//src/help",
         "//src/metrics",
         "//src/output",
-        "//src/parse",
         "//src/plz",
         "//src/plzinit",
         "//src/process",

src/core/package.go

@@ -34,23 +34,46 @@ type Package struct {
 	targets map[string]*BuildTarget
 	// Set of output files from rules.
 	Outputs map[string]*BuildTarget
+	// Includes metadata from parsing the package BUILD file.
+	Metadata PackageMetadata
 	// Protects access to above
 	mutex sync.RWMutex
 }
 
-// NewPackage constructs a new package with the given name.
-func NewPackage(name string) *Package {
-	return NewPackageSubrepo(name, "")
+// PackageOptions is a functional option type for configuring a new Package.
+type PackageOptions func(*Package)
+
+// WithPackageSubrepo returns a PackageOptions that sets the subrepo name for a new Package.
+func WithPackageSubrepo(name string) PackageOptions {
+	return func(p *Package) {
+		p.SubrepoName = name
+	}
 }
 
-// NewPackageSubrepo constructs a new package with the given name and subrepo.
-func NewPackageSubrepo(name, subrepo string) *Package {
-	return &Package{
-		Name:        name,
-		SubrepoName: subrepo,
-		targets:     map[string]*BuildTarget{},
-		Outputs:     map[string]*BuildTarget{},
+// WithPackageMetadata returns a PackageOptions that enables tracking of
+// metadata (like statement positions and subinclude mappings) for the Package.
+// This is required for features like 'plz export'.
+func WithPackageMetadata() PackageOptions {
+	return func(p *Package) {
+		p.Metadata = newPackageMetadata()
+	}
+}
+
+// NewPackage constructs a new package with the given name, and enables additional features
+// given the PackageOptions provided.
+func NewPackage(name string, options ...PackageOptions) *Package {
+	pkg := &Package{
+		Name:    name,
+		targets: map[string]*BuildTarget{},
+		Outputs: map[string]*BuildTarget{},
+		// Defaults to noop to avoid storing metadata for most operations
+		Metadata: newNoopPackageMetadata(),
+	}
+
+	for _, option := range options {
+		option(pkg)
 	}
+	return pkg
 }
 
 // Target returns the target with the given name, or nil if this package doesn't have one.

src/core/package_metadata.go

@@ -0,0 +1,183 @@
+package core
+
+import (
+	"slices"
+	"sync"
+)
+
+// BuildStatement represents the start and end byte positions of a parsed statement in a BUILD file.
+type BuildStatement struct {
+	Start, End int
+}
+
+// Len returns the byte length of the build statement.
+func (bs *BuildStatement) Len() int64 {
+	return int64(bs.End - bs.Start)
+}
+
+// StartPos returns the starting byte position of the build statement.
+func (bs *BuildStatement) StartPos() int64 {
+	return int64(bs.Start)
+}
+
+// BuildStatements is a slice of BuildStatement that implements sort.Interface.
+type BuildStatements []BuildStatement
+
+func (s BuildStatements) Len() int           { return len(s) }
+func (s BuildStatements) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
+func (s BuildStatements) Less(i, j int) bool { return s[i].StartPos() < s[j].StartPos() }
+
+// BuildStatementProvider defines a closure that generates new build statements.
+// It is used as an argument in PackageMetadata methods to defer evaluation, avoiding
+// unnecessary computation when using the no-op implementation.
+type BuildStatementProvider func() BuildStatement
+
+// SubincludesLabelProvider defines a closure that generates labels for a subinclude statement.
+// It is used as an argument in PackageMetadata methods to defer evaluation, avoiding
+// unnecessary computation when using the no-op implementation.
+type SubincludesLabelProvider func() BuildLabels
+
+// PackageMetadata stores metadata about parsed BUILD files, mapping statements and subincludes
+// to their respective targets. This supports additional logic for operations such as `plz export`
+// but should be disabled for most operations by using the no-op implementation to avoid the overhead.
+type PackageMetadata interface {
+	// RegisterStatementTarget records that the given build target was created as a result of the
+	// given statement being executed. This should only be called for statements in BUILD files.
+	RegisterStatementTarget(target *BuildTarget, stmtProvider BuildStatementProvider)
+	// RegisterRequiredSubinclude records that the given build target requires the given subinclude
+	// labels to be built. This is used to track which subinclude statements contributed to a target's
+	// definition. This should only be called for statements in BUILD files.
+	RegisterRequiredSubinclude(target *BuildTarget, labelProvider SubincludesLabelProvider)
+	// RegisterSubincludeStatement records that the given subinclude statement (provided by stmtProvider)
+	// includes the given build label. This should only be called for statements in BUILD files.
+	RegisterSubincludeStatement(label BuildLabel, stmtProvider BuildStatementProvider)
+	// FindStatement returns the build statement that was responsible for generating the given target.
+	// Returns nil if the target was not found in the recorded metadata.
+	FindStatement(target *BuildTarget) *BuildStatement
+	// FindTargets returns all build targets that were generated by the given build statement.
+	// Returns an empty slice if no targets were found for the given statement.
+	FindTargets(stmt *BuildStatement) []*BuildTarget
+	// FindRequiredSubincludes returns all subinclude labels that were required by the given target.
+	// The return value is empty if no subinclude information was found for the target.
+	FindRequiredSubincludes(target *BuildTarget) BuildLabels
+	// GetSubincludedLabels returns all build labels that were included by the given subinclude statement.
+	// Returns the labels or an empty slice if the statement wasn't found.
+	GetSubincludedLabels(stmt *BuildStatement) BuildLabels
+}
+
+// packageMetadataImpl is the default implementation of the PackageMetadata interface.
+// It uses in-memory maps to track the relationships between BUILD file statements,
+// subincludes, and the build targets they define.
+type packageMetadataImpl struct {
+	// StmtToTarget maps each build statement (identified by its byte range in a BUILD file)
+	// to the targets it produced. Since a single statement (like a custom target or loop)
+	// can produce multiple targets, this is a one-to-many mapping.
+	StmtToTarget map[BuildStatement][]*BuildTarget
+	// TargetToSubinclude tracks the subinclude labels that were required for a target's
+	// definition. This allows mapping a target back to the subincluded labels required for building
+	// the target.
+	TargetToSubinclude map[*BuildTarget]BuildLabels
+	// LabelsPerSubincludeStmt maps a subinclude statement (identified by its position
+	// in the BUILD file) to the labels it explicitly subincludes.
+	LabelsPerSubincludeStmt map[BuildStatement]BuildLabels
+
+	// mutex protects concurrent access to the metadata maps during the parallel
+	// parsing of BUILD files.
+	mutex sync.RWMutex
+}
+
+func newPackageMetadata() PackageMetadata {
+	return &packageMetadataImpl{
+		StmtToTarget:            map[BuildStatement][]*BuildTarget{},
+		TargetToSubinclude:      map[*BuildTarget]BuildLabels{},
+		LabelsPerSubincludeStmt: map[BuildStatement]BuildLabels{},
+	}
+}
+
+func (m *packageMetadataImpl) RegisterStatementTarget(target *BuildTarget, stmtProvider BuildStatementProvider) {
+	stmt := stmtProvider()
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.StmtToTarget[stmt] = append(m.StmtToTarget[stmt], target)
+}
+
+func (m *packageMetadataImpl) RegisterRequiredSubinclude(target *BuildTarget, labelProvider SubincludesLabelProvider) {
+	labels := labelProvider()
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.TargetToSubinclude[target] = append(m.TargetToSubinclude[target], labels...)
+}
+
+func (m *packageMetadataImpl) RegisterSubincludeStatement(label BuildLabel, stmtProvider BuildStatementProvider) {
+	stmt := stmtProvider()
+
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.LabelsPerSubincludeStmt[stmt] = append(m.LabelsPerSubincludeStmt[stmt], label)
+}
+
+func (m *packageMetadataImpl) FindStatement(target *BuildTarget) *BuildStatement {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	for stmt, targets := range m.StmtToTarget {
+		if slices.Contains(targets, target) {
+			return &stmt
+		}
+	}
+	return nil
+}
+
+func (m *packageMetadataImpl) FindTargets(stmt *BuildStatement) []*BuildTarget {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return m.StmtToTarget[*stmt]
+}
+
+func (m *packageMetadataImpl) FindRequiredSubincludes(target *BuildTarget) BuildLabels {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return m.TargetToSubinclude[target]
+}
+
+func (m *packageMetadataImpl) GetSubincludedLabels(stmt *BuildStatement) BuildLabels {
+	m.mutex.RLock()
+	defer m.mutex.RUnlock()
+
+	return m.LabelsPerSubincludeStmt[*stmt]
+}
+
+// noopPackageMetadata implements the PackageMetadata interface with no-op methods. This is used to
+// avoid the overhead of parsing metadata for operations that don't depend on it.
+type noopPackageMetadata struct{}
+
+func newNoopPackageMetadata() PackageMetadata {
+	return &noopPackageMetadata{}
+}
+
+func (n *noopPackageMetadata) RegisterStatementTarget(target *BuildTarget, stmtProvider BuildStatementProvider) {
+}
+func (n *noopPackageMetadata) RegisterRequiredSubinclude(target *BuildTarget, labelProvider SubincludesLabelProvider) {
+}
+func (n *noopPackageMetadata) RegisterSubincludeStatement(label BuildLabel, stmtProvider BuildStatementProvider) {
+}
+func (n *noopPackageMetadata) FindStatement(target *BuildTarget) *BuildStatement {
+	log.Fatalf("Metadata not tracked, using no-op implementation.")
+	return nil
+}
+func (n *noopPackageMetadata) FindTargets(stmt *BuildStatement) []*BuildTarget {
+	log.Fatalf("Metadata not tracked, using no-op implementation.")
+	return nil
+}
+func (n *noopPackageMetadata) FindRequiredSubincludes(target *BuildTarget) BuildLabels {
+	log.Fatalf("Metadata not tracked, using no-op implementation.")
+	return nil
+}
+func (n *noopPackageMetadata) GetSubincludedLabels(stmt *BuildStatement) BuildLabels {
+	log.Fatalf("Metadata not tracked, using no-op implementation.")
+	return nil
+}

src/core/state.go

@@ -240,6 +240,8 @@ type BuildState struct {
 	// NeedDebugDeps is true if we're doing a `plz debug` and we need to build the debug tools and
 	// data
 	NeedDebugDeps bool
+	// ParseMetadata is true if we want to store build file metadata
+	ParseMetadata bool
 
 	// initOnce is used to control loading the subrepo .plzconfig
 	initOnce *sync.Once

src/export/BUILD

@@ -4,10 +4,23 @@ go_library(
     pgo_file = "//:pgo",
     visibility = ["PUBLIC"],
     deps = [
+        "///third_party/go/github.com_please-build_buildtools//build",
         "//src/cli/logging",
         "//src/core",
         "//src/fs",
-        "//src/gc",
         "//src/parse",
+        "//src/parse/asp",
+    ],
+)
+
+go_test(
+    name = "export_test",
+    srcs = ["export_test.go"],
+    data = ["test_data"],
+    deps = [
+        ":export",
+        "///third_party/go/github.com_stretchr_testify//assert",
+        "//src/core",
+        "//src/parse/asp",
     ],
 )