Optimize neotest in coverage mode

pkg/neotest/basic.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math/big"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/pkg/config/netmode"
@@ -32,8 +33,16 @@ type Executor struct {
 	Validator     Signer
 	Committee     Signer
 	CommitteeHash util.Uint160
+
 	// collectCoverage is true if coverage is being collected when running this executor.
 	collectCoverage bool
+	// scheduleCoverageReport is a cleanup function that schedules the report of
+	// collected coverage data for this Executor.
+	scheduleCoverageReport     func()
+	scheduleCoverageReportOnce sync.Once
+
+	coverageLock sync.RWMutex
+	rawCoverage  map[util.Uint160]*scriptRawCoverage
 }
 
 // NewExecutor creates a new executor instance from the provided blockchain and committee.
@@ -43,13 +52,20 @@ func NewExecutor(t testing.TB, bc *core.Blockchain, validator, committee Signer)
 	checkMultiSigner(t, validator)
 	checkMultiSigner(t, committee)
 
-	return &Executor{
+	e := &Executor{
 		Chain:           bc,
 		Validator:       validator,
 		Committee:       committee,
 		CommitteeHash:   committee.ScriptHash(),
 		collectCoverage: isCoverageEnabled(t),
+		rawCoverage:     make(map[util.Uint160]*scriptRawCoverage),
 	}
+
+	if e.collectCoverage {
+		e.scheduleCoverageReport = func() { t.Cleanup(func() { e.reportCoverage(t) }) }
+	}
+
+	return e
 }
 
 // TopBlock returns the block with the highest index.
@@ -148,7 +164,7 @@ func (e *Executor) DeployContract(t testing.TB, c *Contract, data any) util.Uint
 // data is an optional argument to `_deploy`.
 // It returns the hash of the deploy transaction.
 func (e *Executor) DeployContractBy(t testing.TB, signer Signer, c *Contract, data any) util.Uint256 {
-	e.trackCoverage(t, c)
+	e.trackCoverage(c)
 	tx := e.NewDeployTxBy(t, signer, c, data)
 	e.AddNewBlock(t, tx)
 	e.CheckHalt(t, tx.Hash())
@@ -168,19 +184,17 @@ func (e *Executor) DeployContractBy(t testing.TB, signer Signer, c *Contract, da
 // DeployContractCheckFAULT compiles and deploys a contract to the bc using the validator
 // account. It checks that the deploy transaction FAULTed with the specified error.
 func (e *Executor) DeployContractCheckFAULT(t testing.TB, c *Contract, data any, errMessage string) {
-	e.trackCoverage(t, c)
+	e.trackCoverage(c)
 	tx := e.NewDeployTx(t, c, data)
 	e.AddNewBlock(t, tx)
 	e.CheckFault(t, tx.Hash(), errMessage)
 }
 
 // trackCoverage switches on coverage tracking for provided script if `go test` is running with coverage enabled.
-func (e *Executor) trackCoverage(t testing.TB, c *Contract) {
+func (e *Executor) trackCoverage(c *Contract) {
 	if e.collectCoverage {
 		addScriptToCoverage(c)
-		t.Cleanup(func() {
-			reportCoverage(t)
-		})
+		e.scheduleCoverageReportOnce.Do(e.scheduleCoverageReport)
 	}
 }
 
@@ -417,7 +431,7 @@ func (e *Executor) TestInvoke(tx *transaction.Transaction) (*vm.VM, error) {
 	ic, _ := e.Chain.GetTestVM(trigger.Application, &ttx, b)
 
 	if e.collectCoverage {
-		ic.VM.SetOnExecHook(coverageHook)
+		ic.VM.SetOnExecHook(e.coverageHook)
 	}
 
 	defer ic.Finalize()

pkg/neotest/client.go

@@ -64,7 +64,7 @@ func (c *ContractInvoker) TestInvokeScript(t testing.TB, script []byte, signers
 	t.Cleanup(ic.Finalize)
 
 	if c.collectCoverage {
-		ic.VM.SetOnExecHook(coverageHook)
+		ic.VM.SetOnExecHook(c.coverageHook)
 	}
 
 	ic.VM.LoadWithFlags(tx.Script, callflag.All)
@@ -83,7 +83,7 @@ func (c *ContractInvoker) TestInvoke(t testing.TB, method string, args ...any) (
 	t.Cleanup(ic.Finalize)
 
 	if c.collectCoverage {
-		ic.VM.SetOnExecHook(coverageHook)
+		ic.VM.SetOnExecHook(c.coverageHook)
 	}
 
 	ic.VM.LoadWithFlags(tx.Script, callflag.All)

pkg/neotest/coverage.go

@@ -13,7 +13,6 @@ import (
 
 	"github.com/nspcc-dev/neo-go/pkg/compiler"
 	"github.com/nspcc-dev/neo-go/pkg/util"
-	"github.com/nspcc-dev/neo-go/pkg/vm"
 	"github.com/nspcc-dev/neo-go/pkg/vm/opcode"
 )
 
@@ -118,17 +117,32 @@ func isCoverageEnabled(t testing.TB) bool {
 	return coverageEnabled
 }
 
-var coverageHook vm.OnExecHook = func(scriptHash util.Uint160, offset int, opcode opcode.Opcode) {
-	coverageLock.Lock()
-	defer coverageLock.Unlock()
-	if cov, ok := rawCoverage[scriptHash]; ok {
+func (e *Executor) coverageHook(scriptHash util.Uint160, offset int, _ opcode.Opcode) {
+	e.coverageLock.Lock()
+	defer e.coverageLock.Unlock()
+	if cov, ok := e.rawCoverage[scriptHash]; ok {
 		cov.offsetsVisited = append(cov.offsetsVisited, offset)
 	}
 }
 
-func reportCoverage(t testing.TB) {
+func (e *Executor) reportCoverage(t testing.TB) {
 	coverageLock.Lock()
 	defer coverageLock.Unlock()
+
+	e.coverageLock.RLock()
+	for h, cov := range e.rawCoverage {
+		fullCov := rawCoverage[h]
+		if fullCov == nil {
+			fullCov = &scriptRawCoverage{
+				debugInfo: cov.debugInfo,
+			}
+			rawCoverage[h] = fullCov
+		}
+
+		fullCov.offsetsVisited = append(fullCov.offsetsVisited, cov.offsetsVisited...)
+	}
+	e.coverageLock.RUnlock()
+
 	f, err := os.Create(coverProfile)
 	if err != nil {
 		t.Fatalf("coverage: can't create file '%s' to write coverage report", coverProfile)