Merge pull request #143013 from xigang/fix-dag-visit-early-stop

Fix early termination in gonum graph traversal

third_party/forked/gonum/graph/simple/directed_acyclic.go

@@ -31,10 +31,11 @@ func (g *DirectedAcyclicGraph) From(n graph.Node) []graph.Node {
 
 	fid := n.ID()
 	nodes := make([]graph.Node, 0, g.UndirectedGraph.edges[n.ID()].Len())
-	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		if edge.From().ID() == fid {
 			nodes = append(nodes, g.UndirectedGraph.nodes[edge.To().ID()])
 		}
+		return true
 	})
 	return nodes
 }
@@ -44,12 +45,11 @@ func (g *DirectedAcyclicGraph) VisitFrom(n graph.Node, visitor func(neighbor gra
 		return
 	}
 	fid := n.ID()
-	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		if edge.From().ID() == fid {
-			if !visitor(g.UndirectedGraph.nodes[edge.To().ID()]) {
-				return
-			}
+			return visitor(g.UndirectedGraph.nodes[edge.To().ID()])
 		}
+		return true
 	})
 }
 
@@ -60,10 +60,11 @@ func (g *DirectedAcyclicGraph) To(n graph.Node) []graph.Node {
 
 	tid := n.ID()
 	nodes := make([]graph.Node, 0, g.UndirectedGraph.edges[n.ID()].Len())
-	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		if edge.To().ID() == tid {
 			nodes = append(nodes, g.UndirectedGraph.nodes[edge.From().ID()])
 		}
+		return true
 	})
 	return nodes
 }
@@ -73,11 +74,10 @@ func (g *DirectedAcyclicGraph) VisitTo(n graph.Node, visitor func(neighbor graph
 		return
 	}
 	tid := n.ID()
-	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.UndirectedGraph.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		if edge.To().ID() == tid {
-			if !visitor(g.UndirectedGraph.nodes[edge.From().ID()]) {
-				return
-			}
+			return visitor(g.UndirectedGraph.nodes[edge.From().ID()])
 		}
+		return true
 	})
 }

third_party/forked/gonum/graph/simple/directed_acyclic_test.go

