prog: fix a data race

The race initially showed up on the new benchmark (see race report below).
The race indicated a wrong call passed to replaceArg,
as the result we sanitized the wrong call and left the new call un-sanitized.

Fix this.
Add test that exposes this.
Run benchmarks in race mode during presubmit
(benchmarks have higher chances of uncovering races than tests).

WARNING: DATA RACE
Write at 0x00c42000d3f0 by goroutine 18:
  github.com/google/syzkaller/sys/linux.(*arch).sanitizeCall()
      sys/linux/init.go:155 +0x256
  github.com/google/syzkaller/sys/linux.(*arch).(github.com/google/syzkaller/sys/linux.sanitizeCall)-fm()
      sys/linux/init.go:42 +0x4b
  github.com/google/syzkaller/prog.(*Prog).replaceArg()
      prog/prog.go:357 +0x239
  github.com/google/syzkaller/prog.generateHints.func2()
      prog/hints.go:105 +0x124
  github.com/google/syzkaller/prog.checkConstArg()
      prog/hints.go:128 +0xf3
  github.com/google/syzkaller/prog.generateHints()
      prog/hints.go:120 +0x495
  github.com/google/syzkaller/prog.(*Prog).MutateWithHints.func1()
      prog/hints.go:72 +0x67
  github.com/google/syzkaller/prog.foreachSubargImpl.func1()
      prog/analysis.go:86 +0x9f
  github.com/google/syzkaller/prog.foreachSubargImpl()
      prog/analysis.go:104 +0xc8
  github.com/google/syzkaller/prog.foreachArgArray()
      prog/analysis.go:113 +0x89
  github.com/google/syzkaller/prog.foreachArg()
      prog/analysis.go:121 +0x50
  github.com/google/syzkaller/prog.(*Prog).MutateWithHints()
      prog/hints.go:71 +0x18e
  github.com/google/syzkaller/prog.BenchmarkHints.func1()
      prog/hints_test.go:477 +0x77
  testing.(*B).RunParallel.func1()
      testing/benchmark.go:626 +0x156

Previous read at 0x00c42000d3f0 by goroutine 17:
  github.com/google/syzkaller/prog.clone()
      prog/clone.go:38 +0xbaa
  github.com/google/syzkaller/prog.(*Prog).cloneImpl()
      prog/clone.go:21 +0x17f
  github.com/google/syzkaller/prog.generateHints()
      prog/hints.go:95 +0xd0
  github.com/google/syzkaller/prog.(*Prog).MutateWithHints.func1()
      prog/hints.go:72 +0x67
  github.com/google/syzkaller/prog.foreachSubargImpl.func1()
      prog/analysis.go:86 +0x9f
  github.com/google/syzkaller/prog.foreachSubargImpl()
      prog/analysis.go:104 +0xc8
  github.com/google/syzkaller/prog.foreachArgArray()
      prog/analysis.go:113 +0x89
  github.com/google/syzkaller/prog.foreachArg()
      prog/analysis.go:121 +0x50
  github.com/google/syzkaller/prog.(*Prog).MutateWithHints()
      prog/hints.go:71 +0x18e
  github.com/google/syzkaller/prog.BenchmarkHints.func1()
      prog/hints_test.go:477 +0x77
  testing.(*B).RunParallel.func1()
      testing/benchmark.go:626 +0x156

4 files changed, 98 insertions, 22 deletions

diff --git a/Makefile b/Makefile
index 4ef3b4585..5ea34c242 100644
--- a/Makefile
+++ b/Makefile
@@ -219,7 +219,7 @@ tidy:
 
 test:
 	$(GO) test -short ./...
-	$(GO) test -short -race ./...
+	$(GO) test -short -race -bench=.* ./...
 
 arch:
 	env GOOG=darwin GOARCH=amd64 go install github.com/google/syzkaller/syz-manager
