-
Notifications
You must be signed in to change notification settings - Fork 482
/
walk.go
134 lines (109 loc) · 3.54 KB
/
walk.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package dag
// WalkFunc is a function that gets invoked when walking a Graph. Walking will
// stop if WalkFunc returns a non-nil error.
type WalkFunc func(n Node) error
// Walk performs a depth-first walk of outgoing edges for all nodes in start,
// invoking the provided fn for each node. Walk returns the error returned by
// fn.
//
// Nodes unreachable from start will not be passed to fn.
func Walk(g *Graph, start []Node, fn WalkFunc) error {
var (
visited = make(nodeSet)
unchecked = make([]Node, 0, len(start))
)
// Prefill the set of unchecked nodes with our start set.
unchecked = append(unchecked, start...)
// Iterate through unchecked nodes, visiting each in turn and adding outgoing
// edges to the unchecked list until all reachable nodes have been processed.
for len(unchecked) > 0 {
check := unchecked[len(unchecked)-1]
unchecked = unchecked[:len(unchecked)-1]
if visited.Has(check) {
continue
}
visited.Add(check)
if err := fn(check); err != nil {
return err
}
for n := range g.outEdges[check] {
unchecked = append(unchecked, n)
}
}
return nil
}
// WalkReverse performs a depth-first walk of incoming edges for all nodes in
// start, invoking the provided fn for each node. Walk returns the error
// returned by fn.
//
// Nodes unreachable from start will not be passed to fn.
func WalkReverse(g *Graph, start []Node, fn WalkFunc) error {
var (
visited = make(nodeSet)
unchecked = make([]Node, 0, len(start))
)
// Prefill the set of unchecked nodes with our start set.
unchecked = append(unchecked, start...)
// Iterate through unchecked nodes, visiting each in turn and adding incoming
// edges to the unchecked list until all reachable nodes have been processed.
for len(unchecked) > 0 {
check := unchecked[len(unchecked)-1]
unchecked = unchecked[:len(unchecked)-1]
if visited.Has(check) {
continue
}
visited.Add(check)
if err := fn(check); err != nil {
return err
}
for n := range g.inEdges[check] {
unchecked = append(unchecked, n)
}
}
return nil
}
// WalkTopological performs a topological walk of all nodes in start in
// dependency order: a node will not be visited until its outgoing edges are
// visited first.
//
// Nodes will not be passed to fn if they are not reachable from start or if
// not all of their outgoing edges are reachable from start.
func WalkTopological(g *Graph, start []Node, fn WalkFunc) error {
// NOTE(rfratto): WalkTopological is an implementation of Kahn's algorithm
// which leaves g unmodified.
var (
visited = make(nodeSet)
unchecked = make([]Node, 0, len(start))
remainingDeps = make(map[Node]int)
)
// Pre-fill the set of nodes to check from the start list.
unchecked = append(unchecked, start...)
for len(unchecked) > 0 {
check := unchecked[len(unchecked)-1]
unchecked = unchecked[:len(unchecked)-1]
if visited.Has(check) {
continue
}
visited.Add(check)
if err := fn(check); err != nil {
return err
}
// Iterate through the incoming edges to check and queue nodes if we're the
// last edge to be walked.
for n := range g.inEdges[check] {
// remainingDeps starts with the number of edges, and we subtract one for
// each outgoing edge that's visited.
if _, ok := remainingDeps[n]; !ok {
remainingDeps[n] = len(g.outEdges[n])
}
remainingDeps[n]--
// Only enqueue the incoming edge once all of its outgoing edges have
// been consumed. This prevents it from being visited before its
// dependencies.
if remainingDeps[n] == 0 {
unchecked = append(unchecked, n)
}
}
}
return nil
}