graph: implement VisitBreadth/VisitDepth

graph/graph.go

@@ -141,6 +141,15 @@ func (g Graph) addDirty(edgeID string, e Edge) {
 	g.vOuts.add(e.Tail.ID, edgeID)
 }
 
+func (g Graph) estSize() int {
+	lvIns := len(g.vIns)
+	lvOuts := len(g.vOuts)
+	if lvIns > lvOuts {
+		return lvIns
+	}
+	return lvOuts
+}
+
 // Del returns a new Graph instance without the given Edge in it. If the
 // original Graph didn't have that Edge this returns the original Graph.
 func (g Graph) Del(e Edge) Graph {
@@ -242,7 +251,7 @@ type Node struct {
 	Value
 
 	// All Edges in the Graph with this Node's Value as their Head and Tail,
-	// respectively.
+	// respectively. These should not be expected to be deterministic.
 	Ins, Outs []Edge
 }
 
@@ -318,7 +327,86 @@ func (g Graph) Traverse(start Value, next func(n Node) (Value, bool)) {
 	}
 }
 
-// TODO VisitBreadth/VisitDepth
+// VisitBreadth is like Traverse, except that each Node is only visited once,
+// and the order of visited Nodes is determined by traversing each Node's output
+// Edges breadth-wise.
+//
+// If the boolean returned from the callback function is false, or the start
+// Value has no edges in the Graph, traversal stops and this method returns.
+//
+// The exact order of Nodes visited is _not_ deterministic.
+func (g Graph) VisitBreadth(start Value, callback func(n Node) bool) {
+	visited := map[string]bool{}
+	toVisit := make([]Value, 0, g.estSize())
+	toVisit = append(toVisit, start)
+
+	for {
+		if len(toVisit) == 0 {
+			return
+		}
+
+		// shift val off front
+		val := toVisit[0]
+		toVisit = toVisit[1:]
+		if visited[val.ID] {
+			continue
+		}
+		node, ok := g.Node(val)
+		if !ok {
+			continue
+		} else if !callback(node) {
+			return
+		}
+		visited[val.ID] = true
+		for _, edge := range node.Outs {
+			if visited[edge.Head.ID] {
+				continue
+			}
+			toVisit = append(toVisit, edge.Head)
+		}
+	}
+}
+
+// VisitDepth is like Traverse, except that each Node is only visited once,
+// and the order of visited Nodes is determined by traversing each Node's output
+// Edges depth-wise.
+//
+// If the boolean returned from the callback function is false, or the start
+// Value has no edges in the Graph, traversal stops and this method returns.
+//
+// The exact order of Nodes visited is _not_ deterministic.
+func (g Graph) VisitDepth(start Value, callback func(n Node) bool) {
+	// VisitDepth is actually the same as VisitBreadth, only you read off the
+	// toVisit list from back-to-front
+	visited := map[string]bool{}
+	toVisit := make([]Value, 0, g.estSize())
+	toVisit = append(toVisit, start)
+
+	for {
+		if len(toVisit) == 0 {
+			return
+		}
+
+		val := toVisit[0]
+		toVisit = toVisit[:len(toVisit)-1] // pop val off back
+		if visited[val.ID] {
+			continue
+		}
+		node, ok := g.Node(val)
+		if !ok {
+			continue
+		} else if !callback(node) {
+			return
+		}
+		visited[val.ID] = true
+		for _, edge := range node.Outs {
+			if visited[edge.Head.ID] {
+				continue
+			}
+			toVisit = append(toVisit, edge.Head)
+		}
+	}
+}
 
 func (g Graph) edgesShared(g2 Graph) bool {
 	for id := range g2.m {

graph/graph_test.go

@@ -258,3 +258,105 @@ func TestDisjoinUnion(t *T) {
 		t.Fatal(err)
 	}
 }
+
+func TestVisitBreadth(t *T) {
+	t.Parallel()
+	type state struct {
+		g Graph
+		// each rank describes the set of values (by ID) which should be
+		// visited in that rank. Within a rank the values will be visited in any
+		// order
+		ranks []map[string]bool
+	}
+
+	thisRank := func(s state) map[string]bool {
+		return s.ranks[len(s.ranks)-1]
+	}
+
+	prevRank := func(s state) map[string]bool {
+		return s.ranks[len(s.ranks)-2]
+	}
+
+	randFromRank := func(s state, rankPickFn func(state) map[string]bool) Value {
+		rank := rankPickFn(s)
+		rankL := make([]string, 0, len(rank))
+		for id := range rank {
+			rankL = append(rankL, id)
+		}
+		return strV(mrand.Element(rankL, nil).(string))
+	}
+
+	type params struct {
+		newRank bool
+		e       Edge
+	}
+
+	chk := mchk.Checker{
+		Init: func() mchk.State {
+			return state{
+				ranks: []map[string]bool{
+					map[string]bool{"start": true},
+					map[string]bool{},
+				},
+			}
+		},
+		Next: func(ss mchk.State) mchk.Action {
+			s := ss.(state)
+			var p params
+			p.newRank = len(thisRank(s)) > 0 && mrand.Intn(10) == 0
+			if p.newRank {
+				p.e.Head = strV(mrand.Hex(3))
+				p.e.Tail = randFromRank(s, thisRank)
+			} else {
+				p.e.Head = strV(mrand.Hex(2))
+				p.e.Tail = randFromRank(s, prevRank)
+			}
+			return mchk.Action{Params: p}
+		},
+		Apply: func(ss mchk.State, a mchk.Action) (mchk.State, error) {
+			s, p := ss.(state), a.Params.(params)
+			if p.newRank {
+				s.ranks = append(s.ranks, map[string]bool{})
+			}
+			if !s.g.Has(p.e.Head) {
+				thisRank(s)[p.e.Head.ID] = true
+			}
+			s.g = s.g.Add(p.e)
+
+			// check the visit
+			var err error
+			expRanks := s.ranks
+			currRank := map[string]bool{}
+			s.g.VisitBreadth(strV("start"), func(n Node) bool {
+				currRank[n.Value.ID] = true
+				if len(currRank) != len(expRanks[0]) {
+					return true
+				}
+				if err = massert.Equal(expRanks[0], currRank).Assert(); err != nil {
+					return false
+				}
+				expRanks = expRanks[1:]
+				currRank = map[string]bool{}
+				return true
+			})
+			if err != nil {
+				return nil, err
+			}
+
+			err = massert.All(
+				massert.Len(expRanks, 0),
+				massert.Len(currRank, 0),
+			).Assert()
+			return s, err
+		},
+		DontMinimize: true,
+	}
+
+	if err := chk.RunCase(); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := chk.RunFor(5 * time.Second); err != nil {
+		t.Fatal(err)
+	}
+}