diff --git a/prog/hints.go b/prog/hints.go
index 065e9f8ba..dabb722da 100644
--- a/prog/hints.go
+++ b/prog/hints.go
@@ -69,11 +69,11 @@ func (p *Prog) MutateWithHints(callIndex int, comps CompMap, exec func(newP *Pro
 		return
 	}
 	foreachArg(c, func(arg, _ Arg, _ *[]Arg) {
-		generateHints(p, comps, c, arg, exec)
+		generateHints(p, comps, callIndex, arg, exec)
 	})
 }
 
-func generateHints(p *Prog, compMap CompMap, c *Call, arg Arg, exec func(p *Prog)) {
+func generateHints(p *Prog, compMap CompMap, callIndex int, arg Arg, exec func(p *Prog)) {
 	if arg.Type().Dir() == DirOut {
 		return
 	}
@@ -93,6 +93,7 @@ func generateHints(p *Prog, compMap CompMap, c *Call, arg Arg, exec func(p *Prog
 	}
 
 	newP, argMap := p.cloneImpl(true)
+	newCall := newP.Calls[callIndex]
 	validateExec := func() {
 		if err := newP.validate(); err != nil {
 			panic(fmt.Sprintf("invalid hints candidate: %v", err))
@@ -102,9 +103,9 @@ func generateHints(p *Prog, compMap CompMap, c *Call, arg Arg, exec func(p *Prog
 	var originalArg Arg
 	constArgCandidate := func(newArg Arg) {
 		oldArg := argMap[arg]
-		newP.replaceArg(c, oldArg, newArg, nil)
+		newP.replaceArg(newCall, oldArg, newArg, nil)
 		validateExec()
-		newP.replaceArg(c, oldArg, originalArg, nil)
+		newP.replaceArg(newCall, oldArg, originalArg, nil)
 	}
 
 	dataArgCandidate := func() {
diff --git a/prog/hints_test.go b/prog/hints_test.go
index b5e35c9c9..b2999ad14 100644
--- a/prog/hints_test.go
+++ b/prog/hints_test.go
@@ -6,6 +6,7 @@ package prog
 import (
 	"encoding/hex"
 	"fmt"
+	"math/rand"
 	"reflect"
 	"sort"
 	"testing"
@@ -450,3 +451,31 @@ func TestHintsData(t *testing.T) {
 		}
 	}
 }
+
+func BenchmarkHints(b *testing.B) {
+	target, err := GetTarget("linux", "amd64")
+	if err != nil {
+		b.Fatal(err)
+	}
+	rs := rand.NewSource(0)
+	r := newRand(target, rs)
+	p := target.Generate(rs, 30, nil)
+	comps := make([]CompMap, len(p.Calls))
+	for i, c := range p.Calls {
+		vals := extractValues(c)
+		for j := 0; j < 5; j++ {
+			vals[r.randInt()] = true
+		}
+		comps[i] = make(CompMap)
+		for v := range vals {
+			comps[i].AddComp(v, r.randInt())
+		}
+	}
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			for i := range p.Calls {
+				p.MutateWithHints(i, comps[i], func(p1 *Prog) {})
+			}
+		}
+	})
+}
diff --git a/prog/prog.go b/prog/prog.go
index 31241c570..913c62a7d 100644
--- a/prog/prog.go
+++ b/prog/prog.go
@@ -324,6 +324,9 @@ func (p *Prog) insertBefore(c *Call, calls []*Call) {
 
 // replaceArg replaces arg with arg1 in call c in program p, and inserts calls before arg call.
 func (p *Prog) replaceArg(c *Call, arg, arg1 Arg, calls []*Call) {
+	if debug {
+		p.replaceArgCheck(c, arg, arg1, calls)
+	}
 	for _, c := range calls {
 		p.Target.SanitizeCall(c)
 	}
@@ -332,19 +335,7 @@ func (p *Prog) replaceArg(c *Call, arg, arg1 Arg, calls []*Call) {
 	case *ConstArg:
 		*a = *arg1.(*ConstArg)
 	case *ResultArg:
-		// Remove link from `a.Res` to `arg`.
-		if a.Res != nil {
-			delete(*a.Res.(ArgUsed).Used(), arg)
-		}
-		// Copy all fields from `arg1` to `arg` except for the list of args that use `arg`.
-		used := *arg.(ArgUsed).Used()
-		*a = *arg1.(*ResultArg)
-		*arg.(ArgUsed).Used() = used
-		// Make the link in `a.Res` (which is now `Res` of `arg1`) to point to `arg` instead of `arg1`.
-		if a.Res != nil {
-			delete(*a.Res.(ArgUsed).Used(), arg1)
-			(*a.Res.(ArgUsed).Used())[arg] = true
-		}
+		replaceResultArg(a, arg1.(*ResultArg))
 	case *PointerArg:
 		*a = *arg1.(*PointerArg)
 	case *UnionArg:
@@ -357,6 +348,60 @@ func (p *Prog) replaceArg(c *Call, arg, arg1 Arg, calls []*Call) {
 	p.Target.SanitizeCall(c)
 }
 
+func replaceResultArg(arg, arg1 *ResultArg) {
+	// Remove link from `a.Res` to `arg`.
+	if arg.Res != nil {
+		delete(*arg.Res.(ArgUsed).Used(), arg)
+	}
+	// Copy all fields from `arg1` to `arg` except for the list of args that use `arg`.
+	used := *arg.Used()
+	*arg = *arg1
+	*arg.Used() = used
+	// Make the link in `arg.Res` (which is now `Res` of `arg1`) to point to `arg` instead of `arg1`.
+	if arg.Res != nil {
+		delete(*arg.Res.(ArgUsed).Used(), arg1)
+		(*arg.Res.(ArgUsed).Used())[arg] = true
+	}
+}
+
+// replaceArgCheck checks that c and arg belog to p.
+func (p *Prog) replaceArgCheck(c *Call, arg, arg1 Arg, calls []*Call) {
+	foundCall, foundArg := false, false
+	for _, c0 := range p.Calls {
+		if c0 == c {
+			if foundCall {
+				panic("duplicate call")
+			}
+			foundCall = true
+		}
+		for _, newC := range calls {
+			if c0 == newC {
+				panic("call is already in prog")
+			}
+		}
+		foreachArg(c0, func(arg0, _ Arg, _ *[]Arg) {
+			if arg0 == arg {
+				if c0 != c {
+					panic("wrong call")
+				}
+				if foundArg {
+					panic("duplicate arg")
+				}
+				foundArg = true
+			}
+			if arg0 == arg1 {
+				panic("arg is already in prog")
+			}
+		})
+	}
+	if !foundCall {
+		panic("call is not in prog")
+	}
+	if !foundArg {
+		panic("arg is not in prog")
+	}
+}
+
 // removeArg removes all references to/from arg0 of call c from p.
 func (p *Prog) removeArg(c *Call, arg0 Arg) {
 	foreachSubarg(arg0, func(arg, _ Arg, _ *[]Arg) {
@@ -368,11 +413,12 @@ func (p *Prog) removeArg(c *Call, arg0 Arg) {
 		}
 		if used, ok := arg.(ArgUsed); ok {
 			for arg1 := range *used.Used() {
-				if _, ok := arg1.(*ResultArg); !ok {
+				a1, ok := arg1.(*ResultArg)
+				if !ok {
 					panic("use references not ArgResult")
 				}
 				arg2 := MakeResultArg(arg1.Type(), nil, arg1.Type().Default())
-				p.replaceArg(c, arg1, arg2, nil)
+				replaceResultArg(a1, arg2.(*ResultArg))
 			}
 		}
 	})
@@ -381,10 +427,10 @@ func (p *Prog) removeArg(c *Call, arg0 Arg) {
 // removeCall removes call idx from p.
 func (p *Prog) removeCall(idx int) {
 	c := p.Calls[idx]
-	copy(p.Calls[idx:], p.Calls[idx+1:])
-	p.Calls = p.Calls[:len(p.Calls)-1]
 	for _, arg := range c.Args {
 		p.removeArg(c, arg)
 	}
 	p.removeArg(c, c.Ret)
+	copy(p.Calls[idx:], p.Calls[idx+1:])
+	p.Calls = p.Calls[:len(p.Calls)-1]
 }