tools/syz-trace2syz/proggen: replace memoryTracker with prog.memAlloc

2 files changed, 39 insertions, 8 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
+}