1 files changed, 196 insertions, 0 deletions
diff --git a/src/cmd/covdata/subtractintersect.go b/src/cmd/covdata/subtractintersect.go
new file mode 100644
index 0000000..5d71e3d
--- /dev/null
+++ b/src/cmd/covdata/subtractintersect.go
@@ -0,0 +1,196 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package main
+
+// This file contains functions and apis to support the "subtract" and
+// "intersect" subcommands of "go tool covdata".
+
+import (
+	"flag"
+	"fmt"
+	"internal/coverage"
+	"internal/coverage/decodecounter"
+	"internal/coverage/decodemeta"
+	"internal/coverage/pods"
+	"os"
+	"strings"
+)
+
+// makeSubtractIntersectOp creates a subtract or intersect operation.
+// 'mode' here must be either "subtract" or "intersect".
+func makeSubtractIntersectOp(mode string) covOperation {
+	outdirflag = flag.String("o", "", "Output directory to write")
+	s := &sstate{
+		mode:  mode,
+		mm:    newMetaMerge(),
+		inidx: -1,
+	}
+	return s
+}
+
+// sstate holds state needed to implement subtraction and intersection
+// operations on code coverage data files. This type provides methods
+// to implement the CovDataVisitor interface, and is designed to be
+// used in concert with the CovDataReader utility, which abstracts
+// away most of the grubby details of reading coverage data files.
+type sstate struct {
+	mm    *metaMerge
+	inidx int
+	mode  string
+	// Used only for intersection; keyed by pkg/fn ID, it keeps track of
+	// just the set of functions for which we have data in the current
+	// input directory.
+	imm map[pkfunc]struct{}
+}
+
+func (s *sstate) Usage(msg string) {
+	if len(msg) > 0 {
+		fmt.Fprintf(os.Stderr, "error: %s\n", msg)
+	}
+	fmt.Fprintf(os.Stderr, "usage: go tool covdata %s -i=dir1,dir2 -o=<dir>\n\n", s.mode)
+	flag.PrintDefaults()
+	fmt.Fprintf(os.Stderr, "\nExamples:\n\n")
+	op := "from"
+	if s.mode == intersectMode {
+		op = "with"
+	}
+	fmt.Fprintf(os.Stderr, "  go tool covdata %s -i=dir1,dir2 -o=outdir\n\n", s.mode)
+	fmt.Fprintf(os.Stderr, "  \t%ss dir2 %s dir1, writing result\n", s.mode, op)
+	fmt.Fprintf(os.Stderr, "  \tinto output dir outdir.\n")
+	os.Exit(2)
+}
+
+func (s *sstate) Setup() {
+	if *indirsflag == "" {
+		usage("select input directories with '-i' option")
+	}
+	indirs := strings.Split(*indirsflag, ",")
+	if s.mode == subtractMode && len(indirs) != 2 {
+		usage("supply exactly two input dirs for subtract operation")
+	}
+	if *outdirflag == "" {
+		usage("select output directory with '-o' option")
+	}
+}
+
+func (s *sstate) BeginPod(p pods.Pod) {
+	s.mm.beginPod()
+}
+
+func (s *sstate) EndPod(p pods.Pod) {
+	const pcombine = false
+	s.mm.endPod(pcombine)
+}
+
+func (s *sstate) EndCounters() {
+	if s.imm != nil {
+		s.pruneCounters()
+	}
+}
+
+// pruneCounters performs a function-level partial intersection using the
+// current POD counter data (s.mm.pod.pmm) and the intersected data from
+// PODs in previous dirs (s.imm).
+func (s *sstate) pruneCounters() {
+	pkeys := make([]pkfunc, 0, len(s.mm.pod.pmm))
+	for k := range s.mm.pod.pmm {
+		pkeys = append(pkeys, k)
+	}
+	// Remove anything from pmm not found in imm. We don't need to
+	// go the other way (removing things from imm not found in pmm)
+	// since we don't add anything to imm if there is no pmm entry.
+	for _, k := range pkeys {
+		if _, found := s.imm[k]; !found {
+			delete(s.mm.pod.pmm, k)
+		}
+	}
+	s.imm = nil
+}
+
+func (s *sstate) BeginCounterDataFile(cdf string, cdr *decodecounter.CounterDataReader, dirIdx int) {
+	dbgtrace(2, "visiting counter data file %s diridx %d", cdf, dirIdx)
+	if s.inidx != dirIdx {
+		if s.inidx > dirIdx {
+			// We're relying on having data files presented in
+			// the order they appear in the inputs (e.g. first all
+			// data files from input dir 0, then dir 1, etc).
+			panic("decreasing dir index, internal error")
+		}
+		if dirIdx == 0 {
+			// No need to keep track of the functions in the first
+			// directory, since that info will be replicated in
+			// s.mm.pod.pmm.
+			s.imm = nil
+		} else {
+			// We're now starting to visit the Nth directory, N != 0.
+			if s.mode == intersectMode {
+				if s.imm != nil {
+					s.pruneCounters()
+				}
+				s.imm = make(map[pkfunc]struct{})
+			}
+		}
+		s.inidx = dirIdx
+	}
+}
+
+func (s *sstate) EndCounterDataFile(cdf string, cdr *decodecounter.CounterDataReader, dirIdx int) {
+}
+
+func (s *sstate) VisitFuncCounterData(data decodecounter.FuncPayload) {
+	key := pkfunc{pk: data.PkgIdx, fcn: data.FuncIdx}
+
+	if *verbflag >= 5 {
+		fmt.Printf("ctr visit fid=%d pk=%d inidx=%d data.Counters=%+v\n", data.FuncIdx, data.PkgIdx, s.inidx, data.Counters)
+	}
+
+	// If we're processing counter data from the initial (first) input
+	// directory, then just install it into the counter data map
+	// as usual.
+	if s.inidx == 0 {
+		s.mm.visitFuncCounterData(data)
+		return
+	}
+
+	// If we're looking at counter data from a dir other than
+	// the first, then perform the intersect/subtract.
+	if val, ok := s.mm.pod.pmm[key]; ok {
+		if s.mode == subtractMode {
+			for i := 0; i < len(data.Counters); i++ {
+				if data.Counters[i] != 0 {
+					val.Counters[i] = 0
+				}
+			}
+		} else if s.mode == intersectMode {
+			s.imm[key] = struct{}{}
+			for i := 0; i < len(data.Counters); i++ {
+				if data.Counters[i] == 0 {
+					val.Counters[i] = 0
+				}
+			}
+		}
+	}
+}
+
+func (s *sstate) VisitMetaDataFile(mdf string, mfr *decodemeta.CoverageMetaFileReader) {
+	if s.mode == intersectMode {
+		s.imm = make(map[pkfunc]struct{})
+	}
+	s.mm.visitMetaDataFile(mdf, mfr)
+}
+
+func (s *sstate) BeginPackage(pd *decodemeta.CoverageMetaDataDecoder, pkgIdx uint32) {
+	s.mm.visitPackage(pd, pkgIdx, false)
+}
+
+func (s *sstate) EndPackage(pd *decodemeta.CoverageMetaDataDecoder, pkgIdx uint32) {
+}
+
+func (s *sstate) VisitFunc(pkgIdx uint32, fnIdx uint32, fd *coverage.FuncDesc) {
+	s.mm.visitFunc(pkgIdx, fnIdx, fd, s.mode, false)
+}
+
+func (s *sstate) Finish() {
+}