Merge "Preserve depset structure from bazel aquery"

android/bazel_handler.go

@@ -28,6 +28,7 @@ import (
 
 	"android/soong/bazel/cquery"
 	"android/soong/shared"
+	"github.com/google/blueprint"
 
 	"android/soong/bazel"
 )
@@ -101,6 +102,9 @@ type BazelContext interface {
 
 	// Returns build statements which should get registered to reflect Bazel's outputs.
 	BuildStatementsToRegister() []bazel.BuildStatement
+
+	// Returns the depsets defined in Bazel's aquery response.
+	AqueryDepsets() []bazel.AqueryDepset
 }
 
 type bazelRunner interface {
@@ -128,6 +132,9 @@ type bazelContext struct {
 
 	// Build statements which should get registered to reflect Bazel's outputs.
 	buildStatements []bazel.BuildStatement
+
+	// Depsets which should be used for Bazel's build statements.
+	depsets []bazel.AqueryDepset
 }
 
 var _ BazelContext = &bazelContext{}
@@ -175,6 +182,10 @@ func (m MockBazelContext) BuildStatementsToRegister() []bazel.BuildStatement {
 	return []bazel.BuildStatement{}
 }
 
+func (m MockBazelContext) AqueryDepsets() []bazel.AqueryDepset {
+	return []bazel.AqueryDepset{}
+}
+
 var _ BazelContext = MockBazelContext{}
 
 func (bazelCtx *bazelContext) GetOutputFiles(label string, cfgKey configKey) ([]string, bool) {
@@ -236,6 +247,10 @@ func (m noopBazelContext) BuildStatementsToRegister() []bazel.BuildStatement {
 	return []bazel.BuildStatement{}
 }
 
+func (m noopBazelContext) AqueryDepsets() []bazel.AqueryDepset {
+	return []bazel.AqueryDepset{}
+}
+
 func NewBazelContext(c *config) (BazelContext, error) {
 	// TODO(cparsons): Assess USE_BAZEL=1 instead once "mixed Soong/Bazel builds"
 	// are production ready.
@@ -746,7 +761,7 @@ func (context *bazelContext) InvokeBazel() error {
 		return err
 	}
 
-	context.buildStatements, err = bazel.AqueryBuildStatements([]byte(aqueryOutput))
+	context.buildStatements, context.depsets, err = bazel.AqueryBuildStatements([]byte(aqueryOutput))
 	if err != nil {
 		return err
 	}
@@ -772,6 +787,10 @@ func (context *bazelContext) BuildStatementsToRegister() []bazel.BuildStatement
 	return context.buildStatements
 }
 
+func (context *bazelContext) AqueryDepsets() []bazel.AqueryDepset {
+	return context.depsets
+}
+
 func (context *bazelContext) OutputBase() string {
 	return context.paths.outputBase
 }
@@ -804,6 +823,23 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 		ctx.AddNinjaFileDeps(file)
 	}
 
+	for _, depset := range ctx.Config().BazelContext.AqueryDepsets() {
+		var outputs []Path
+		for _, depsetDepId := range depset.TransitiveDepSetIds {
+			otherDepsetName := bazelDepsetName(depsetDepId)
+			outputs = append(outputs, PathForPhony(ctx, otherDepsetName))
+		}
+		for _, artifactPath := range depset.DirectArtifacts {
+			outputs = append(outputs, PathForBazelOut(ctx, artifactPath))
+		}
+		thisDepsetName := bazelDepsetName(depset.Id)
+		ctx.Build(pctx, BuildParams{
+			Rule:      blueprint.Phony,
+			Outputs:   []WritablePath{PathForPhony(ctx, thisDepsetName)},
+			Implicits: outputs,
+		})
+	}
+
 	// Register bazel-owned build statements (obtained from the aquery invocation).
 	for index, buildStatement := range ctx.Config().BazelContext.BuildStatementsToRegister() {
 		if len(buildStatement.Command) < 1 {
@@ -838,6 +874,10 @@ func (c *bazelSingleton) GenerateBuildActions(ctx SingletonContext) {
 		for _, inputPath := range buildStatement.InputPaths {
 			cmd.Implicit(PathForBazelOut(ctx, inputPath))
 		}
+		for _, inputDepsetId := range buildStatement.InputDepsetIds {
+			otherDepsetName := bazelDepsetName(inputDepsetId)
+			cmd.Implicit(PathForPhony(ctx, otherDepsetName))
+		}
 
 		if depfile := buildStatement.Depfile; depfile != nil {
 			cmd.ImplicitDepFile(PathForBazelOut(ctx, *depfile))
@@ -883,3 +923,7 @@ func GetConfigKey(ctx ModuleContext) configKey {
 		osType: ctx.Os(),
 	}
 }
+
+func bazelDepsetName(depsetId int) string {
+	return fmt.Sprintf("bazel_depset_%d", depsetId)
+}

bazel/aquery.go

@@ -43,6 +43,18 @@ type KeyValuePair struct {
 	Value string
 }
 
+// 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.
+// 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
+}
+
 // depSetOfFiles contains relevant portions of Bazel's aquery proto, DepSetOfFiles.
 // Represents a data structure containing one or more files. Depsets in Bazel are an efficient
 // data structure for storing large numbers of file paths.
@@ -79,21 +91,21 @@ type BuildStatement struct {
 	Command      string
 	Depfile      *string
 	OutputPaths  []string
-	InputPaths   []string
 	SymlinkPaths []string
 	Env          []KeyValuePair
 	Mnemonic     string
+
+	// Inputs of this build statement, either as unexpanded depsets or expanded
+	// 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
 }
 
 // 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 middleman artifact Id 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
-	// that action instead.
-	middlemanIdToDepsetIds map[int][]int
 	// Maps depset Id to depset struct.
 	depsetIdToDepset map[int]depSetOfFiles
 	// depsetIdToArtifactIdsCache is a memoization of depset flattening, because flattening
@@ -132,12 +144,11 @@ func newAqueryHandler(aqueryResult actionGraphContainer) (*aqueryArtifactHandler
 		artifactIdToPath[artifact.Id] = artifactPath
 	}
 
-	depsetIdToDepset := map[int]depSetOfFiles{}
+	// Map middleman artifact Id to input artifact depset ID.
-	for _, depset := range aqueryResult.DepSetOfFiles {
+	// Middleman artifacts are treated as "substitute" artifacts for mixed builds. For example,
-		depsetIdToDepset[depset.Id] = depset
+	// 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
-
+	// that action instead.
-	// Do a pass through all actions to identify which artifacts are middleman artifacts.
 	middlemanIdToDepsetIds := map[int][]int{}
 	for _, actionEntry := range aqueryResult.Actions {
 		if actionEntry.Mnemonic == "Middleman" {
@@ -146,14 +157,64 @@ func newAqueryHandler(aqueryResult actionGraphContainer) (*aqueryArtifactHandler
 			}
 		}
 	}
+
+	// Store all depset IDs to validate all depset links are resolvable.
+	depsetIds := map[int]bool{}
+	for _, depset := range aqueryResult.DepSetOfFiles {
+		depsetIds[depset.Id] = true
+	}
+
+	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{
-		middlemanIdToDepsetIds:     middlemanIdToDepsetIds,
 		depsetIdToDepset:           depsetIdToDepset,
 		depsetIdToArtifactIdsCache: map[int][]int{},
 		artifactIdToPath:           artifactIdToPath,
 	}, nil
 }
 
+// getInputPaths flattens the depsets of the given IDs and returns all transitive
+// input paths contained in these depsets.
+// This is a potentially expensive operation, and should not be invoked except
+// for actions which need specialized input handling.
 func (a *aqueryArtifactHandler) getInputPaths(depsetIds []int) ([]string, error) {
 	inputPaths := []string{}
 
@@ -163,48 +224,15 @@ func (a *aqueryArtifactHandler) getInputPaths(depsetIds []int) ([]string, error)
 			return nil, err
 		}
 		for _, inputId := range inputArtifacts {
-			if middlemanInputDepsetIds, isMiddlemanArtifact := a.middlemanIdToDepsetIds[inputId]; isMiddlemanArtifact {
+			inputPath, exists := a.artifactIdToPath[inputId]
-				// Add all inputs from middleman actions which created middleman artifacts which are
+			if !exists {
-				// in the inputs for this action.
+				return nil, fmt.Errorf("undefined input artifactId %d", inputId)
-				swappedInputPaths, err := a.getInputPaths(middlemanInputDepsetIds)
-				if err != nil {
-					return nil, err
-				}
-				inputPaths = append(inputPaths, swappedInputPaths...)
-			} else {
-				inputPath, exists := a.artifactIdToPath[inputId]
-				if !exists {
-					return nil, fmt.Errorf("undefined input artifactId %d", inputId)
-				}
-				inputPaths = append(inputPaths, inputPath)
 			}
+			inputPaths = append(inputPaths, inputPath)
 		}
 	}
 
-	// TODO(b/197135294): Clean up this custom runfiles handling logic when
+	return inputPaths, nil
-	// SourceSymlinkManifest and SymlinkTree actions are supported.
-	filteredInputPaths := filterOutPy3wrapperAndManifestFileFromInputPaths(inputPaths)
-
-	return filteredInputPaths, nil
-}
-
-// 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.
-func filterOutPy3wrapperAndManifestFileFromInputPaths(inputPaths []string) []string {
-	filteredInputPaths := []string{}
-	for _, path := range inputPaths {
-		if strings.HasSuffix(path, py3wrapperFileName) || manifestFilePattern.MatchString(path) {
-			continue
-		}
-		filteredInputPaths = append(filteredInputPaths, path)
-	}
-	return filteredInputPaths
 }
 
 func (a *aqueryArtifactHandler) artifactIdsFromDepsetId(depsetId int) ([]int, error) {
@@ -227,115 +255,233 @@ func (a *aqueryArtifactHandler) artifactIdsFromDepsetId(depsetId int) ([]int, er
 	}
 }
 
-// AqueryBuildStatements returns an array of BuildStatements which should be registered (and output
+// AqueryBuildStatements returns a slice of BuildStatements and a slice of AqueryDepset
-// to a ninja file) to correspond one-to-one with the given action graph json proto (from a bazel
+// which should be registered (and output  to a ninja file) to correspond with Bazel's
-// aquery invocation).
+// action graph, as described by the given action graph json proto.
-func AqueryBuildStatements(aqueryJsonProto []byte) ([]BuildStatement, error) {
+// 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 {
-		return nil, err
+		return nil, nil, err
 	}
 	aqueryHandler, err := newAqueryHandler(aqueryResult)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	for _, actionEntry := range aqueryResult.Actions {
 		if shouldSkipAction(actionEntry) {
 			continue
 		}
-		outputPaths := []string{}
-		var depfile *string
-		for _, outputId := range actionEntry.OutputIds {
-			outputPath, exists := aqueryHandler.artifactIdToPath[outputId]
-			if !exists {
-				return nil, fmt.Errorf("undefined outputId %d", outputId)
-			}
-			ext := filepath.Ext(outputPath)
-			if ext == ".d" {
-				if depfile != nil {
-					return nil, fmt.Errorf("found multiple potential depfiles %q, %q", *depfile, outputPath)
-				} else {
-					depfile = &outputPath
-				}
-			} else {
-				outputPaths = append(outputPaths, outputPath)
-			}
-		}
-		inputPaths, err := aqueryHandler.getInputPaths(actionEntry.InputDepSetIds)
-		if err != nil {
-			return nil, err
-		}
-
-		buildStatement := BuildStatement{
-			Command:     strings.Join(proptools.ShellEscapeListIncludingSpaces(actionEntry.Arguments), " "),
-			Depfile:     depfile,
-			OutputPaths: outputPaths,
-			InputPaths:  inputPaths,
-			Env:         actionEntry.EnvironmentVariables,
-			Mnemonic:    actionEntry.Mnemonic,
-		}
 
+		var buildStatement BuildStatement
 		if isSymlinkAction(actionEntry) {
-			if len(inputPaths) != 1 || len(outputPaths) != 1 {
+			buildStatement, err = aqueryHandler.symlinkActionBuildStatement(actionEntry)
-				return nil, fmt.Errorf("Expect 1 input and 1 output to symlink action, got: input %q, output %q", inputPaths, outputPaths)
-			}
-			out := outputPaths[0]
-			outDir := proptools.ShellEscapeIncludingSpaces(filepath.Dir(out))
-			out = proptools.ShellEscapeIncludingSpaces(out)
-			in := filepath.Join("$PWD", proptools.ShellEscapeIncludingSpaces(inputPaths[0]))
-			// Use absolute paths, because some soong actions don't play well with relative paths (for example, `cp -d`).
-			buildStatement.Command = fmt.Sprintf("mkdir -p %[1]s && rm -f %[2]s && ln -sf %[3]s %[2]s", outDir, out, in)
-			buildStatement.SymlinkPaths = outputPaths[:]
 		} else if isTemplateExpandAction(actionEntry) && len(actionEntry.Arguments) < 1 {
-			if len(outputPaths) != 1 {
+			buildStatement, err = aqueryHandler.templateExpandActionBuildStatement(actionEntry)
-				return nil, fmt.Errorf("Expect 1 output to template expand action, got: output %q", outputPaths)
-			}
-			expandedTemplateContent := expandTemplateContent(actionEntry)
-			// The expandedTemplateContent is escaped for being used in double quotes and shell unescape,
-			// and the new line characters (\n) are also changed to \\n which avoids some Ninja escape on \n, which might
-			// change \n to space and mess up the format of Python programs.
-			// sed is used to convert \\n back to \n before saving to output file.
-			// 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])
-			buildStatement.Command = command
 		} else if isPythonZipperAction(actionEntry) {
-			if len(inputPaths) < 1 || len(outputPaths) != 1 {
+			buildStatement, err = aqueryHandler.pythonZipperActionBuildStatement(actionEntry, buildStatements)
-				return nil, fmt.Errorf("Expect 1+ input and 1 output to python zipper action, got: input %q, output %q", inputPaths, outputPaths)
-			}
-			buildStatement.InputPaths, buildStatement.Command = removePy3wrapperScript(buildStatement)
-			buildStatement.Command = addCommandForPyBinaryRunfilesDir(buildStatement, inputPaths[0], outputPaths[0])
-			// Add the python zip file as input of the corresponding python binary stub script in Ninja build statements.
-			// In Ninja build statements, the outputs of dependents of a python binary have python binary stub script as input,
-			// which is not sufficient without the python zip file from which runfiles directory is created for py_binary.
-			//
-			// The following logic relies on that Bazel aquery output returns actions in the order that
-			// PythonZipper is after TemplateAction of creating Python binary stub script. If later Bazel doesn't return actions
-			// in that order, the following logic might not find the build statement generated for Python binary
-			// stub script and the build might fail. So the check of pyBinaryFound is added to help debug in case later Bazel might change aquery output.
-			// See go/python-binary-host-mixed-build for more details.
-			pythonZipFilePath := outputPaths[0]
-			pyBinaryFound := false
-			for i, _ := range buildStatements {
-				if len(buildStatements[i].OutputPaths) == 1 && buildStatements[i].OutputPaths[0]+".zip" == pythonZipFilePath {
-					buildStatements[i].InputPaths = append(buildStatements[i].InputPaths, pythonZipFilePath)
-					pyBinaryFound = true
-				}
-			}
-			if !pyBinaryFound {
-				return nil, fmt.Errorf("Could not find the correspondinging Python binary stub script of PythonZipper: %q", outputPaths)
-			}
 		} else if len(actionEntry.Arguments) < 1 {
-			return nil, fmt.Errorf("received action with no command: [%v]", buildStatement)
+			return nil, nil, fmt.Errorf("received action with no command: [%s]", actionEntry.Mnemonic)
+		} else {
+			buildStatement, err = aqueryHandler.normalActionBuildStatement(actionEntry)
+		}
+
+		if err != nil {
+			return nil, nil, err
 		}
 		buildStatements = append(buildStatements, buildStatement)
 	}
 
-	return buildStatements, nil
+	// 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)
+		}
+		aqueryDepset := AqueryDepset{
+			Id:                  depset.Id,
+			DirectArtifacts:     directPaths,
+			TransitiveDepSetIds: depset.TransitiveDepSetIds,
+		}
+		depsets = append(depsets, aqueryDepset)
+	}
+	return buildStatements, depsets, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) validateInputDepsets(inputDepsetIds []int) ([]int, error) {
+	// Validate input depsets correspond to real depsets.
+	for _, depsetId := range inputDepsetIds {
+		if _, exists := aqueryHandler.depsetIdToDepset[depsetId]; !exists {
+			return nil, fmt.Errorf("undefined input depsetId %d", depsetId)
+		}
+	}
+	return inputDepsetIds, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) normalActionBuildStatement(actionEntry action) (BuildStatement, error) {
+	command := strings.Join(proptools.ShellEscapeListIncludingSpaces(actionEntry.Arguments), " ")
+	inputDepsetIds, err := aqueryHandler.validateInputDepsets(actionEntry.InputDepSetIds)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+	outputPaths, depfile, err := aqueryHandler.getOutputPaths(actionEntry)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+
+	buildStatement := BuildStatement{
+		Command:        command,
+		Depfile:        depfile,
+		OutputPaths:    outputPaths,
+		InputDepsetIds: inputDepsetIds,
+		Env:            actionEntry.EnvironmentVariables,
+		Mnemonic:       actionEntry.Mnemonic,
+	}
+	return buildStatement, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) pythonZipperActionBuildStatement(actionEntry action, prevBuildStatements []BuildStatement) (BuildStatement, error) {
+	inputPaths, err := aqueryHandler.getInputPaths(actionEntry.InputDepSetIds)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+	outputPaths, depfile, err := aqueryHandler.getOutputPaths(actionEntry)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+
+	if len(inputPaths) < 1 || len(outputPaths) != 1 {
+		return BuildStatement{}, fmt.Errorf("Expect 1+ input and 1 output to python zipper action, got: input %q, output %q", inputPaths, outputPaths)
+	}
+	command := strings.Join(proptools.ShellEscapeListIncludingSpaces(actionEntry.Arguments), " ")
+	inputPaths, command = removePy3wrapperScript(inputPaths, command)
+	command = addCommandForPyBinaryRunfilesDir(command, inputPaths[0], outputPaths[0])
+	// Add the python zip file as input of the corresponding python binary stub script in Ninja build statements.
+	// In Ninja build statements, the outputs of dependents of a python binary have python binary stub script as input,
+	// which is not sufficient without the python zip file from which runfiles directory is created for py_binary.
+	//
+	// The following logic relies on that Bazel aquery output returns actions in the order that
+	// PythonZipper is after TemplateAction of creating Python binary stub script. If later Bazel doesn't return actions
+	// in that order, the following logic might not find the build statement generated for Python binary
+	// stub script and the build might fail. So the check of pyBinaryFound is added to help debug in case later Bazel might change aquery output.
+	// See go/python-binary-host-mixed-build for more details.
+	pythonZipFilePath := outputPaths[0]
+	pyBinaryFound := false
+	for i, _ := range prevBuildStatements {
+		if len(prevBuildStatements[i].OutputPaths) == 1 && prevBuildStatements[i].OutputPaths[0]+".zip" == pythonZipFilePath {
+			prevBuildStatements[i].InputPaths = append(prevBuildStatements[i].InputPaths, pythonZipFilePath)
+			pyBinaryFound = true
+		}
+	}
+	if !pyBinaryFound {
+		return BuildStatement{}, fmt.Errorf("Could not find the correspondinging Python binary stub script of PythonZipper: %q", outputPaths)
+	}
+
+	buildStatement := BuildStatement{
+		Command:     command,
+		Depfile:     depfile,
+		OutputPaths: outputPaths,
+		InputPaths:  inputPaths,
+		Env:         actionEntry.EnvironmentVariables,
+		Mnemonic:    actionEntry.Mnemonic,
+	}
+	return buildStatement, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) templateExpandActionBuildStatement(actionEntry action) (BuildStatement, error) {
+	outputPaths, depfile, err := aqueryHandler.getOutputPaths(actionEntry)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+	if len(outputPaths) != 1 {
+		return BuildStatement{}, fmt.Errorf("Expect 1 output to template expand action, got: output %q", outputPaths)
+	}
+	expandedTemplateContent := expandTemplateContent(actionEntry)
+	// The expandedTemplateContent is escaped for being used in double quotes and shell unescape,
+	// and the new line characters (\n) are also changed to \\n which avoids some Ninja escape on \n, which might
+	// change \n to space and mess up the format of Python programs.
+	// sed is used to convert \\n back to \n before saving to output file.
+	// 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)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+
+	buildStatement := BuildStatement{
+		Command:        command,
+		Depfile:        depfile,
+		OutputPaths:    outputPaths,
+		InputDepsetIds: inputDepsetIds,
+		Env:            actionEntry.EnvironmentVariables,
+		Mnemonic:       actionEntry.Mnemonic,
+	}
+	return buildStatement, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) symlinkActionBuildStatement(actionEntry action) (BuildStatement, error) {
+	outputPaths, depfile, err := aqueryHandler.getOutputPaths(actionEntry)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+
+	inputPaths, err := aqueryHandler.getInputPaths(actionEntry.InputDepSetIds)
+	if err != nil {
+		return BuildStatement{}, err
+	}
+	if len(inputPaths) != 1 || len(outputPaths) != 1 {
+		return BuildStatement{}, fmt.Errorf("Expect 1 input and 1 output to symlink action, got: input %q, output %q", inputPaths, outputPaths)
+	}
+	out := outputPaths[0]
+	outDir := proptools.ShellEscapeIncludingSpaces(filepath.Dir(out))
+	out = proptools.ShellEscapeIncludingSpaces(out)
+	in := filepath.Join("$PWD", proptools.ShellEscapeIncludingSpaces(inputPaths[0]))
+	// Use absolute paths, because some soong actions don't play well with relative paths (for example, `cp -d`).
+	command := fmt.Sprintf("mkdir -p %[1]s && rm -f %[2]s && ln -sf %[3]s %[2]s", outDir, out, in)
+	symlinkPaths := outputPaths[:]
+
+	buildStatement := BuildStatement{
+		Command:      command,
+		Depfile:      depfile,
+		OutputPaths:  outputPaths,
+		InputPaths:   inputPaths,
+		Env:          actionEntry.EnvironmentVariables,
+		Mnemonic:     actionEntry.Mnemonic,
+		SymlinkPaths: symlinkPaths,
+	}
+	return buildStatement, nil
+}
+
+func (aqueryHandler *aqueryArtifactHandler) getOutputPaths(actionEntry action) (outputPaths []string, depfile *string, err error) {
+	for _, outputId := range actionEntry.OutputIds {
+		outputPath, exists := aqueryHandler.artifactIdToPath[outputId]
+		if !exists {
+			err = fmt.Errorf("undefined outputId %d", outputId)
+			return
+		}
+		ext := filepath.Ext(outputPath)
+		if ext == ".d" {
+			if depfile != nil {
+				err = fmt.Errorf("found multiple potential depfiles %q, %q", *depfile, outputPath)
+				return
+			} else {
+				depfile = &outputPath
+			}
+		} else {
+			outputPaths = append(outputPaths, outputPath)
+		}
+	}
+	return
 }
 
 // expandTemplateContent substitutes the tokens in a template.
