5 files changed, 49 insertions, 86 deletions
diff --git a/prog/prog.go b/prog/prog.go
index 69c76c019..575680882 100644
--- a/prog/prog.go
+++ b/prog/prog.go
@@ -20,14 +20,6 @@ type Call struct {
 	Comment string
 }
 
-func (p *Prog) Finalize() error {
-	for _, c := range p.Calls {
-		p.Target.assignSizesCall(c)
-		p.Target.SanitizeCall(c)
-	}
-	return p.validate()
-}
-
 type Arg interface {
 	Type() Type
 	Size() uint64
diff --git a/prog/target.go b/prog/target.go
index db8b1835e..f99c45c5b 100644
--- a/prog/target.go
+++ b/prog/target.go
@@ -236,3 +236,42 @@ func (g *Gen) MutateArg(arg0 Arg) (calls []*Call) {
 	}
 	return calls
 }
+
+type ProgGen struct {
+	target *Target
+	ma     *memAlloc
+	p      *Prog
+}
+
+func MakeProgGen(target *Target) *ProgGen {
+	return &ProgGen{
+		target: target,
+		ma:     newMemAlloc(target.NumPages * target.PageSize),
+		p: &Prog{
+			Target: target,
+		},
+	}
+}
+
+func (pg *ProgGen) Append(c *Call) error {
+	pg.target.assignSizesCall(c)
+	pg.target.SanitizeCall(c)
+	pg.p.Calls = append(pg.p.Calls, c)
+	return nil
+}
+
+func (pg *ProgGen) Allocate(size uint64) uint64 {
+	return pg.ma.alloc(nil, size)
+}
+
+func (pg *ProgGen) Finalize() (*Prog, error) {
+	if err := pg.p.validate(); err != nil {
+		return nil, err
+	}
+	if _, err := pg.p.SerializeForExec(make([]byte, ExecBufferSize)); err != nil {
+		return nil, err
+	}
+	p := pg.p
+	pg.p = nil
+	return p, nil
+}
diff --git a/tools/syz-trace2syz/proggen/context.go b/tools/syz-trace2syz/proggen/context.go
index f81b70674..8283cfd8f 100644
--- a/tools/syz-trace2syz/proggen/context.go
+++ b/tools/syz-trace2syz/proggen/context.go
@@ -10,24 +10,20 @@ import (
 
 // Context stores metadata related to a syzkaller program
 type Context struct {
+	pg                *prog.ProgGen
 	ReturnCache       returnCache
-	Prog              *prog.Prog
 	CurrentStraceCall *parser.Syscall
 	CurrentSyzCall    *prog.Call
 	CurrentStraceArg  parser.IrType
 	Target            *prog.Target
-	Tracker           *memoryTracker
 	callSelector      *callSelector
 }
 
 func newContext(target *prog.Target) *Context {
 	return &Context{
+		pg:           prog.MakeProgGen(target),
 		ReturnCache:  newRCache(),
-		Tracker:      newTracker(),
 		Target:       target,
 		callSelector: newCallSelector(),
-		Prog: &prog.Prog{
-			Target: target,
-		},
 	}
 }
diff --git a/tools/syz-trace2syz/proggen/memory_tracker.go b/tools/syz-trace2syz/proggen/memory_tracker.go
deleted file mode 100644
index 5e6c6ddb8..000000000
--- a/tools/syz-trace2syz/proggen/memory_tracker.go
+++ /dev/null
@@ -1,55 +0,0 @@
-// Copyright 2018 syzkaller project authors. All rights reserved.
-// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
-
-package proggen
-
-import (
-	"fmt"
-
-	"github.com/google/syzkaller/prog"
-)
-
-const (
-	memAllocMaxMem = 16 << 20
-)
-
-type allocation struct {
-	numBytes uint64
-	arg      *prog.PointerArg
-}
-
-// TODO: Replace memory tracker with memAlloc in prog package.
-type memoryTracker struct {
-	allocations map[*prog.Call][]*allocation
-}
-
-func newTracker() *memoryTracker {
-	return &memoryTracker{
-		allocations: make(map[*prog.Call][]*allocation),
-	}
-}
-
-func (m *memoryTracker) addAllocation(call *prog.Call, size uint64, arg *prog.PointerArg) {
-	m.allocations[call] = append(m.allocations[call], &allocation{
-		arg:      arg,
-		numBytes: size,
-	})
-}
-
-func (m *memoryTracker) fillOutPtrArgs(p *prog.Prog) error {
-	var offset uint64
-	for _, call := range p.Calls {
-		for _, a := range m.allocations[call] {
-			a.arg.Address = offset
-			offset += a.numBytes
-
-			if a.arg.Address >= memAllocMaxMem {
-				return fmt.Errorf("unable to allocate space to store arg: %#v"+
-					"in Call: %v. Required memory is larger than what we allow."+
-					"Offending address: %v",
-					a.arg, call, a.arg.Address)
-			}
-		}
-	}
-	return nil
-}
diff --git a/tools/syz-trace2syz/proggen/proggen.go b/tools/syz-trace2syz/proggen/proggen.go
index 66a0cf842..967167fe9 100644
--- a/tools/syz-trace2syz/proggen/proggen.go
+++ b/tools/syz-trace2syz/proggen/proggen.go
@@ -70,20 +70,15 @@ func genProg(trace *parser.Trace, target *prog.Target) *prog.Prog {
 		if call == nil {
 			continue
 		}
-		ctx.Prog.Calls = append(ctx.Prog.Calls, call)
-	}
-	if err := ctx.Tracker.fillOutPtrArgs(ctx.Prog); err != nil {
-		log.Logf(1, "failed to fill out memory: %v, skipping this prog", err)
-		return nil
+		if err := ctx.pg.Append(call); err != nil {
+			log.Fatalf("%v", err)
+		}
 	}
-	if err := ctx.Prog.Finalize(); err != nil {
+	p, err := ctx.pg.Finalize()
+	if err != nil {
 		log.Fatalf("error validating program: %v", err)
 	}
-	if _, err := ctx.Prog.SerializeForExec(make([]byte, prog.ExecBufferSize)); err != nil {
-		log.Logf(1, "prog is too large")
-		return nil
-	}
-	return ctx.Prog
+	return p
 }
 
 func genCall(ctx *Context) *prog.Call {
@@ -172,9 +167,7 @@ func genVma(syzType *prog.VmaType, _ parser.IrType, ctx *Context) prog.Arg {
 	if syzType.RangeBegin != 0 || syzType.RangeEnd != 0 {
 		npages = syzType.RangeEnd
 	}
-	arg := prog.MakeVmaPointerArg(syzType, 0, npages)
-	ctx.Tracker.addAllocation(ctx.CurrentSyzCall, ctx.Target.PageSize, arg)
-	return arg
+	return prog.MakeVmaPointerArg(syzType, ctx.pg.Allocate(ctx.Target.PageSize), npages)
 }
 
 func genArray(syzType *prog.ArrayType, traceType parser.IrType, ctx *Context) prog.Arg {
@@ -380,9 +373,7 @@ func parseProc(syzType *prog.ProcType, traceType parser.IrType, ctx *Context) pr
 }
 
 func addr(ctx *Context, syzType prog.Type, size uint64, data prog.Arg) prog.Arg {
-	arg := prog.MakePointerArg(syzType, uint64(0), data)
-	ctx.Tracker.addAllocation(ctx.CurrentSyzCall, size, arg)
-	return arg
+	return prog.MakePointerArg(syzType, ctx.pg.Allocate(size), data)
 }
 
 func reorderStructFields(syzType *prog.StructType, traceType *parser.GroupType, ctx *Context) {