1 files changed, 223 insertions, 0 deletions
diff --git a/src/cmd/go/internal/mvs/graph.go b/src/cmd/go/internal/mvs/graph.go
new file mode 100644
index 0000000..c5de486
--- /dev/null
+++ b/src/cmd/go/internal/mvs/graph.go
@@ -0,0 +1,223 @@
+// Copyright 2020 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 mvs
+
+import (
+	"fmt"
+
+	"golang.org/x/mod/module"
+)
+
+// Graph implements an incremental version of the MVS algorithm, with the
+// requirements pushed by the caller instead of pulled by the MVS traversal.
+type Graph struct {
+	cmp   func(v1, v2 string) int
+	roots []module.Version
+
+	required map[module.Version][]module.Version
+
+	isRoot   map[module.Version]bool // contains true for roots and false for reachable non-roots
+	selected map[string]string       // path → version
+}
+
+// NewGraph returns an incremental MVS graph containing only a set of root
+// dependencies and using the given max function for version strings.
+//
+// The caller must ensure that the root slice is not modified while the Graph
+// may be in use.
+func NewGraph(cmp func(v1, v2 string) int, roots []module.Version) *Graph {
+	g := &Graph{
+		cmp:      cmp,
+		roots:    roots[:len(roots):len(roots)],
+		required: make(map[module.Version][]module.Version),
+		isRoot:   make(map[module.Version]bool),
+		selected: make(map[string]string),
+	}
+
+	for _, m := range roots {
+		g.isRoot[m] = true
+		if g.cmp(g.Selected(m.Path), m.Version) < 0 {
+			g.selected[m.Path] = m.Version
+		}
+	}
+
+	return g
+}
+
+// Require adds the information that module m requires all modules in reqs.
+// The reqs slice must not be modified after it is passed to Require.
+//
+// m must be reachable by some existing chain of requirements from g's target,
+// and Require must not have been called for it already.
+//
+// If any of the modules in reqs has the same path as g's target,
+// the target must have higher precedence than the version in req.
+func (g *Graph) Require(m module.Version, reqs []module.Version) {
+	// To help catch disconnected-graph bugs, enforce that all required versions
+	// are actually reachable from the roots (and therefore should affect the
+	// selected versions of the modules they name).
+	if _, reachable := g.isRoot[m]; !reachable {
+		panic(fmt.Sprintf("%v is not reachable from any root", m))
+	}
+
+	// Truncate reqs to its capacity to avoid aliasing bugs if it is later
+	// returned from RequiredBy and appended to.
+	reqs = reqs[:len(reqs):len(reqs)]
+
+	if _, dup := g.required[m]; dup {
+		panic(fmt.Sprintf("requirements of %v have already been set", m))
+	}
+	g.required[m] = reqs
+
+	for _, dep := range reqs {
+		// Mark dep reachable, regardless of whether it is selected.
+		if _, ok := g.isRoot[dep]; !ok {
+			g.isRoot[dep] = false
+		}
+
+		if g.cmp(g.Selected(dep.Path), dep.Version) < 0 {
+			g.selected[dep.Path] = dep.Version
+		}
+	}
+}
+
+// RequiredBy returns the slice of requirements passed to Require for m, if any,
+// with its capacity reduced to its length.
+// If Require has not been called for m, RequiredBy(m) returns ok=false.
+//
+// The caller must not modify the returned slice, but may safely append to it
+// and may rely on it not to be modified.
+func (g *Graph) RequiredBy(m module.Version) (reqs []module.Version, ok bool) {
+	reqs, ok = g.required[m]
+	return reqs, ok
+}
+
+// Selected returns the selected version of the given module path.
+//
+// If no version is selected, Selected returns version "none".
+func (g *Graph) Selected(path string) (version string) {
+	v, ok := g.selected[path]
+	if !ok {
+		return "none"
+	}
+	return v
+}
+
+// BuildList returns the selected versions of all modules present in the Graph,
+// beginning with the selected versions of each module path in the roots of g.
+//
+// The order of the remaining elements in the list is deterministic
+// but arbitrary.
+func (g *Graph) BuildList() []module.Version {
+	seenRoot := make(map[string]bool, len(g.roots))
+
+	var list []module.Version
+	for _, r := range g.roots {
+		if seenRoot[r.Path] {
+			// Multiple copies of the same root, with the same or different versions,
+			// are a bit of a degenerate case: we will take the transitive
+			// requirements of both roots into account, but only the higher one can
+			// possibly be selected. However — especially given that we need the
+			// seenRoot map for later anyway — it is simpler to support this
+			// degenerate case than to forbid it.
+			continue
+		}
+
+		if v := g.Selected(r.Path); v != "none" {
+			list = append(list, module.Version{Path: r.Path, Version: v})
+		}
+		seenRoot[r.Path] = true
+	}
+	uniqueRoots := list
+
+	for path, version := range g.selected {
+		if !seenRoot[path] {
+			list = append(list, module.Version{Path: path, Version: version})
+		}
+	}
+	module.Sort(list[len(uniqueRoots):])
+
+	return list
+}
+
+// WalkBreadthFirst invokes f once, in breadth-first order, for each module
+// version other than "none" that appears in the graph, regardless of whether
+// that version is selected.
+func (g *Graph) WalkBreadthFirst(f func(m module.Version)) {
+	var queue []module.Version
+	enqueued := make(map[module.Version]bool)
+	for _, m := range g.roots {
+		if m.Version != "none" {
+			queue = append(queue, m)
+			enqueued[m] = true
+		}
+	}
+
+	for len(queue) > 0 {
+		m := queue[0]
+		queue = queue[1:]
+
+		f(m)
+
+		reqs, _ := g.RequiredBy(m)
+		for _, r := range reqs {
+			if !enqueued[r] && r.Version != "none" {
+				queue = append(queue, r)
+				enqueued[r] = true
+			}
+		}
+	}
+}
+
+// FindPath reports a shortest requirement path starting at one of the roots of
+// the graph and ending at a module version m for which f(m) returns true, or
+// nil if no such path exists.
+func (g *Graph) FindPath(f func(module.Version) bool) []module.Version {
+	// firstRequires[a] = b means that in a breadth-first traversal of the
+	// requirement graph, the module version a was first required by b.
+	firstRequires := make(map[module.Version]module.Version)
+
+	queue := g.roots
+	for _, m := range g.roots {
+		firstRequires[m] = module.Version{}
+	}
+
+	for len(queue) > 0 {
+		m := queue[0]
+		queue = queue[1:]
+
+		if f(m) {
+			// Construct the path reversed (because we're starting from the far
+			// endpoint), then reverse it.
+			path := []module.Version{m}
+			for {
+				m = firstRequires[m]
+				if m.Path == "" {
+					break
+				}
+				path = append(path, m)
+			}
+
+			i, j := 0, len(path)-1
+			for i < j {
+				path[i], path[j] = path[j], path[i]
+				i++
+				j--
+			}
+
+			return path
+		}
+
+		reqs, _ := g.RequiredBy(m)
+		for _, r := range reqs {
+			if _, seen := firstRequires[r]; !seen {
+				queue = append(queue, r)
+				firstRequires[r] = m
+			}
+		}
+	}
+
+	return nil
+}