Add ability to replace all resources of a type

Currently only works for shaders
UpdateResourceBinary should read the resource data from
standard input and output the modified resource data to
standard output.
Also modifying replace_resource to consume spir-v
disassembly to replace a single shader for consistency

cmd/gapit/flags.go

@@ -155,12 +155,13 @@ type (
 		CommandFilterFlags
 	}
 	ReplaceResourceFlags struct {
-		Gapis           GapisFlags
-		Gapir           GapirFlags
-		Handle          string `help:"required. handle of the resource to replace"`
-		ResourcePath    string `help:"required. file path for the new resource"`
-		At              int    `help:"command index to replace the resource at"`
-		OutputTraceFile string `help:"file name for the updated trace"`
+		Gapis                GapisFlags
+		Gapir                GapirFlags
+		Handle               string `help:"required. handle of the resource to replace"`
+		ResourcePath         string `help:"file path for the new resource"`
+		At                   int    `help:"command index to replace the resource(s) at"`
+		UpdateResourceBinary string `help:"shaders only. binary to run for every shader; consumes resource data from standard input and writes to standard output"`
+		OutputTraceFile      string `help:"file name for the updated trace"`
 	}
 	StateFlags struct {
 		Gapis  GapisFlags

cmd/gapit/replace_resource.go

@@ -18,14 +18,18 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"io"
 	"io/ioutil"
+	"os/exec"
 	"path/filepath"
 	"strings"
+	"syscall"
 
 	"github.com/google/gapid/core/app"
 	"github.com/google/gapid/core/log"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/service"
+	"github.com/google/gapid/gapis/service/path"
 )
 
 type replaceResourceVerb struct{ ReplaceResourceFlags }
@@ -50,13 +54,13 @@ func (verb *replaceResourceVerb) Run(ctx context.Context, flags flag.FlagSet) er
 		return nil
 	}
 
