Reserve the memory ranges needed by the initial commands.

This will instruct the allocator to not allocate over the memory
ranges needed by the initial commands.

cmd/linearize_trace/main.go

@@ -64,8 +64,7 @@ func run(ctx context.Context) error {
 		return err
 	}
 
-	initialCmds := capt.GetInitialCommands(ctx)
-	_ = initialCmds
+	initialCmds, _ := capt.GetInitialCommands(ctx)
 
 	log.I(ctx, "Generated %v initial commands", len(initialCmds))
 

gapis/api/gles/state_builder.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/gapid/core/data/compare"
 	"github.com/google/gapid/core/log"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/memory"
 	"github.com/google/gapid/gapis/replay/value"
@@ -34,16 +35,18 @@ type stateBuilder struct {
 	cb       CommandBuilder   // Default command builder for thread 0
 	preCmd   []func(api.Cmd)  // Actions to be done on the next written command
 	seen     map[interface{}]bool
+	memoryIntervals interval.U64RangeList
 }
 
-func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) []api.Cmd {
+func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) ([]api.Cmd, interval.U64RangeList) {
 	newState := api.NewStateWithAllocator(memory.NewBasicAllocator(value.ValidMemoryRanges), oldState.MemoryLayout)
 	sb := &stateBuilder{
 		ctx:      ctx,
 		oldState: oldState,
 		newState: newState,
 		cb:       CommandBuilder{Thread: 0},
 		seen:     map[interface{}]bool{},
+		memoryIntervals: interval.U64RangeList{},
 	}
 
 	// Ensure that all pool IDs are distinct between the old state and new state.
@@ -99,7 +102,7 @@ func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) []a
 		log.W(ctx, "Initial state: %v", d)
 	})
 
