Merge "Deterministic aquery details in mixed builds" am: 58c43e9694

Original change: https://android-review.googlesource.com/c/platform/build/soong/+/2097673 Change-Id: I0b47451d66a970de114628385ca856d77d02a78f Signed-off-by: Automerger Merge Worker <android-build-automerger-merge-worker@system.gserviceaccount.com>
2022-05-13 18:28:30 +00:00
parent ec2519b0b9 58c43e9694
commit 1bf08694c8
3 changed files with 224 additions and 147 deletions
--- a/android/bazel_handler.go
+++ b/android/bazel_handler.go
@@ -825,14 +825,14 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {

 	for _, depset := range ctx.Config().BazelContext.AqueryDepsets() {
 		var outputs []Path
-		for _, depsetDepId := range depset.TransitiveDepSetIds {
-			otherDepsetName := bazelDepsetName(depsetDepId)
+		for _, depsetDepHash := range depset.TransitiveDepSetHashes {
+			otherDepsetName := bazelDepsetName(depsetDepHash)
 			outputs = append(outputs, PathForPhony(ctx, otherDepsetName))
 		}
 		for _, artifactPath := range depset.DirectArtifacts {
 			outputs = append(outputs, PathForBazelOut(ctx, artifactPath))
 		}
-		thisDepsetName := bazelDepsetName(depset.Id)
+		thisDepsetName := bazelDepsetName(depset.ContentHash)
 		ctx.Build(pctx, BuildParams{
 			Rule:      blueprint.Phony,
 			Outputs:   []WritablePath{PathForPhony(ctx, thisDepsetName)},
@@ -874,8 +874,8 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 		for _, inputPath := range buildStatement.InputPaths {
 			cmd.Implicit(PathForBazelOut(ctx, inputPath))
 		}
-		for _, inputDepsetId := range buildStatement.InputDepsetIds {
-			otherDepsetName := bazelDepsetName(inputDepsetId)
+		for _, inputDepsetHash := range buildStatement.InputDepsetHashes {
+			otherDepsetName := bazelDepsetName(inputDepsetHash)
 			cmd.Implicit(PathForPhony(ctx, otherDepsetName))
 		}

@@ -924,6 +924,6 @@ func GetConfigKey(ctx ModuleContext) configKey {
 	}
 }

-func bazelDepsetName(depsetId int) string {
-	return fmt.Sprintf("bazel_depset_%d", depsetId)
+func bazelDepsetName(contentHash string) string {
+	return fmt.Sprintf("bazel_depset_%s", contentHash)
 }
--- a/bazel/aquery.go
+++ b/bazel/aquery.go
@@ -15,10 +15,13 @@
 package bazel

 import (
+	"crypto/sha256"
 	"encoding/json"
 	"fmt"
 	"path/filepath"
+	"reflect"
 	"regexp"
+	"sort"
 	"strings"

 	"github.com/google/blueprint/proptools"
@@ -44,15 +47,16 @@ type KeyValuePair struct {
 }

 // AqueryDepset is a depset definition from Bazel's aquery response. This is
-// akin to the `depSetOfFiles` in the response proto, except that direct
-// artifacts are enumerated by full path instead of by ID.
+// akin to the `depSetOfFiles` in the response proto, except:
+//   * direct artifacts are enumerated by full path instead of by ID
+//   * has a hash of the depset contents, instead of an int ID (for determinism)
 // A depset is a data structure for efficient transitive handling of artifact
 // paths. A single depset consists of one or more artifact paths and one or
 // more "child" depsets.
 type AqueryDepset struct {
-	Id                  int
-	DirectArtifacts     []string
-	TransitiveDepSetIds []int
+	ContentHash            string
+	DirectArtifacts        []string
+	TransitiveDepSetHashes []string
 }

 // depSetOfFiles contains relevant portions of Bazel's aquery proto, DepSetOfFiles.
@@ -99,19 +103,24 @@ type BuildStatement struct {
 	// input paths. There should be no overlap between these fields; an input
 	// path should either be included as part of an unexpanded depset or a raw
 	// input path string, but not both.
-	InputDepsetIds []int
-	InputPaths     []string
+	InputDepsetHashes []string
+	InputPaths        []string
 }

 // A helper type for aquery processing which facilitates retrieval of path IDs from their
 // less readable Bazel structures (depset and path fragment).
 type aqueryArtifactHandler struct {
-	// Maps depset Id to depset struct.
-	depsetIdToDepset map[int]depSetOfFiles
+	// Maps depset id to AqueryDepset, a representation of depset which is
+	// post-processed for middleman artifact handling, unhandled artifact
+	// dropping, content hashing, etc.
+	depsetIdToAqueryDepset map[int]AqueryDepset
+	// Maps content hash to AqueryDepset.
+	depsetHashToAqueryDepset map[string]AqueryDepset
+
 	// depsetIdToArtifactIdsCache is a memoization of depset flattening, because flattening
 	// may be an expensive operation.
-	depsetIdToArtifactIdsCache map[int][]int
-	// Maps artifact Id to fully expanded path.
+	depsetHashToArtifactPathsCache map[string][]string
+	// Maps artifact ContentHash to fully expanded path.
 	artifactIdToPath map[int]string
 }

@@ -127,7 +136,7 @@ var TemplateActionOverriddenTokens = map[string]string{
 var manifestFilePattern = regexp.MustCompile(".*/.+\\.runfiles/MANIFEST$")

 // The file name of py3wrapper.sh, which is used by py_binary targets.
-var py3wrapperFileName = "/py3wrapper.sh"
+const py3wrapperFileName = "/py3wrapper.sh"

 func newAqueryHandler(aqueryResult actionGraphContainer) (*aqueryArtifactHandler, error) {
 	pathFragments := map[int]pathFragment{}
@@ -144,7 +153,7 @@ func newAqueryHandler(aqueryResult actionGraphContainer) (*aqueryArtifactHandler
 		artifactIdToPath[artifact.Id] = artifactPath
 	}

-	// Map middleman artifact Id to input artifact depset ID.
+	// Map middleman artifact ContentHash to input artifact depset ID.
 	// Middleman artifacts are treated as "substitute" artifacts for mixed builds. For example,
 	// if we find a middleman action which has outputs [foo, bar], and output [baz_middleman], then,
 	// for each other action which has input [baz_middleman], we add [foo, bar] to the inputs for
@@ -165,50 +174,84 @@ func newAqueryHandler(aqueryResult actionGraphContainer) (*aqueryArtifactHandler
 	}

 	depsetIdToDepset := map[int]depSetOfFiles{}
-	// Validate and adjust aqueryResult.DepSetOfFiles values.
 	for _, depset := range aqueryResult.DepSetOfFiles {
-		filteredArtifactIds := []int{}
-		for _, artifactId := range depset.DirectArtifactIds {
-			path, pathExists := artifactIdToPath[artifactId]
-			if !pathExists {
-				return nil, fmt.Errorf("undefined input artifactId %d", artifactId)
-			}
-			// Filter out any inputs which are universally dropped, and swap middleman
-			// artifacts with their corresponding depsets.
-			if depsetsToUse, isMiddleman := middlemanIdToDepsetIds[artifactId]; isMiddleman {
-				// Swap middleman artifacts with their corresponding depsets and drop the middleman artifacts.
-				depset.TransitiveDepSetIds = append(depset.TransitiveDepSetIds, depsetsToUse...)
-			} else if strings.HasSuffix(path, py3wrapperFileName) || manifestFilePattern.MatchString(path) {
-				// Drop these artifacts.
-				// See go/python-binary-host-mixed-build for more details.
-				// 1) For py3wrapper.sh, there is no action for creating py3wrapper.sh in the aquery output of
-				// Bazel py_binary targets, so there is no Ninja build statements generated for creating it.
-				// 2) For MANIFEST file, SourceSymlinkManifest action is in aquery output of Bazel py_binary targets,
-				// but it doesn't contain sufficient information so no Ninja build statements are generated
-				// for creating it.
-				// So in mixed build mode, when these two are used as input of some Ninja build statement,
-				// since there is no build statement to create them, they should be removed from input paths.
-				// TODO(b/197135294): Clean up this custom runfiles handling logic when
-				// SourceSymlinkManifest and SymlinkTree actions are supported.
-			} else {
-				// TODO(b/216194240): Filter out bazel tools.
-				filteredArtifactIds = append(filteredArtifactIds, artifactId)
-			}
-		}
-		depset.DirectArtifactIds = filteredArtifactIds
-		for _, childDepsetId := range depset.TransitiveDepSetIds {
-			if _, exists := depsetIds[childDepsetId]; !exists {
-				return nil, fmt.Errorf("undefined input depsetId %d (referenced by depsetId %d)", childDepsetId, depset.Id)
-			}
-		}
 		depsetIdToDepset[depset.Id] = depset
 	}

-	return &aqueryArtifactHandler{
-		depsetIdToDepset:           depsetIdToDepset,
-		depsetIdToArtifactIdsCache: map[int][]int{},
-		artifactIdToPath:           artifactIdToPath,
-	}, nil
+	aqueryHandler := aqueryArtifactHandler{
+		depsetIdToAqueryDepset:         map[int]AqueryDepset{},
+		depsetHashToAqueryDepset:       map[string]AqueryDepset{},
+		depsetHashToArtifactPathsCache: map[string][]string{},
+		artifactIdToPath:               artifactIdToPath,
+	}
+
+	// Validate and adjust aqueryResult.DepSetOfFiles values.
+	for _, depset := range aqueryResult.DepSetOfFiles {
+		_, err := aqueryHandler.populateDepsetMaps(depset, middlemanIdToDepsetIds, depsetIdToDepset)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return &aqueryHandler, nil
+}
+
+// Ensures that the handler's depsetIdToAqueryDepset map contains an entry for the given
+// depset.
+func (a *aqueryArtifactHandler) populateDepsetMaps(depset depSetOfFiles, middlemanIdToDepsetIds map[int][]int, depsetIdToDepset map[int]depSetOfFiles) (AqueryDepset, error) {
+	if aqueryDepset, containsDepset := a.depsetIdToAqueryDepset[depset.Id]; containsDepset {
+		return aqueryDepset, nil
+	}
+	transitiveDepsetIds := depset.TransitiveDepSetIds
+	directArtifactPaths := []string{}
+	for _, artifactId := range depset.DirectArtifactIds {
+		path, pathExists := a.artifactIdToPath[artifactId]
+		if !pathExists {
+			return AqueryDepset{}, fmt.Errorf("undefined input artifactId %d", artifactId)
+		}
+		// Filter out any inputs which are universally dropped, and swap middleman
+		// artifacts with their corresponding depsets.
+		if depsetsToUse, isMiddleman := middlemanIdToDepsetIds[artifactId]; isMiddleman {
+			// Swap middleman artifacts with their corresponding depsets and drop the middleman artifacts.
+			transitiveDepsetIds = append(transitiveDepsetIds, depsetsToUse...)
+		} else if strings.HasSuffix(path, py3wrapperFileName) || manifestFilePattern.MatchString(path) {
+			// Drop these artifacts.
+			// See go/python-binary-host-mixed-build for more details.
+			// 1) For py3wrapper.sh, there is no action for creating py3wrapper.sh in the aquery output of
+			// Bazel py_binary targets, so there is no Ninja build statements generated for creating it.
+			// 2) For MANIFEST file, SourceSymlinkManifest action is in aquery output of Bazel py_binary targets,
+			// but it doesn't contain sufficient information so no Ninja build statements are generated
+			// for creating it.
+			// So in mixed build mode, when these two are used as input of some Ninja build statement,
+			// since there is no build statement to create them, they should be removed from input paths.
+			// TODO(b/197135294): Clean up this custom runfiles handling logic when
+			// SourceSymlinkManifest and SymlinkTree actions are supported.
+		} else {
+			// TODO(b/216194240): Filter out bazel tools.
+			directArtifactPaths = append(directArtifactPaths, path)
+		}
+	}
+
+	childDepsetHashes := []string{}
+	for _, childDepsetId := range transitiveDepsetIds {
+		childDepset, exists := depsetIdToDepset[childDepsetId]
+		if !exists {
+			return AqueryDepset{}, fmt.Errorf("undefined input depsetId %d (referenced by depsetId %d)", childDepsetId, depset.Id)
+		}
+		childAqueryDepset, err := a.populateDepsetMaps(childDepset, middlemanIdToDepsetIds, depsetIdToDepset)
+		if err != nil {
+			return AqueryDepset{}, err
+		}
+		childDepsetHashes = append(childDepsetHashes, childAqueryDepset.ContentHash)
+	}
+	aqueryDepset := AqueryDepset{
+		ContentHash:            depsetContentHash(directArtifactPaths, childDepsetHashes),
+		DirectArtifacts:        directArtifactPaths,
+		TransitiveDepSetHashes: childDepsetHashes,
+	}
+	a.depsetIdToAqueryDepset[depset.Id] = aqueryDepset
+	a.depsetHashToAqueryDepset[aqueryDepset.ContentHash] = aqueryDepset
+	return aqueryDepset, nil
 }

 // getInputPaths flattens the depsets of the given IDs and returns all transitive
@@ -219,15 +262,12 @@ func (a *aqueryArtifactHandler) getInputPaths(depsetIds []int) ([]string, error)
 	inputPaths := []string{}

 	for _, inputDepSetId := range depsetIds {
-		inputArtifacts, err := a.artifactIdsFromDepsetId(inputDepSetId)
+		depset := a.depsetIdToAqueryDepset[inputDepSetId]
+		inputArtifacts, err := a.artifactPathsFromDepsetHash(depset.ContentHash)
 		if err != nil {
 			return nil, err
 		}
-		for _, inputId := range inputArtifacts {
-			inputPath, exists := a.artifactIdToPath[inputId]
-			if !exists {
-				return nil, fmt.Errorf("undefined input artifactId %d", inputId)
-			}
+		for _, inputPath := range inputArtifacts {
 			inputPaths = append(inputPaths, inputPath)
 		}
 	}
@@ -235,23 +275,23 @@ func (a *aqueryArtifactHandler) getInputPaths(depsetIds []int) ([]string, error)
 	return inputPaths, nil
 }

-func (a *aqueryArtifactHandler) artifactIdsFromDepsetId(depsetId int) ([]int, error) {
-	if result, exists := a.depsetIdToArtifactIdsCache[depsetId]; exists {
+func (a *aqueryArtifactHandler) artifactPathsFromDepsetHash(depsetHash string) ([]string, error) {
+	if result, exists := a.depsetHashToArtifactPathsCache[depsetHash]; exists {
 		return result, nil
 	}
-	if depset, exists := a.depsetIdToDepset[depsetId]; exists {
-		result := depset.DirectArtifactIds
-		for _, childId := range depset.TransitiveDepSetIds {
-			childArtifactIds, err := a.artifactIdsFromDepsetId(childId)
+	if depset, exists := a.depsetHashToAqueryDepset[depsetHash]; exists {
+		result := depset.DirectArtifacts
+		for _, childHash := range depset.TransitiveDepSetHashes {
+			childArtifactIds, err := a.artifactPathsFromDepsetHash(childHash)
 			if err != nil {
 				return nil, err
 			}
 			result = append(result, childArtifactIds...)
 		}
-		a.depsetIdToArtifactIdsCache[depsetId] = result
+		a.depsetHashToArtifactPathsCache[depsetHash] = result
 		return result, nil
 	} else {
-		return nil, fmt.Errorf("undefined input depsetId %d", depsetId)
+		return nil, fmt.Errorf("undefined input depset hash %d", depsetHash)
 	}
 }

@@ -261,9 +301,6 @@ func (a *aqueryArtifactHandler) artifactIdsFromDepsetId(depsetId int) ([]int, er
 // BuildStatements are one-to-one with actions in the given action graph, and AqueryDepsets
 // are one-to-one with Bazel's depSetOfFiles objects.
 func AqueryBuildStatements(aqueryJsonProto []byte) ([]BuildStatement, []AqueryDepset, error) {
-	buildStatements := []BuildStatement{}
-	depsets := []AqueryDepset{}
-
 	var aqueryResult actionGraphContainer
 	err := json.Unmarshal(aqueryJsonProto, &aqueryResult)
 	if err != nil {
@@ -274,6 +311,8 @@ func AqueryBuildStatements(aqueryJsonProto []byte) ([]BuildStatement, []AqueryDe
 		return nil, nil, err
 	}

+	var buildStatements []BuildStatement
+
 	for _, actionEntry := range aqueryResult.Actions {
 		if shouldSkipAction(actionEntry) {
 			continue
@@ -298,38 +337,75 @@ func AqueryBuildStatements(aqueryJsonProto []byte) ([]BuildStatement, []AqueryDe
 		buildStatements = append(buildStatements, buildStatement)
 	}

-	// Iterate over depset IDs in the initial aquery order to preserve determinism.
-	for _, depset := range aqueryResult.DepSetOfFiles {
-		// Use the depset in the aqueryHandler, as this contains the augmented depsets.
-		depset = aqueryHandler.depsetIdToDepset[depset.Id]
-		directPaths := []string{}
-		for _, artifactId := range depset.DirectArtifactIds {
-			pathString := aqueryHandler.artifactIdToPath[artifactId]
-			directPaths = append(directPaths, pathString)
+	depsetsByHash := map[string]AqueryDepset{}
+	depsets := []AqueryDepset{}
+	for _, aqueryDepset := range aqueryHandler.depsetIdToAqueryDepset {
+		if prevEntry, hasKey := depsetsByHash[aqueryDepset.ContentHash]; hasKey {
+			// Two depsets collide on hash. Ensure that their contents are identical.
+			if !reflect.DeepEqual(aqueryDepset, prevEntry) {
+				return nil, nil, fmt.Errorf("Two different depsets have the same hash: %v, %v", prevEntry, aqueryDepset)
+			}
+		} else {
+			depsetsByHash[aqueryDepset.ContentHash] = aqueryDepset
+			depsets = append(depsets, aqueryDepset)
 		}
-		aqueryDepset := AqueryDepset{
-			Id:                  depset.Id,
-			DirectArtifacts:     directPaths,
-			TransitiveDepSetIds: depset.TransitiveDepSetIds,
-		}
-		depsets = append(depsets, aqueryDepset)
 	}
+
+	// Build Statements and depsets must be sorted by their content hash to
+	// preserve determinism between builds (this will result in consistent ninja file
+	// output). Note they are not sorted by their original IDs nor their Bazel ordering,
+	// as Bazel gives nondeterministic ordering / identifiers in aquery responses.
+	sort.Slice(buildStatements, func(i, j int) bool {
+		// For build statements, compare output lists. In Bazel, each output file
+		// may only have one action which generates it, so this will provide
+		// a deterministic ordering.
+		outputs_i := buildStatements[i].OutputPaths
+		outputs_j := buildStatements[j].OutputPaths
+		if len(outputs_i) != len(outputs_j) {
+			return len(outputs_i) < len(outputs_j)
+		}
+		if len(outputs_i) == 0 {
+			// No outputs for these actions, so compare commands.
+			return buildStatements[i].Command < buildStatements[j].Command
+		}
+		// There may be multiple outputs, but the output ordering is deterministic.
+		return outputs_i[0] < outputs_j[0]
+	})
+	sort.Slice(depsets, func(i, j int) bool {
+		return depsets[i].ContentHash < depsets[j].ContentHash
+	})
 	return buildStatements, depsets, nil
 }

-func (aqueryHandler *aqueryArtifactHandler) validateInputDepsets(inputDepsetIds []int) ([]int, error) {
-	// Validate input depsets correspond to real depsets.
+// depsetContentHash computes and returns a SHA256 checksum of the contents of
+// the given depset. This content hash may serve as the depset's identifier.
+// Using a content hash for an identifier is superior for determinism. (For example,
+// using an integer identifier which depends on the order in which the depsets are
+// created would result in nondeterministic depset IDs.)
+func depsetContentHash(directPaths []string, transitiveDepsetHashes []string) string {
+	h := sha256.New()
+	// Use newline as delimiter, as paths cannot contain newline.
+	h.Write([]byte(strings.Join(directPaths, "\n")))
+	h.Write([]byte(strings.Join(transitiveDepsetHashes, "\n")))
+	fullHash := fmt.Sprintf("%016x", h.Sum(nil))
+	return fullHash
+}
+
+func (aqueryHandler *aqueryArtifactHandler) depsetContentHashes(inputDepsetIds []int) ([]string, error) {
+	hashes := []string{}
 	for _, depsetId := range inputDepsetIds {
-		if _, exists := aqueryHandler.depsetIdToDepset[depsetId]; !exists {
+		if aqueryDepset, exists := aqueryHandler.depsetIdToAqueryDepset[depsetId]; !exists {
 			return nil, fmt.Errorf("undefined input depsetId %d", depsetId)
+		} else {
+			hashes = append(hashes, aqueryDepset.ContentHash)
 		}
 	}
-	return inputDepsetIds, nil
+	return hashes, nil
 }

 func (aqueryHandler *aqueryArtifactHandler) normalActionBuildStatement(actionEntry action) (BuildStatement, error) {
 	command := strings.Join(proptools.ShellEscapeListIncludingSpaces(actionEntry.Arguments), " ")
-	inputDepsetIds, err := aqueryHandler.validateInputDepsets(actionEntry.InputDepSetIds)
+	inputDepsetHashes, err := aqueryHandler.depsetContentHashes(actionEntry.InputDepSetIds)
 	if err != nil {
 		return BuildStatement{}, err
 	}
@@ -339,12 +415,12 @@ func (aqueryHandler *aqueryArtifactHandler) normalActionBuildStatement(actionEnt
 	}

 	buildStatement := BuildStatement{
-		Command:        command,
-		Depfile:        depfile,
-		OutputPaths:    outputPaths,
-		InputDepsetIds: inputDepsetIds,
-		Env:            actionEntry.EnvironmentVariables,
-		Mnemonic:       actionEntry.Mnemonic,
+		Command:           command,
+		Depfile:           depfile,
+		OutputPaths:       outputPaths,
+		InputDepsetHashes: inputDepsetHashes,
+		Env:               actionEntry.EnvironmentVariables,
+		Mnemonic:          actionEntry.Mnemonic,
 	}
 	return buildStatement, nil
 }
@@ -413,18 +489,18 @@ func (aqueryHandler *aqueryArtifactHandler) templateExpandActionBuildStatement(a
 	// See go/python-binary-host-mixed-build for more details.
 	command := fmt.Sprintf(`/bin/bash -c 'echo "%[1]s" | sed "s/\\\\n/\\n/g" > %[2]s && chmod a+x %[2]s'`,
 		escapeCommandlineArgument(expandedTemplateContent), outputPaths[0])
-	inputDepsetIds, err := aqueryHandler.validateInputDepsets(actionEntry.InputDepSetIds)
+	inputDepsetHashes, err := aqueryHandler.depsetContentHashes(actionEntry.InputDepSetIds)
 	if err != nil {
 		return BuildStatement{}, err
 	}

 	buildStatement := BuildStatement{
-		Command:        command,
-		Depfile:        depfile,
-		OutputPaths:    outputPaths,
-		InputDepsetIds: inputDepsetIds,
-		Env:            actionEntry.EnvironmentVariables,
-		Mnemonic:       actionEntry.Mnemonic,
+		Command:           command,
+		Depfile:           depfile,
+		OutputPaths:       outputPaths,
+		InputDepsetHashes: inputDepsetHashes,
+		Env:               actionEntry.EnvironmentVariables,
+		Mnemonic:          actionEntry.Mnemonic,
 	}
 	return buildStatement, nil
 }
--- a/bazel/aquery_test.go
+++ b/bazel/aquery_test.go
@@ -234,7 +234,6 @@ func TestAqueryMultiArchGenrule(t *testing.T) {
 				OutputPaths: []string{
 					fmt.Sprintf("bazel-out/sourceroot/k8-fastbuild/bin/bionic/libc/syscalls-%s.S", arch),
 				},
-				InputDepsetIds: []int{1},
 				Env: []KeyValuePair{
 					KeyValuePair{Key: "PATH", Value: "/bin:/usr/bin:/usr/local/bin"},
 				},
@@ -248,9 +247,12 @@ func TestAqueryMultiArchGenrule(t *testing.T) {
 		"../sourceroot/bionic/libc/tools/gensyscalls.py",
 		"../bazel_tools/tools/genrule/genrule-setup.sh",
 	}
-	actualFlattenedInputs := flattenDepsets([]int{1}, actualDepsets)
-	if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
-		t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
+	// In this example, each depset should have the same expected inputs.
+	for _, actualDepset := range actualDepsets {
+		actualFlattenedInputs := flattenDepsets([]string{actualDepset.ContentHash}, actualDepsets)
+		if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
+			t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
+		}
 	}
 }

@@ -746,10 +748,9 @@ func TestTransitiveInputDepsets(t *testing.T) {

 	expectedBuildStatements := []BuildStatement{
 		BuildStatement{
-			Command:        "/bin/bash -c 'touch bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out'",
-			OutputPaths:    []string{"bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out"},
-			InputDepsetIds: []int{1},
-			Mnemonic:       "Action",
+			Command:     "/bin/bash -c 'touch bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out'",
+			OutputPaths: []string{"bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out"},
+			Mnemonic:    "Action",
 		},
 	}
 	assertBuildStatements(t, expectedBuildStatements, actualbuildStatements)
@@ -763,7 +764,8 @@ func TestTransitiveInputDepsets(t *testing.T) {
 	}
 	expectedFlattenedInputs = append(expectedFlattenedInputs, "bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_root")

-	actualFlattenedInputs := flattenDepsets([]int{1}, actualDepsets)
+	actualDepsetHashes := actualbuildStatements[0].InputDepsetHashes
+	actualFlattenedInputs := flattenDepsets(actualDepsetHashes, actualDepsets)
 	if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
 		t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
 	}
@@ -844,38 +846,24 @@ func TestMiddlemenAction(t *testing.T) {
 		t.Fatalf("Expected %d build statements, got %d", expected, len(actualBuildStatements))
 	}

+	expectedDepsetFiles := [][]string{
+		[]string{"middleinput_one", "middleinput_two"},
+		[]string{"middleinput_one", "middleinput_two", "maininput_one", "maininput_two"},
+	}
+	assertFlattenedDepsets(t, actualDepsets, expectedDepsetFiles)
+
 	bs := actualBuildStatements[0]
 	if len(bs.InputPaths) > 0 {
 		t.Errorf("Expected main action raw inputs to be empty, but got %q", bs.InputPaths)
 	}

-	expectedInputDepsets := []int{2}
-	if !reflect.DeepEqual(bs.InputDepsetIds, expectedInputDepsets) {
-		t.Errorf("Expected main action depset IDs %v, but got %v", expectedInputDepsets, bs.InputDepsetIds)
-	}
-
 	expectedOutputs := []string{"output"}
 	if !reflect.DeepEqual(bs.OutputPaths, expectedOutputs) {
 		t.Errorf("Expected main action outputs %q, but got %q", expectedOutputs, bs.OutputPaths)
 	}

-	expectedAllDepsets := []AqueryDepset{
-		{
-			Id:              1,
-			DirectArtifacts: []string{"middleinput_one", "middleinput_two"},
-		},
-		{
-			Id:                  2,
-			DirectArtifacts:     []string{"maininput_one", "maininput_two"},
-			TransitiveDepSetIds: []int{1},
-		},
-	}
-	if !reflect.DeepEqual(actualDepsets, expectedAllDepsets) {
-		t.Errorf("Expected depsets %v, but got %v", expectedAllDepsets, actualDepsets)
-	}
-
 	expectedFlattenedInputs := []string{"middleinput_one", "middleinput_two", "maininput_one", "maininput_two"}
-	actualFlattenedInputs := flattenDepsets(bs.InputDepsetIds, actualDepsets)
+	actualFlattenedInputs := flattenDepsets(bs.InputDepsetHashes, actualDepsets)

 	if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
 		t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
@@ -883,29 +871,42 @@ func TestMiddlemenAction(t *testing.T) {
 }

 // Returns the contents of given depsets in concatenated post order.
-func flattenDepsets(depsetIdsToFlatten []int, allDepsets []AqueryDepset) []string {
-	depsetsById := map[int]AqueryDepset{}
+func flattenDepsets(depsetHashesToFlatten []string, allDepsets []AqueryDepset) []string {
+	depsetsByHash := map[string]AqueryDepset{}
 	for _, depset := range allDepsets {
-		depsetsById[depset.Id] = depset
+		depsetsByHash[depset.ContentHash] = depset
 	}
 	result := []string{}
-	for _, depsetId := range depsetIdsToFlatten {
-		result = append(result, flattenDepset(depsetId, depsetsById)...)
+	for _, depsetId := range depsetHashesToFlatten {
+		result = append(result, flattenDepset(depsetId, depsetsByHash)...)
 	}
 	return result
 }

 // Returns the contents of a given depset in post order.
-func flattenDepset(depsetIdToFlatten int, allDepsets map[int]AqueryDepset) []string {
-	depset := allDepsets[depsetIdToFlatten]
+func flattenDepset(depsetHashToFlatten string, allDepsets map[string]AqueryDepset) []string {
+	depset := allDepsets[depsetHashToFlatten]
 	result := []string{}
-	for _, depsetId := range depset.TransitiveDepSetIds {
+	for _, depsetId := range depset.TransitiveDepSetHashes {
 		result = append(result, flattenDepset(depsetId, allDepsets)...)
 	}
 	result = append(result, depset.DirectArtifacts...)
 	return result
 }

+func assertFlattenedDepsets(t *testing.T, actualDepsets []AqueryDepset, expectedDepsetFiles [][]string) {
+	t.Helper()
+	if len(actualDepsets) != len(expectedDepsetFiles) {
+		t.Errorf("Expected %d depsets, but got %d depsets", expectedDepsetFiles, actualDepsets)
+	}
+	for i, actualDepset := range actualDepsets {
+		actualFlattenedInputs := flattenDepsets([]string{actualDepset.ContentHash}, actualDepsets)
+		if !reflect.DeepEqual(actualFlattenedInputs, expectedDepsetFiles[i]) {
+			t.Errorf("Expected depset files: %v, but got %v", expectedDepsetFiles[i], actualFlattenedInputs)
+		}
+	}
+}
+
 func TestSimpleSymlink(t *testing.T) {
 	const inputString = `
 {