add coverage support to neotest

pkg/neotest/basic.go

@@ -401,6 +401,10 @@ func TestInvoke(bc *core.Blockchain, tx *transaction.Transaction) (*vm.VM, error
 	ttx := *tx
 	ic, _ := bc.GetTestVM(trigger.Application, &ttx, b)
 
+	if isCoverageEnabled() {
+		ic.VM.SetOnExecHook(coverageHook())
+	}
+
 	defer ic.Finalize()
 
 	ic.VM.LoadWithFlags(tx.Script, callflag.All)

pkg/neotest/compile.go

@@ -35,11 +35,15 @@ func CompileSource(t testing.TB, sender util.Uint160, src io.Reader, opts *compi
 	m, err := compiler.CreateManifest(di, opts)
 	require.NoError(t, err)
 
-	return &Contract{
+	c := Contract{
 		Hash:     state.CreateContractHash(sender, ne.Checksum, m.Name),
 		NEF:      ne,
 		Manifest: m,
 	}
+
+	collectCoverage(t, di, c.Hash)
+
+	return &c
 }
 
 // CompileFile compiles a contract from the file and returns its NEF, manifest and hash.
@@ -77,6 +81,18 @@ func CompileFile(t testing.TB, sender util.Uint160, srcPath string, configPath s
 		NEF:      ne,
 		Manifest: m,
 	}
+
+	collectCoverage(t, di, c.Hash)
+
 	contracts[srcPath] = c
 	return c
 }
+
+func collectCoverage(t testing.TB, di *compiler.DebugInfo, h scriptHash) {
+	if isCoverageEnabled() {
+		rawCoverage[h] = &scriptRawCoverage{debugInfo: di}
+		t.Cleanup(func() {
+			reportCoverage()
+		})
+	}
+}

pkg/neotest/coverage.go

@@ -0,0 +1,153 @@
+package neotest
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"os"
+
+	"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"
+)
+
+var rawCoverage = make(map[scriptHash]*scriptRawCoverage)
+
+var enabled = false
+var coverProfile = ""
+
+type scriptHash = util.Uint160
+
+type scriptRawCoverage struct {
+	debugInfo      *compiler.DebugInfo
+	offsetsVisited []int
+}
+
+type coverBlock struct {
+	startLine uint // Line number for block start.
+	startCol  uint // Column number for block start.
+	endLine   uint // Line number for block end.
+	endCol    uint // Column number for block end.
+	stmts     uint // Number of statements included in this block.
+	counts    uint
+}
+
+type documentName = string
+
+func isCoverageEnabled() bool {
+	if enabled {
+		return true
+	}
+	const coverProfileFlag = "test.coverprofile"
+	flag.VisitAll(func(f *flag.Flag) {
+		if f.Name == coverProfileFlag && f.Value != nil {
+			enabled = true
+			coverProfile = f.Value.String()
+		}
+	})
+	if enabled {
+		// this is needed so go cover tool doesn't overwrite
+		// the file with our coverage when all tests are done
+		flag.Set(coverProfileFlag, "")
+	}
+	return enabled
+}
+
+func coverageHook() vm.OnExecHook {
+	return func(sh scriptHash, offset int, opcode opcode.Opcode) {
+		if cov, ok := rawCoverage[sh]; ok {
+			cov.offsetsVisited = append(cov.offsetsVisited, offset)
+		}
+	}
+}
+
+func reportCoverage() {
+	f, _ := os.Create(coverProfile)
+	defer f.Close()
+	writeCoverageReport(f)
+}
+
+func writeCoverageReport(w io.Writer) {
+	fmt.Fprintf(w, "mode: set\n") // TODO: other mods
+	cover := processCover()
+	for name, blocks := range cover {
+		for _, b := range blocks {
+			c := 0
+			if b.counts > 0 {
+				c = 1
+			}
+			fmt.Fprintf(w, "%s:%d.%d,%d.%d %d %d\n", name,
+				b.startLine, b.startCol,
+				b.endLine, b.endCol,
+				b.stmts,
+				c,
+			)
+		}
+	}
+}
+
+func processCover() map[documentName][]coverBlock {
+	documents := make(map[documentName]struct{})
+	for _, scriptRawCoverage := range rawCoverage {
+		for _, documentName := range scriptRawCoverage.debugInfo.Documents {
+			documents[documentName] = struct{}{}
+		}
+	}
+
+	cover := make(map[documentName][]coverBlock)
+
+	for documentName := range documents {
+		mappedBlocks := make(map[int]*coverBlock)
+
+		for _, scriptRawCoverage := range rawCoverage {
+			di := scriptRawCoverage.debugInfo
+			documentSeqPoints := documentSeqPoints(di, documentName)
+
+			for _, point := range documentSeqPoints {
+				b := coverBlock{
+					startLine: uint(point.StartLine),
+					startCol:  uint(point.StartCol),
+					endLine:   uint(point.EndLine),
+					endCol:    uint(point.EndCol),
+					stmts:     1 + uint(point.EndLine) - uint(point.StartLine),
+					counts:    0,
+				}
+				mappedBlocks[point.Opcode] = &b
+			}
+		}
+
+		for _, scriptRawCoverage := range rawCoverage {
+			di := scriptRawCoverage.debugInfo
+			documentSeqPoints := documentSeqPoints(di, documentName)
+
+			for _, offset := range scriptRawCoverage.offsetsVisited {
+				for _, point := range documentSeqPoints {
+					if point.Opcode == offset {
+						mappedBlocks[point.Opcode].counts++
+					}
+				}
+			}
+		}
+
+		var blocks []coverBlock
+		for _, b := range mappedBlocks {
+			blocks = append(blocks, *b)
+		}
+		cover[documentName] = blocks
+	}
+
+	return cover
+}
+
+func documentSeqPoints(di *compiler.DebugInfo, doc documentName) []compiler.DebugSeqPoint {
+	var res []compiler.DebugSeqPoint
+	for _, methodDebugInfo := range di.Methods {
+		for _, p := range methodDebugInfo.SeqPoints {
+			if di.Documents[p.Document] == doc {
+				res = append(res, p)
+			}
+		}
+	}
+	return res
+}