@@ -372,10 +518,10 @@ func escapeCommandlineArgument(str string) string {
 // removed from input paths and command of creating python zip file.
 // See go/python-binary-host-mixed-build for more details.
 // TODO(b/205879240) remove this after py3wrapper.sh could be created in the mixed build mode.
-func removePy3wrapperScript(bs BuildStatement) (newInputPaths []string, newCommand string) {
+func removePy3wrapperScript(inputPaths []string, command string) (newInputPaths []string, newCommand string) {
 	// Remove from inputs
 	filteredInputPaths := []string{}
-	for _, path := range bs.InputPaths {
+	for _, path := range inputPaths {
 		if !strings.HasSuffix(path, py3wrapperFileName) {
 			filteredInputPaths = append(filteredInputPaths, path)
 		}
@@ -384,7 +530,7 @@ func removePy3wrapperScript(bs BuildStatement) (newInputPaths []string, newComma
 
 	// Remove from command line
 	var re = regexp.MustCompile(`\S*` + py3wrapperFileName)
-	newCommand = re.ReplaceAllString(bs.Command, "")
+	newCommand = re.ReplaceAllString(command, "")
 	return
 }
 
@@ -395,14 +541,14 @@ func removePy3wrapperScript(bs BuildStatement) (newInputPaths []string, newComma
 // so MANIFEST file could not be created, which also blocks the creation of runfiles directory.
 // See go/python-binary-host-mixed-build for more details.
 // TODO(b/197135294) create runfiles directory from MANIFEST file once it can be created from SourceSymlinkManifest action.
-func addCommandForPyBinaryRunfilesDir(bs BuildStatement, zipperCommandPath, zipFilePath string) string {
+func addCommandForPyBinaryRunfilesDir(oldCommand string, zipperCommandPath, zipFilePath string) string {
 	// Unzip the zip file, zipFilePath looks like <python_binary>.zip
 	runfilesDirName := zipFilePath[0:len(zipFilePath)-4] + ".runfiles"
 	command := fmt.Sprintf("%s x %s -d %s", zipperCommandPath, zipFilePath, runfilesDirName)
 	// Create a symbolic link in <python_binary>.runfiles/, which is the expected structure
 	// when running the python binary stub script.
 	command += fmt.Sprintf(" && ln -sf runfiles/__main__ %s", runfilesDirName)
-	return bs.Command + " && " + command
+	return oldCommand + " && " + command
 }
 
 func isSymlinkAction(a action) bool {

bazel/aquery_test.go

@@ -223,7 +223,7 @@ func TestAqueryMultiArchGenrule(t *testing.T) {
     "parentId": 13
   }]
 }`
-	actualbuildStatements, _ := AqueryBuildStatements([]byte(inputString))
+	actualbuildStatements, actualDepsets, _ := AqueryBuildStatements([]byte(inputString))
 	expectedBuildStatements := []BuildStatement{}
 	for _, arch := range []string{"arm", "arm64", "x86", "x86_64"} {
 		expectedBuildStatements = append(expectedBuildStatements,
@@ -234,11 +234,7 @@ func TestAqueryMultiArchGenrule(t *testing.T) {
 				OutputPaths: []string{
 					fmt.Sprintf("bazel-out/sourceroot/k8-fastbuild/bin/bionic/libc/syscalls-%s.S", arch),
 				},
-				InputPaths: []string{
+				InputDepsetIds: []int{1},
-					"../sourceroot/bionic/libc/SYSCALLS.TXT",
-					"../sourceroot/bionic/libc/tools/gensyscalls.py",
-					"../bazel_tools/tools/genrule/genrule-setup.sh",
-				},
 				Env: []KeyValuePair{
 					KeyValuePair{Key: "PATH", Value: "/bin:/usr/bin:/usr/local/bin"},
 				},
@@ -246,6 +242,16 @@ func TestAqueryMultiArchGenrule(t *testing.T) {
 			})
 	}
 	assertBuildStatements(t, expectedBuildStatements, actualbuildStatements)
+
+	expectedFlattenedInputs := []string{
+		"../sourceroot/bionic/libc/SYSCALLS.TXT",
+		"../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)
+	}
 }
 
 func TestInvalidOutputId(t *testing.T) {
@@ -280,11 +286,11 @@ func TestInvalidOutputId(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, "undefined outputId 3")
 }
 
-func TestInvalidInputDepsetId(t *testing.T) {
+func TestInvalidInputDepsetIdFromAction(t *testing.T) {
 	const inputString = `
 {
   "artifacts": [{
@@ -316,10 +322,47 @@ func TestInvalidInputDepsetId(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, "undefined input depsetId 2")
 }
 
+func TestInvalidInputDepsetIdFromDepset(t *testing.T) {
+	const inputString = `
+{
+  "artifacts": [{
+    "id": 1,
+    "pathFragmentId": 1
+  }, {
+    "id": 2,
+    "pathFragmentId": 2
+  }],
+  "actions": [{
+    "targetId": 1,
+    "actionKey": "x",
+    "mnemonic": "x",
+    "arguments": ["touch", "foo"],
+    "inputDepSetIds": [1],
+    "outputIds": [1],
+    "primaryOutputId": 1
+  }],
+  "depSetOfFiles": [{
+    "id": 1,
+    "directArtifactIds": [1, 2],
+    "transitiveDepSetIds": [42]
+  }],
+  "pathFragments": [{
+    "id": 1,
+    "label": "one"
+  }, {
+    "id": 2,
+    "label": "two"
+  }]
+}`
+
+	_, _, err := AqueryBuildStatements([]byte(inputString))
+	assertError(t, err, "undefined input depsetId 42 (referenced by depsetId 1)")
+}
+
 func TestInvalidInputArtifactId(t *testing.T) {
 	const inputString = `
 {
@@ -352,7 +395,7 @@ func TestInvalidInputArtifactId(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, "undefined input artifactId 3")
 }
 
@@ -389,7 +432,7 @@ func TestInvalidPathFragmentId(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, "undefined path fragment id 3")
 }
 
@@ -431,7 +474,7 @@ func TestDepfiles(t *testing.T) {
   }]
 }`
 
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actual, _, err := AqueryBuildStatements([]byte(inputString))
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
 	}
@@ -492,7 +535,7 @@ func TestMultipleDepfiles(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `found multiple potential depfiles "two.d", "other.d"`)
 }
 
@@ -699,23 +742,31 @@ func TestTransitiveInputDepsets(t *testing.T) {
   }]
 }`
 
-	actualbuildStatements, _ := AqueryBuildStatements([]byte(inputString))
+	actualbuildStatements, actualDepsets, _ := AqueryBuildStatements([]byte(inputString))
-	// Inputs for the action are test_{i} from 1 to 20, and test_root. These inputs
+
-	// are given via a deep depset, but the depset is flattened when returned as a
-	// BuildStatement slice.
-	inputPaths := []string{"bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_root"}
-	for i := 1; i < 20; i++ {
-		inputPaths = append(inputPaths, fmt.Sprintf("bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_%d", i))
-	}
 	expectedBuildStatements := []BuildStatement{
 		BuildStatement{
-			Command:     "/bin/bash -c 'touch bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out'",
+			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"},
+			OutputPaths:    []string{"bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_out"},
-			InputPaths:  inputPaths,
+			InputDepsetIds: []int{1},
-			Mnemonic:    "Action",
+			Mnemonic:       "Action",
 		},
 	}
 	assertBuildStatements(t, expectedBuildStatements, actualbuildStatements)
+
+	// Inputs for the action are test_{i} from 1 to 20, and test_root. These inputs
+	// are given via a deep depset, but the depset is flattened when returned as a
+	// BuildStatement slice.
+	expectedFlattenedInputs := []string{}
+	for i := 1; i < 20; i++ {
+		expectedFlattenedInputs = append(expectedFlattenedInputs, fmt.Sprintf("bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_%d", i))
+	}
+	expectedFlattenedInputs = append(expectedFlattenedInputs, "bazel-out/sourceroot/k8-fastbuild/bin/testpkg/test_root")
+
+	actualFlattenedInputs := flattenDepsets([]int{1}, actualDepsets)
+	if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
+		t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
+	}
 }
 
 func TestMiddlemenAction(t *testing.T) {
@@ -785,24 +836,74 @@ func TestMiddlemenAction(t *testing.T) {
   }]
 }`
 
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actualBuildStatements, actualDepsets, err := AqueryBuildStatements([]byte(inputString))
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
 	}
-	if expected := 1; len(actual) != expected {
+	if expected := 1; len(actualBuildStatements) != expected {
-		t.Fatalf("Expected %d build statements, got %d", expected, len(actual))
+		t.Fatalf("Expected %d build statements, got %d", expected, len(actualBuildStatements))
 	}
 
-	bs := actual[0]
+	bs := actualBuildStatements[0]
-	expectedInputs := []string{"middleinput_one", "middleinput_two", "maininput_one", "maininput_two"}
+	if len(bs.InputPaths) > 0 {
-	if !reflect.DeepEqual(bs.InputPaths, expectedInputs) {
+		t.Errorf("Expected main action raw inputs to be empty, but got %q", bs.InputPaths)
-		t.Errorf("Expected main action inputs %q, but got %q", expectedInputs, 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)
+
+	if !reflect.DeepEqual(actualFlattenedInputs, expectedFlattenedInputs) {
+		t.Errorf("Expected flattened inputs %v, but got %v", expectedFlattenedInputs, actualFlattenedInputs)
+	}
+}
+
+// Returns the contents of given depsets in concatenated post order.
+func flattenDepsets(depsetIdsToFlatten []int, allDepsets []AqueryDepset) []string {
+	depsetsById := map[int]AqueryDepset{}
+	for _, depset := range allDepsets {
+		depsetsById[depset.Id] = depset
+	}
+	result := []string{}
+	for _, depsetId := range depsetIdsToFlatten {
+		result = append(result, flattenDepset(depsetId, depsetsById)...)
+	}
+	return result
+}
+
+// Returns the contents of a given depset in post order.
+func flattenDepset(depsetIdToFlatten int, allDepsets map[int]AqueryDepset) []string {
+	depset := allDepsets[depsetIdToFlatten]
+	result := []string{}
+	for _, depsetId := range depset.TransitiveDepSetIds {
+		result = append(result, flattenDepset(depsetId, allDepsets)...)
+	}
+	result = append(result, depset.DirectArtifacts...)
+	return result
 }
 
 func TestSimpleSymlink(t *testing.T) {
@@ -849,7 +950,7 @@ func TestSimpleSymlink(t *testing.T) {
   }]
 }`
 
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actual, _, err := AqueryBuildStatements([]byte(inputString))
 
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
@@ -913,7 +1014,7 @@ func TestSymlinkQuotesPaths(t *testing.T) {
   }]
 }`
 
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actual, _, err := AqueryBuildStatements([]byte(inputString))
 
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
@@ -970,7 +1071,7 @@ func TestSymlinkMultipleInputs(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `Expect 1 input and 1 output to symlink action, got: input ["file" "other_file"], output ["symlink"]`)
 }
 
@@ -1011,7 +1112,7 @@ func TestSymlinkMultipleOutputs(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `Expect 1 input and 1 output to symlink action, got: input ["file"], output ["symlink" "other_symlink"]`)
 }
 
@@ -1045,7 +1146,7 @@ func TestTemplateExpandActionSubstitutions(t *testing.T) {
   }]
 }`
 
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actual, _, err := AqueryBuildStatements([]byte(inputString))
 
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
@@ -1091,7 +1192,7 @@ func TestTemplateExpandActionNoOutput(t *testing.T) {
   }]
 }`
 
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `Expect 1 output to template expand action, got: output []`)
 }
 
@@ -1211,7 +1312,7 @@ func TestPythonZipperActionSuccess(t *testing.T) {
     "label": "python_binary"
   }]
 }`
-	actual, err := AqueryBuildStatements([]byte(inputString))
+	actual, _, err := AqueryBuildStatements([]byte(inputString))
 
 	if err != nil {
 		t.Errorf("Unexpected error %q", err)
@@ -1264,7 +1365,7 @@ func TestPythonZipperActionNoInput(t *testing.T) {
     "label": "python_binary.zip"
   }]
 }`
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `Expect 1+ input and 1 output to python zipper action, got: input [], output ["python_binary.zip"]`)
 }
 
@@ -1360,7 +1461,7 @@ func TestPythonZipperActionNoOutput(t *testing.T) {
     "parentId": 11
   }]
 }`
-	_, err := AqueryBuildStatements([]byte(inputString))
+	_, _, err := AqueryBuildStatements([]byte(inputString))
 	assertError(t, err, `Expect 1+ input and 1 output to python zipper action, got: input ["../bazel_tools/tools/zip/zipper/zipper" "python_binary.py"], output []`)
 }