-	return sb.cmds
+	return sb.cmds, sb.memoryIntervals
 }
 
 func (sb *stateBuilder) E(ctx context.Context, fmt string, args ...interface{}) {
@@ -108,6 +111,8 @@ func (sb *stateBuilder) E(ctx context.Context, fmt string, args ...interface{})
 
 func (sb *stateBuilder) readsData(ctx context.Context, v interface{}) memory.Pointer {
 	tmp := sb.newState.AllocDataOrPanic(ctx, v)
+	rng := tmp.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	sb.preCmd = append(sb.preCmd, func(cmd api.Cmd) {
 		cmd.Extras().GetOrAppendObservations().AddRead(tmp.Data())
 		tmp.Free()
@@ -117,6 +122,8 @@ func (sb *stateBuilder) readsData(ctx context.Context, v interface{}) memory.Poi
 
 func (sb *stateBuilder) readsSlice(ctx context.Context, v U8ˢ) memory.Pointer {
 	tmp := sb.newState.AllocOrPanic(ctx, v.Count())
+	rng := tmp.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	id := v.ResourceID(ctx, sb.oldState)
 	sb.preCmd = append(sb.preCmd, func(cmd api.Cmd) {
 		cmd.Extras().GetOrAppendObservations().AddRead(tmp.Range(), id)
@@ -127,6 +134,8 @@ func (sb *stateBuilder) readsSlice(ctx context.Context, v U8ˢ) memory.Pointer {
 
 func (sb *stateBuilder) writesData(ctx context.Context, v interface{}) memory.Pointer {
 	tmp := sb.newState.AllocDataOrPanic(ctx, v)
+	rng := tmp.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	sb.preCmd = append(sb.preCmd, func(cmd api.Cmd) {
 		cmd.Extras().GetOrAppendObservations().AddWrite(tmp.Data())
 		tmp.Free()

gapis/api/gvr/gvr.go

@@ -18,6 +18,7 @@ import (
 	"context"
 
 	"github.com/google/gapid/core/image"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/api/gles"
 	"github.com/google/gapid/gapis/api/sync"
@@ -38,7 +39,9 @@ func (s *State) Root(ctx context.Context, p *path.State) (path.Node, error) {
 
 func (*State) SetupInitialState(ctx context.Context) {}
 
-func (*State) RebuildState(ctx context.Context, s *api.GlobalState) []api.Cmd { return nil }
+func (*State) RebuildState(ctx context.Context, s *api.GlobalState) ([]api.Cmd, interval.U64RangeList) {
+	return nil, nil
+}
 
 func (c *State) preMutate(ctx context.Context, s *api.GlobalState, cmd api.Cmd) error {
 	return nil

gapis/api/state.go

@@ -24,6 +24,7 @@ import (
 	"github.com/google/gapid/core/data/endian"
 	"github.com/google/gapid/core/data/id"
 	"github.com/google/gapid/core/log"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/core/os/device"
 	"github.com/google/gapid/gapis/database"
 	"github.com/google/gapid/gapis/memory"
@@ -80,7 +81,9 @@ type State interface {
 
 	// RebuildState returns a set of commands which, if executed
 	// on a clean new state, will reproduce the current state.
-	RebuildState(ctx context.Context, s *GlobalState) []Cmd
+	// The segments of memory that were used to create these commands
+	// are returned in the rangeList.
+	RebuildState(ctx context.Context, s *GlobalState) ([]Cmd, interval.U64RangeList)
 }
 
 // NewStateWithEmptyAllocator returns a new, default-initialized State object,

gapis/api/test/test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 
 	"github.com/google/gapid/core/image"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/service/path"
 )
@@ -44,7 +45,9 @@ func (s *State) Root(ctx context.Context, p *path.State) (path.Node, error) {
 
 func (*State) SetupInitialState(ctx context.Context) {}
 
-func (*State) RebuildState(ctx context.Context, s *api.GlobalState) []api.Cmd { return nil }
+func (*State) RebuildState(ctx context.Context, s *api.GlobalState) ([]api.Cmd, interval.U64RangeList) {
+	return nil, nil
+}
 
 func (c *State) preMutate(ctx context.Context, s *api.GlobalState, cmd api.Cmd) error {
 	return nil

gapis/api/vulkan/custom_replay.go

@@ -306,27 +306,25 @@ func (a *VkCreateDevice) Mutate(ctx context.Context, id api.CmdID, s *api.Global
 		newCreateInfoData := s.AllocDataOrPanic(ctx, createInfo)
 		allocated = append(allocated, &newCreateInfoData)
 		createInfoPtr = NewVkDeviceCreateInfoᶜᵖ(newCreateInfoData.Ptr())
-	}
-
-	cb := CommandBuilder{Thread: a.thread}
-	hijack := cb.ReplayCreateVkDevice(a.PhysicalDevice, createInfoPtr, a.PAllocator, a.PDevice, a.Result)
-	hijack.Extras().MustClone(a.Extras().All()...)
 
-	if b != nil {
+		cb := CommandBuilder{Thread: a.thread}
+		hijack := cb.ReplayCreateVkDevice(a.PhysicalDevice, createInfoPtr, a.PAllocator, a.PDevice, a.Result)
+		hijack.Extras().MustClone(a.Extras().All()...)
+
 		for _, d := range allocated {
 			hijack.AddRead(d.Data())
 		}
+
+		err := hijack.Mutate(ctx, id, s, b)
+		if err != nil {
+			return err
+		}
+		// Call the replayRegisterVkDevice() synthetic API function.
+		device := a.PDevice.MustRead(ctx, a, s, b)
+		return cb.ReplayRegisterVkDevice(a.PhysicalDevice, device, a.PCreateInfo).Mutate(ctx, id, s, b)
 	}
-
-	err := hijack.Mutate(ctx, id, s, b)
-
-	if b == nil || err != nil {
-		return err
-	}
-
-	// Call the replayRegisterVkDevice() synthetic API function.
-	device := a.PDevice.MustRead(ctx, a, s, b)
-	return cb.ReplayRegisterVkDevice(a.PhysicalDevice, device, a.PCreateInfo).Mutate(ctx, id, s, b)
+
+	return a.mutate(ctx, id, s, b)
 }
 
 func (a *VkDestroyDevice) Mutate(ctx context.Context, id api.CmdID, s *api.GlobalState, b *builder.Builder) error {

gapis/api/vulkan/replay.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/google/gapid/core/image"
 	"github.com/google/gapid/core/log"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/core/os/device"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/api/transform"
@@ -497,6 +498,7 @@ func (a API) Replay(
 	// Populate the dead-code eliminitation later, only once we are sure
 	// we will need it.
 	dceInfo := dCEInfo{}
+	initMem := interval.U64RangeList{}
 
 	expandedCmds := false
 	numInitialCommands := 0
@@ -518,7 +520,8 @@ func (a API) Replay(
 			numInitialCommands = dceInfo.ft.NumInitialCommands
 		} else {
 			// If the capture contains initial state, prepend the commands to build the state.
-			initialCmds := capture.GetInitialCommands(ctx)
+			initialCmds, im := capture.GetInitialCommands(ctx)
+			initMem = im
 			numInitialCommands = len(initialCmds)
 			if len(initialCmds) > 0 {
 				cmds = append(initialCmds, cmds...)

gapis/api/vulkan/state_rebuilder.go

@@ -18,33 +18,35 @@ import (
 	"context"
 
 	"github.com/google/gapid/core/log"
+	"github.com/google/gapid/core/math/interval"
 	"github.com/google/gapid/gapis/api"
 	"github.com/google/gapid/gapis/memory"
-	"github.com/google/gapid/gapis/replay/value"
 )
 
 type stateBuilder struct {
-	ctx           context.Context
-	s             *State
-	oldState      *api.GlobalState
-	newState      *api.GlobalState
-	cmds          []api.Cmd
-	cb            CommandBuilder
-	readMemories  []*api.AllocResult
-	writeMemories []*api.AllocResult
+	ctx             context.Context
+	s               *State
+	oldState        *api.GlobalState
+	newState        *api.GlobalState
+	cmds            []api.Cmd
+	cb              CommandBuilder
+	readMemories    []*api.AllocResult
+	writeMemories   []*api.AllocResult
+	memoryIntervals interval.U64RangeList
 }
 
 // TODO: wherever possible, use old resources instead of doing full reads on the old pools.
 //       This is especially useful for things that are internal pools, (Shader words for example)
-func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) []api.Cmd {
+func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) ([]api.Cmd, interval.U64RangeList) {
 	// TODO: Debug Info
-	newState := api.NewStateWithAllocator(memory.NewBasicAllocator(value.ValidMemoryRanges), oldState.MemoryLayout)
+	newState := api.NewStateWithAllocator(memory.NewBasicAllocator(oldState.Allocator.FreeList()), oldState.MemoryLayout)
 	sb := &stateBuilder{
-		ctx:      ctx,
-		s:        s,
-		oldState: oldState,
-		newState: newState,
-		cb:       CommandBuilder{Thread: 0},
+		ctx:             ctx,
+		s:               s,
+		oldState:        oldState,
+		newState:        newState,
+		cb:              CommandBuilder{Thread: 0},
+		memoryIntervals: interval.U64RangeList{},
 	}
 	sb.newState.Memory.NewAt(sb.oldState.Memory.NextPoolID())
 
@@ -166,7 +168,7 @@ func (s *State) RebuildState(ctx context.Context, oldState *api.GlobalState) []a
 		sb.createCommandBuffer(s.CommandBuffers.Get(qp), VkCommandBufferLevel_VK_COMMAND_BUFFER_LEVEL_PRIMARY)
 	}
 
-	return sb.cmds
+	return sb.cmds, sb.memoryIntervals
 }
 
 func GetPipelinesInOrder(s *State, compute bool) []VkPipeline {
@@ -220,24 +222,32 @@ func GetPipelinesInOrder(s *State, compute bool) []VkPipeline {
 func (sb *stateBuilder) MustAllocReadData(v ...interface{}) api.AllocResult {
 	allocate_result := sb.newState.AllocDataOrPanic(sb.ctx, v...)
 	sb.readMemories = append(sb.readMemories, &allocate_result)
+	rng := allocate_result.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	return allocate_result
 }
 
 func (sb *stateBuilder) MustAllocWriteData(v ...interface{}) api.AllocResult {
 	allocate_result := sb.newState.AllocDataOrPanic(sb.ctx, v...)
 	sb.writeMemories = append(sb.writeMemories, &allocate_result)
+	rng := allocate_result.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	return allocate_result
 }
 
 func (sb *stateBuilder) MustUnpackReadMap(v interface{}) api.AllocResult {
 	allocate_result, _ := unpackMap(sb.ctx, sb.newState, v)
 	sb.readMemories = append(sb.readMemories, &allocate_result)
+	rng := allocate_result.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	return allocate_result
 }
 
 func (sb *stateBuilder) MustUnpackWriteMap(v interface{}) api.AllocResult {
 	allocate_result, _ := unpackMap(sb.ctx, sb.newState, v)
 	sb.writeMemories = append(sb.writeMemories, &allocate_result)
+	rng := allocate_result.Range()
+	interval.Merge(&sb.memoryIntervals, interval.U64Span{rng.Base, rng.Base + rng.Size}, true)
 	return allocate_result
 }
 
@@ -549,13 +559,10 @@ func (sb *stateBuilder) createDevice(d *DeviceObject) {
 	reordered_queue_creates := map[uint32]VkDeviceQueueCreateInfo{}
 	i := uint32(0)
 	for k, v := range queue_create {
-		log.E(sb.ctx, "Initial queue_creates %+v", queue_priorities[k])
 		v.PQueuePriorities = NewF32ᶜᵖ(sb.MustAllocReadData(queue_priorities[k]).Ptr())
 		reordered_queue_creates[i] = v
 		i++
 	}
-
-	log.E(sb.ctx, "Queues %+v", reordered_queue_creates)
 
 	sb.write(sb.cb.VkCreateDevice(
 		d.PhysicalDevice,

gapis/capture/capture.go

@@ -117,8 +117,11 @@ func NewState(ctx context.Context) (*api.GlobalState, error) {
 }
 
 // NewUninitializedState returns a new, uninitialized State object built for the capture.
-func (c *Capture) NewUninitializedState(ctx context.Context) *api.GlobalState {
+func (c *Capture) NewUninitializedState(ctx context.Context, rngs interval.U64RangeList) *api.GlobalState {
 	freeList := memory.InvertMemoryRanges(c.Observed)
+	for _, r := range rngs {
+		interval.Remove(&freeList, r.Span())
+	}
 	interval.Remove(&freeList, interval.U64Span{Start: 0, End: value.FirstValidAddress})
 	s := api.NewStateWithAllocator(
 		memory.NewBasicAllocator(freeList),
@@ -130,7 +133,7 @@ func (c *Capture) NewUninitializedState(ctx context.Context) *api.GlobalState {
 // NewState returns a new, default-initialized State object built for the
 // capture.
 func (c *Capture) NewState(ctx context.Context) *api.GlobalState {
-	s := c.NewUninitializedState(ctx)
+	s := c.NewUninitializedState(ctx, interval.U64RangeList{})
 	if c.InitialState != nil {
 		for _, m := range c.InitialState.Memory {
 			pool, _ := s.Memory.Get(memory.PoolID(m.Pool))
@@ -150,15 +153,20 @@ func (c *Capture) NewState(ctx context.Context) *api.GlobalState {
 }
 
 // GetInitialCommands returns a set of commands which will setup the initial state.
-func (c *Capture) GetInitialCommands(ctx context.Context) (cmds []api.Cmd) {
+func (c *Capture) GetInitialCommands(ctx context.Context) ([]api.Cmd, interval.U64RangeList) {
+	ranges := interval.U64RangeList{}
+	cmds := []api.Cmd{}
 	// TODO: This can be easily cached for the given capture.
 	if c.InitialState != nil {
 		s := c.NewState(ctx)
 		for _, v := range s.APIs {
-			cmds = append(cmds, v.RebuildState(ctx, s)...)
+			s, r := v.RebuildState(ctx, s)
+			ranges = append(ranges, r...)
+			cmds = append(cmds, s...)
 		}
+		return cmds, ranges
 	}
-	return
+	return nil, interval.U64RangeList{}
 }
 
 // Service returns the service.Capture description for this capture.

gapis/replay/batch.go

@@ -137,8 +137,10 @@ func (m *Manager) execute(
 
 	b := builder.New(replayABI.MemoryLayout)
 
+	_, ranges := c.GetInitialCommands(ctx)
+
 	out := &adapter{
-		state:   c.NewUninitializedState(ctx),
+		state:   c.NewUninitializedState(ctx, ranges),
 		builder: b,
 	}
 

gapis/resolve/dependencygraph/dependency_graph.go

@@ -163,7 +163,7 @@ func (r *DependencyGraphResolvable) Resolve(ctx context.Context) (interface{}, e
 	behaviourProviders := map[api.API]BehaviourProvider{}
 
 	// If the capture contains initial state, prepend the commands to build the state.
-	initialCmds := c.GetInitialCommands(ctx)
+	initialCmds, ranges := c.GetInitialCommands(ctx)
 	if len(initialCmds) > 0 {
 		cmds = append(initialCmds, cmds...)
 	}
@@ -180,7 +180,8 @@ func (r *DependencyGraphResolvable) Resolve(ctx context.Context) (interface{}, e
 		},
 	}
 
-	s := c.NewUninitializedState(ctx)
+	s := c.NewUninitializedState(ctx, ranges)
+
 	dependencyGraphBuildCounter.Time(func() {
 		api.ForeachCmd(ctx, cmds, func(ctx context.Context, index api.CmdID, cmd api.Cmd) error {
 			a := cmd.API()