@@ -5,6 +5,7 @@
 package simple
 
 import (
+	"fmt"
 	"math"
 	"testing"
 
@@ -60,3 +61,72 @@ func TestAcyclicIssue123DirectedGraph(t *testing.T) {
 	n2 := Node(g.NewNodeID())
 	g.AddNode(n2)
 }
+
+func TestAcyclicVisitStopsEarly(t *testing.T) {
+	// 6 edges on the hub promotes its edgeHolder to a map.
+	g := NewDirectedAcyclicGraph(0, math.Inf(1))
+	for i := 1; i <= 3; i++ {
+		g.SetEdge(Edge{F: Node(0), T: Node(i), W: 1})
+	}
+	for i := 4; i <= 6; i++ {
+		g.SetEdge(Edge{F: Node(i), T: Node(0), W: 1})
+	}
+
+	fromCalls := 0
+	g.VisitFrom(Node(0), func(n graph.Node) bool {
+		fromCalls++
+		if n.ID() < 1 || n.ID() > 3 {
+			t.Errorf("VisitFrom visited unexpected node %d", n.ID())
+		}
+		return false
+	})
+	if fromCalls != 1 {
+		t.Errorf("expected VisitFrom visitor to be called once, got %d", fromCalls)
+	}
+
+	toCalls := 0
+	g.VisitTo(Node(0), func(n graph.Node) bool {
+		toCalls++
+		if n.ID() < 4 || n.ID() > 6 {
+			t.Errorf("VisitTo visited unexpected node %d", n.ID())
+		}
+		return false
+	})
+	if toCalls != 1 {
+		t.Errorf("expected VisitTo visitor to be called once, got %d", toCalls)
+	}
+
+	// Returning true must still visit every neighbor in the right direction.
+	fromAll, toAll := 0, 0
+	g.VisitFrom(Node(0), func(graph.Node) bool { fromAll++; return true })
+	g.VisitTo(Node(0), func(graph.Node) bool { toAll++; return true })
+	if fromAll != 3 || toAll != 3 {
+		t.Errorf("expected 3 from and 3 to neighbors, got %d and %d", fromAll, toAll)
+	}
+}
+
+// BenchmarkAcyclicVisitStopsEarly measures a visitor that stops after the
+// first neighbor, which is what VisitingDepthFirst.Walk does once it finds
+// its target. Cost should not scale with the hub's degree.
+func BenchmarkAcyclicVisitStopsEarly(b *testing.B) {
+	for _, degree := range []int{2, 100, 10000} {
+		// Hub 0 has degree outgoing edges and degree incoming edges.
+		g := NewDirectedAcyclicGraph(0, math.Inf(1))
+		for i := 1; i <= degree; i++ {
+			g.SetEdge(Edge{F: Node(0), T: Node(i), W: 1})
+			g.SetEdge(Edge{F: Node(degree + i), T: Node(0), W: 1})
+		}
+		stop := func(graph.Node) bool { return false }
+
+		b.Run(fmt.Sprintf("VisitFrom/degree=%d", degree), func(b *testing.B) {
+			for b.Loop() {
+				g.VisitFrom(Node(0), stop)
+			}
+		})
+		b.Run(fmt.Sprintf("VisitTo/degree=%d", degree), func(b *testing.B) {
+			for b.Loop() {
+				g.VisitTo(Node(0), stop)
+			}
+		})
+	}
+}

third_party/forked/gonum/graph/simple/edgeholder.go

@@ -4,8 +4,9 @@ import "k8s.io/kubernetes/third_party/forked/gonum/graph"
 
 // edgeHolder represents a set of edges, with no more than one edge to or from a particular neighbor node
 type edgeHolder interface {
-	// Visit invokes visitor with each edge and the id of the neighbor node in the edge
-	Visit(visitor func(neighbor int, edge graph.Edge))
+	// Visit invokes visitor with each edge and the id of the neighbor node in the edge,
+	// stopping as soon as visitor returns false
+	Visit(visitor func(neighbor int, edge graph.Edge) (shouldContinue bool))
 	// Delete removes edges to or from the specified neighbor
 	Delete(neighbor int) edgeHolder
 	// Set stores the edge to or from the specified neighbor
@@ -22,12 +23,14 @@ type sliceEdgeHolder struct {
 	edges []graph.Edge
 }
 
-func (e *sliceEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge)) {
+func (e *sliceEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge) (shouldContinue bool)) {
 	for _, edge := range e.edges {
-		if edge.From().ID() == e.self {
-			visitor(edge.To().ID(), edge)
-		} else {
-			visitor(edge.From().ID(), edge)
+		id := edge.From().ID()
+		if id == e.self {
+			id = edge.To().ID()
+		}
+		if !visitor(id, edge) {
+			return
 		}
 	}
 }
@@ -100,9 +103,11 @@ func (e *sliceEdgeHolder) Len() int {
 // mapEdgeHolder holds a map of neighbors to edges
 type mapEdgeHolder map[int]graph.Edge
 
-func (e mapEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge)) {
+func (e mapEdgeHolder) Visit(visitor func(neighbor int, edge graph.Edge) (shouldContinue bool)) {
 	for neighbor, edge := range e {
-		visitor(neighbor, edge)
+		if !visitor(neighbor, edge) {
+			return
+		}
 	}
 }
 func (e mapEdgeHolder) Delete(neighbor int) edgeHolder {

third_party/forked/gonum/graph/simple/edgeholder_test.go

@@ -18,7 +18,7 @@ func TestEdgeHolder(t *testing.T) {
 	if n, ok := holder.Get(2); ok || n != nil {
 		t.Errorf("expected nil,false")
 	}
-	holder.Visit(func(_ int, _ graph.Edge) { t.Errorf("unexpected call to visitor") })
+	holder.Visit(func(_ int, _ graph.Edge) (shouldContinue bool) { t.Errorf("unexpected call to visitor"); return true })
 	holder = holder.Delete(2)
 
 	// Insert an edge to ourselves
@@ -30,7 +30,10 @@ func TestEdgeHolder(t *testing.T) {
 		t.Errorf("expected edge to ourselves, got %#v", n)
 	}
 	neighbors := []int{}
-	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	holder.Visit(func(neighbor int, _ graph.Edge) (shouldContinue bool) {
+		neighbors = append(neighbors, neighbor)
+		return true
+	})
 	if !reflect.DeepEqual(neighbors, []int{1}) {
 		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
 	}
@@ -46,7 +49,10 @@ func TestEdgeHolder(t *testing.T) {
 		t.Errorf("expected edge from us to another node, got %#v", n)
 	}
 	neighbors = []int{}
-	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	holder.Visit(func(neighbor int, _ graph.Edge) (shouldContinue bool) {
+		neighbors = append(neighbors, neighbor)
+		return true
+	})
 	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4}) {
 		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
 	}
@@ -62,7 +68,10 @@ func TestEdgeHolder(t *testing.T) {
 		t.Errorf("expected reversed edge, got %#v", n)
 	}
 	neighbors = []int{}
-	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	holder.Visit(func(neighbor int, _ graph.Edge) (shouldContinue bool) {
+		neighbors = append(neighbors, neighbor)
+		return true
+	})
 	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4}) {
 		t.Errorf("expected a single visit to ourselves, got %v", neighbors)
 	}
@@ -87,7 +96,10 @@ func TestEdgeHolder(t *testing.T) {
 		t.Errorf("expected new edge, got %#v", n)
 	}
 	neighbors = []int{}
-	holder.Visit(func(neighbor int, _ graph.Edge) { neighbors = append(neighbors, neighbor) })
+	holder.Visit(func(neighbor int, _ graph.Edge) (shouldContinue bool) {
+		neighbors = append(neighbors, neighbor)
+		return true
+	})
 	sort.Ints(neighbors) // sort, map order is random
 	if !reflect.DeepEqual(neighbors, []int{1, 2, 3, 4, 5}) {
 		t.Errorf("expected 1,2,3,4,5, got %v", neighbors)
@@ -102,3 +114,22 @@ func TestEdgeHolder(t *testing.T) {
 		t.Errorf("expected 0")
 	}
 }
+
+func TestEdgeHolderVisitStopsEarly(t *testing.T) {
+	// 3 edges stays a sliceEdgeHolder; 6 edges promotes to a mapEdgeHolder.
+	for _, numEdges := range []int{3, 6} {
+		holder := edgeHolder(&sliceEdgeHolder{self: 1})
+		for i := 2; i < 2+numEdges; i++ {
+			holder = holder.Set(i, Edge{F: Node(1), T: Node(i)})
+		}
+
+		calls := 0
+		holder.Visit(func(_ int, _ graph.Edge) (shouldContinue bool) {
+			calls++
+			return false
+		})
+		if calls != 1 {
+			t.Errorf("%T with %d edges: expected visitor to be called once, got %d", holder, numEdges, calls)
+		}
+	}
+}

third_party/forked/gonum/graph/simple/undirected.go

@@ -80,8 +80,9 @@ func (g *UndirectedGraph) RemoveNode(n graph.Node) {
 	}
 	delete(g.nodes, n.ID())
 
-	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		g.edges[neighbor] = g.edges[neighbor].Delete(n.ID())
+		return true
 	})
 	delete(g.edges, n.ID())
 
@@ -159,15 +160,16 @@ func (g *UndirectedGraph) Edges() []graph.Edge {
 
 	seen := make(map[[2]int]struct{})
 	for _, u := range g.edges {
-		u.Visit(func(neighbor int, e graph.Edge) {
+		u.Visit(func(neighbor int, e graph.Edge) (shouldContinue bool) {
 			uid := e.From().ID()
 			vid := e.To().ID()
 			if _, ok := seen[[2]int{uid, vid}]; ok {
-				return
+				return true
 			}
 			seen[[2]int{uid, vid}] = struct{}{}
 			seen[[2]int{vid, uid}] = struct{}{}
 			edges = append(edges, e)
+			return true
 		})
 	}
 
@@ -182,9 +184,10 @@ func (g *UndirectedGraph) From(n graph.Node) []graph.Node {
 
 	nodes := make([]graph.Node, g.edges[n.ID()].Len())
 	i := 0
-	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) {
+	g.edges[n.ID()].Visit(func(neighbor int, edge graph.Edge) (shouldContinue bool) {
 		nodes[i] = g.nodes[neighbor]
 		i++
+		return true
 	})
 
 	return nodes