-	if verb.Handle == "" {
-		app.Usage(ctx, "-handle argument is required")
+	if verb.Handle == verb.UpdateResourceBinary {
+		app.Usage(ctx, "only one of -handle or -updateresourcebinary arguments is required")
 		return nil
 	}
 
-	if verb.ResourcePath == "" {
-		app.Usage(ctx, "-resourcepath argument is required")
+	if verb.Handle != "" && verb.ResourcePath == "" {
+		app.Usage(ctx, "-resourcepath argument is required if -handle is specified")
 		return nil
 	}
 
@@ -90,36 +94,102 @@ func (verb *replaceResourceVerb) Run(ctx context.Context, flags flag.FlagSet) er
 		verb.At = int(boxedCapture.(*service.Capture).NumCommands) - 1
 	}
 
+	var resourcePath *path.Any
+	var resourceData interface{}
+
 	for _, types := range resources.GetTypes() {
 		if types.Type == api.ResourceType_ShaderResource {
-			var matchedResource *service.Resource
-			for _, v := range types.GetResources() {
-				if strings.Contains(v.GetHandle(), verb.Handle) {
-					if matchedResource != nil {
-						return fmt.Errorf("Multiple resources matched: %s, %s", matchedResource.GetHandle(), v.GetHandle())
+			switch {
+			case verb.Handle != "":
+				matchedResource, err := resources.FindSingle(func(t api.ResourceType, r service.Resource) bool {
+					return strings.Contains(r.GetHandle(), verb.Handle)
+				})
+				if err != nil {
+					return err
+				}
+				resourcePath = capture.Command(uint64(verb.At)).ResourceAfter(matchedResource.ID).Path()
+				newResourceBytes, err := ioutil.ReadFile(verb.ResourcePath)
+				if err != nil {
+					return log.Errf(ctx, err, "Could not read resource file %s", verb.ResourcePath)
+				}
+				resourceData = api.NewResourceData(&api.Shader{Type: api.ShaderType_Spirv, Source: string(newResourceBytes)})
+			case verb.UpdateResourceBinary != "":
+				resources := types.GetResources()
+				ids := make([]*path.ID, len(resources))
+				resourcesSource := make([]*api.ResourceData, len(resources))
+				for i, v := range resources {
+					ids[i] = v.ID
+					resourcePath := capture.Command(uint64(verb.At)).ResourceAfter(v.ID)
+					rd, err := client.Get(ctx, resourcePath.Path())
+					if err != nil {
+						log.Errf(ctx, err, "Could not get data for shader: %v", v)
+						return err
+					}
+					newData, err := verb.getNewResourceData(ctx, rd.(*api.ResourceData))
+					if err != nil {
+						log.Errf(ctx, err, "Could not update the shader: %v", v)
+						return err
 					}
-					matchedResource = v
+					resourcesSource[i] = api.NewResourceData(&api.Shader{Type: api.ShaderType_Spirv, Source: string(newData)})
 				}
+				resourceData = api.NewMultiResourceData(resourcesSource)
+				resourcePath = capture.Command(uint64(verb.At)).ResourcesAfter(ids).Path()
 			}
-			resourcePath := capture.Command(uint64(verb.At)).ResourceAfter(matchedResource.ID)
-			newResourceBytes, err := ioutil.ReadFile(verb.ResourcePath)
-			if err != nil {
-				return log.Errf(ctx, err, "Could not read resource file %s", verb.ResourcePath)
-			}
+		}
+	}
+	newResourcePath, err := client.Set(ctx, resourcePath, resourceData)
+	if err != nil {
+		return log.Errf(ctx, err, "Could not update resource data: %v", resourcePath)
+	}
+	newCapture := path.FindCapture(newResourcePath.Node())
+	newCaptureFilepath, err := filepath.Abs(verb.OutputTraceFile)
+	err = client.SaveCapture(ctx, newCapture, newCaptureFilepath)
 
-			newResourceData := api.NewResourceData(&api.Shader{Type: api.ShaderType_SpirvBinary, Source: string(newResourceBytes)})
-			newResourcePath, err := client.Set(ctx, resourcePath.Path(), newResourceData)
-			if err != nil {
-				return log.Errf(ctx, err, "Could not update data for shader: %v", matchedResource)
-			}
-			newCapture := newResourcePath.GetResourceData().GetAfter().GetCapture()
-			newCaptureFilepath, err := filepath.Abs(verb.OutputTraceFile)
-			err = client.SaveCapture(ctx, newCapture, newCaptureFilepath)
+	log.I(ctx, "Capture written to: %v", newCaptureFilepath)
+	return nil
+}
+
+// getNewResourceData runs the update resource binary on the old resource data
+// and returns the newly generated resource data
+func (verb *replaceResourceVerb) getNewResourceData(ctx context.Context, resourceData *api.ResourceData) (string, error) {
+	updateCmd := exec.Command(verb.UpdateResourceBinary)
+	stdin, err := updateCmd.StdinPipe()
+	if err != nil {
+		return "", log.Errf(ctx, err, "Failed to get stdinput pipe")
+	}
+
+	stdout, err := updateCmd.StdoutPipe()
+	if err != nil {
+		return "", log.Errf(ctx, err, "Failed to get stdoutput pipe")
+	}
+
+	stderr, err := updateCmd.StderrPipe()
+	if err != nil {
+		return "", log.Errf(ctx, err, "Failed to get stderr pipe")
+	}
+
+	if err = updateCmd.Start(); err != nil {
+		return "", log.Errf(ctx, err, "Could not start the update resource command")
+	}
+	shaderSource := resourceData.GetShader().GetSource()
+	io.WriteString(stdin, shaderSource)
+	stdin.Close()
+	newResourceBytes, err := ioutil.ReadAll(stdout)
+	if err != nil {
+		return "", err
+	}
+	errorMsg, err := ioutil.ReadAll(stderr)
+	if err != nil {
+		return "", err
+	}
 
-			log.I(ctx, "Capture written to: %v", newCaptureFilepath)
-			return nil
+	if err = updateCmd.Wait(); err != nil {
+		if exitError, ok := err.(*exec.ExitError); ok {
+			return "", fmt.Errorf("Error generating new resource data; status %d; error message: %s",
+				exitError.Sys().(syscall.WaitStatus).ExitStatus(), errorMsg)
 		}
+		return "", err
 	}
 
-	return fmt.Errorf("Failed to find the resource with the handle %s", verb.Handle)
+	return string(newResourceBytes), err
 }

gapis/api/gles/resources.go

@@ -321,9 +321,9 @@ func (s Shaderʳ) SetResourceData(
 	}
 	resourceID := resourceIDs[s]
 
-	resource := resources.Find(s.ResourceType(ctx), resourceID)
-	if resource != nil {
-		return fmt.Errorf("Couldn't find resource")
+	resource, err := resources.Find(s.ResourceType(ctx), resourceID)
+	if err != nil {
+		return err
 	}
 
 	c, err := capture.ResolveFromPath(ctx, capturePath)

gapis/api/resource.go

@@ -51,8 +51,8 @@ type Resource interface {
 
 // ResourceMeta represents resource with a state information obtained during building.
 type ResourceMeta struct {
-	Resource Resource    // Resolved resource.
-	IDMap    ResourceMap // Map for resolved resources to ids.
+	Resources []Resource  // Resolved resource.
+	IDMap     ResourceMap // Map for resolved resources to ids.
 }
 
 // ReplaceCallback is called from SetResourceData to propagate changes to current command stream.
@@ -100,3 +100,7 @@ func NewResourceData(data interface{}) *ResourceData {
 		panic(fmt.Errorf("%T is not a ResourceData type", data))
 	}
 }
+
+func NewMultiResourceData(resources []*ResourceData) *MultiResourceData {
+	return &MultiResourceData{Resources: resources}
+}

gapis/api/service.proto

@@ -127,6 +127,11 @@ message ResourceData {
   }
 }
 
+// MultiResourceData represents the state  of resources at a single point in a capture
+message MultiResourceData {
+  repeated ResourceData resources = 1;
+}
+
 // Texture represents a texture resource.
 message Texture {
   oneof type {

gapis/api/vulkan/resources.go

@@ -765,9 +765,9 @@ func (shader ShaderModuleObjectʳ) SetResourceData(
 	}
 	resourceID := resourceIDs[shader]
 
-	resource := resources.Find(shader.ResourceType(ctx), resourceID)
-	if resource == nil {
-		return fmt.Errorf("Couldn't find resource")
+	resource, err := resources.Find(shader.ResourceType(ctx), resourceID)
+	if err != nil {
+		return err
 	}
 
 	c, err := capture.ResolveFromPath(ctx, at.Capture)

gapis/resolve/resolvables.proto

@@ -91,7 +91,7 @@ message AllResourceDataResolvable {
 }
 
 message ResourceMetaResolvable {
-  path.ID ID = 1;
+  repeated path.ID IDs = 1;
   path.Command after = 2;
 }
 

gapis/resolve/resource_meta.go

@@ -24,8 +24,8 @@ import (
 )
 
 // ResourceMeta returns the metadata for the specified resource.
-func ResourceMeta(ctx context.Context, id *path.ID, after *path.Command) (*api.ResourceMeta, error) {
-	obj, err := database.Build(ctx, &ResourceMetaResolvable{ID: id, After: after})
+func ResourceMeta(ctx context.Context, ids []*path.ID, after *path.Command) (*api.ResourceMeta, error) {
+	obj, err := database.Build(ctx, &ResourceMetaResolvable{IDs: ids, After: after})
 	if err != nil {
 		return nil, err
 	}
@@ -42,14 +42,18 @@ func (r *ResourceMetaResolvable) Resolve(ctx context.Context) (interface{}, erro
 	if !ok {
 		return nil, fmt.Errorf("Cannot resolve resources at command: %v", r.After)
 	}
-	id := r.ID.ID()
-	val, ok := res.resources[id]
-	if !ok {
-		return nil, fmt.Errorf("Could not find resource %v", id)
+	ids := r.IDs
+	values := make([]api.Resource, len(ids))
+	for i, id := range ids {
+		val, ok := res.resources[id.ID()]
+		if !ok {
+			return nil, fmt.Errorf("Could not find resource %v", id.ID())
+		}
+		values[i] = val
 	}
 	result := &api.ResourceMeta{
-		IDMap:    res.resourceMap,
-		Resource: val,
+		IDMap:     res.resourceMap,
+		Resources: values,
 	}
 	return result, nil
 }

gapis/resolve/set.go

@@ -63,8 +63,8 @@ func change(ctx context.Context, p path.Node, val interface{}) (path.Node, error
 	case *path.Report:
 		return nil, fmt.Errorf("Reports are immutable")
 
-	case *path.ResourceData:
-		meta, err := ResourceMeta(ctx, p.ID, p.After)
+	case *path.MultiResourceData:
+		meta, err := ResourceMeta(ctx, p.IDs, p.After)
 		if err != nil {
 			return nil, err
 		}
@@ -85,12 +85,63 @@ func change(ctx context.Context, p path.Node, val interface{}) (path.Node, error
 			cmds[where] = with.(api.Cmd)
 		}
 
+		data, ok := val.(*api.MultiResourceData)
+		if !ok {
+			return nil, fmt.Errorf("Expected ResourceData, got %T", val)
+		}
+
+		for i, resource := range meta.Resources {
+			if err := resource.SetResourceData(ctx, p.After, data.Resources[i], meta.IDMap, replaceCommands); err != nil {
+				return nil, err
+			}
+		}
+
+		// Store the new command list
+		c, err := changeCommands(ctx, p.After.Capture, cmds)
+		if err != nil {
+			return nil, err
+		}
+
+		return &path.MultiResourceData{
+			IDs: p.IDs, // TODO: Shouldn't this change?
+			After: &path.Command{
+				Capture: c,
+				Indices: p.After.Indices,
+			},
+		}, nil
+
+	case *path.ResourceData:
+		meta, err := ResourceMeta(ctx, []*path.ID{p.ID}, p.After)
+		if err != nil {
+			return nil, err
+		}
+
+		cmdIdx := p.After.Indices[0]
+		// If we change resource data, subcommands do not affect this, so change
+		// the main command.
+
+		oldCmds, err := NCmds(ctx, p.After.Capture, cmdIdx+1)
+		if err != nil {
+			return nil, err
+		}
+
+		cmds := make([]api.Cmd, len(oldCmds))
+		copy(cmds, oldCmds)
+
+		replaceCommands := func(where uint64, with interface{}) {
+			cmds[where] = with.(api.Cmd)
+		}
+
 		data, ok := val.(*api.ResourceData)
 		if !ok {
 			return nil, fmt.Errorf("Expected ResourceData, got %T", val)
 		}
 
-		if err := meta.Resource.SetResourceData(ctx, p.After, data, meta.IDMap, replaceCommands); err != nil {
+		if len(meta.Resources) != 1 {
+			return nil, fmt.Errorf("Expected a single resource, got %d", len(meta.Resources))
+		}
+
+		if err := meta.Resources[0].SetResourceData(ctx, p.After, data, meta.IDMap, replaceCommands); err != nil {
 			return nil, err
 		}