1414package dag
1515
1616import (
17+ "errors"
1718 "fmt"
1819 "sort"
1920 "strings"
21+
22+ "golang.org/x/exp/maps"
2023)
2124
2225// Vertex represents a node/vertex in a directed acyclic graph.
2326type Vertex struct {
2427 // ID is a unique identifier for the node
2528 ID string
26- // Edges stores the IDs of the nodes that this node has an outgoing edge to.
27- // In kro, this would be the children of a resource.
28- Edges map [string ]struct {}
29+ // Order records the original order, and is used to preserve the original user-provided ordering as far as posible.
30+ Order int
31+ // DependsOn stores the IDs of the nodes that this node depends on.
32+ // If we depend on another vertex, we must appear after that vertex in the topological sort.
33+ DependsOn map [string ]struct {}
34+ }
35+
36+ func (v Vertex ) String () string {
37+ dependsOn := strings .Join (maps .Keys (v .DependsOn ), "," )
38+ return fmt .Sprintf ("Vertex[ID: %s, Order: %d, DependsOn: %s]" , v .ID , v .Order , dependsOn )
2939}
3040
3141// DirectedAcyclicGraph represents a directed acyclic graph
@@ -42,134 +52,128 @@ func NewDirectedAcyclicGraph() *DirectedAcyclicGraph {
4252}
4353
4454// AddVertex adds a new node to the graph.
45- func (d * DirectedAcyclicGraph ) AddVertex (id string ) error {
55+ func (d * DirectedAcyclicGraph ) AddVertex (id string , order int ) error {
4656 if _ , exists := d .Vertices [id ]; exists {
4757 return fmt .Errorf ("node %s already exists" , id )
4858 }
4959 d .Vertices [id ] = & Vertex {
50- ID : id ,
51- Edges : make (map [string ]struct {}),
60+ ID : id ,
61+ Order : order ,
62+ DependsOn : make (map [string ]struct {}),
5263 }
5364 return nil
5465}
5566
5667type CycleError struct {
57- From , to string
58- Cycle []string
68+ Cycle []string
5969}
6070
6171func (e * CycleError ) Error () string {
62- return fmt .Sprintf ("Cannot add edge from %s to %s. This would create a cycle: %s" , e . From , e . to , formatCycle (e .Cycle ))
72+ return fmt .Sprintf ("graph contains a cycle: %s" , formatCycle (e .Cycle ))
6373}
6474
6575func formatCycle (cycle []string ) string {
6676 return strings .Join (cycle , " -> " )
6777}
6878
69- // AddEdge adds a directed edge from one node to another.
70- func (d * DirectedAcyclicGraph ) AddEdge (from , to string ) error {
79+ // AsCycleError returns the (potentially wrapped) CycleError, or nil if it is not a CycleError.
80+ func AsCycleError (err error ) * CycleError {
81+ cycleError := & CycleError {}
82+ if errors .As (err , & cycleError ) {
83+ return cycleError
84+ }
85+ return nil
86+ }
87+
88+ // AddDependencies adds a set of dependencies to the "from" vertex.
89+ // This indicates that all the vertexes in "dependencies" must occur before "from".
90+ func (d * DirectedAcyclicGraph ) AddDependencies (from string , dependencies []string ) error {
7191 fromNode , fromExists := d .Vertices [from ]
72- _ , toExists := d .Vertices [to ]
7392 if ! fromExists {
7493 return fmt .Errorf ("node %s does not exist" , from )
7594 }
76- if ! toExists {
77- return fmt .Errorf ("node %s does not exist" , to )
78- }
79- if from == to {
80- return fmt .Errorf ("self references are not allowed" )
81- }
8295
83- fromNode .Edges [to ] = struct {}{}
96+ for _ , dependency := range dependencies {
97+ _ , toExists := d .Vertices [dependency ]
98+ if ! toExists {
99+ return fmt .Errorf ("node %s does not exist" , dependency )
100+ }
101+ if from == dependency {
102+ return fmt .Errorf ("self references are not allowed" )
103+ }
104+ fromNode .DependsOn [dependency ] = struct {}{}
105+ }
84106
85107 // Check if the graph is still a DAG
86- hasCycle , cycle := d .HasCycle ()
108+ hasCycle , cycle := d .hasCycle ()
87109 if hasCycle {
88110 // Ehmmm, we have a cycle, let's remove the edge we just added
89- delete (fromNode .Edges , to )
111+ for _ , dependency := range dependencies {
112+ delete (fromNode .DependsOn , dependency )
113+ }
90114 return & CycleError {
91- From : from ,
92- to : to ,
93115 Cycle : cycle ,
94116 }
95117 }
96118
97119 return nil
98120}
99121
122+ // TopologicalSort returns the vertexes of the graph, respecting topological ordering first,
123+ // and preserving order of nodes within each "depth" of the topological ordering.
100124func (d * DirectedAcyclicGraph ) TopologicalSort () ([]string , error ) {
101- if cyclic , _ := d .HasCycle (); cyclic {
102- return nil , fmt .Errorf ("graph has a cycle" )
103- }
104-
105125 visited := make (map [string ]bool )
106126 var order []string
107127
108- // Get a sorted list of all vertices
109- vertices := d .GetVertices ()
128+ // Make a list of vertices, sorted by Order
129+ vertices := make ([]* Vertex , 0 , len (d .Vertices ))
130+ for _ , vertex := range d .Vertices {
131+ vertices = append (vertices , vertex )
132+ }
133+ sort .Slice (vertices , func (i , j int ) bool {
134+ return vertices [i ].Order < vertices [j ].Order
135+ })
110136
111- var dfs func (string )
112- dfs = func (node string ) {
113- visited [node ] = true
137+ for len (visited ) < len (vertices ) {
138+ progress := false
114139
115- // Sort the neighbors to ensure deterministic order
116- neighbors := make ([]string , 0 , len (d .Vertices [node ].Edges ))
117- for neighbor := range d .Vertices [node ].Edges {
118- neighbors = append (neighbors , neighbor )
119- }
120- sort .Strings (neighbors )
140+ for _ , vertex := range vertices {
141+ if visited [vertex .ID ] {
142+ continue
143+ }
121144
122- for _ , neighbor := range neighbors {
123- if ! visited [neighbor ] {
124- dfs (neighbor )
145+ allDependenciesReady := true
146+ for dep := range vertex .DependsOn {
147+ if ! visited [dep ] {
148+ allDependenciesReady = false
149+ break
150+ }
151+ }
152+ if ! allDependenciesReady {
153+ continue
125154 }
155+
156+ order = append (order , vertex .ID )
157+ visited [vertex .ID ] = true
158+ progress = true
126159 }
127- order = append (order , node )
128- }
129160
130- // Visit nodes in a deterministic order
131- for _ , node := range vertices {
132- if ! visited [node ] {
133- dfs (node )
161+ if ! progress {
162+ hasCycle , cycle := d .hasCycle ()
163+ if ! hasCycle {
164+ // Unexpected!
165+ return nil , & CycleError {}
166+ }
167+ return nil , & CycleError {
168+ Cycle : cycle ,
169+ }
134170 }
135171 }
136172
137173 return order , nil
138174}
139175
140- // GetVertices returns the nodes in the graph in sorted alphabetical
141- // order.
142- func (d * DirectedAcyclicGraph ) GetVertices () []string {
143- nodes := make ([]string , 0 , len (d .Vertices ))
144- for node := range d .Vertices {
145- nodes = append (nodes , node )
146- }
147-
148- // Ensure deterministic order. This is important for TopologicalSort
149- // to return a deterministic result.
150- sort .Strings (nodes )
151- return nodes
152- }
153-
154- // GetEdges returns the edges in the graph in sorted order...
155- func (d * DirectedAcyclicGraph ) GetEdges () [][2 ]string {
156- var edges [][2 ]string
157- for from , node := range d .Vertices {
158- for to := range node .Edges {
159- edges = append (edges , [2 ]string {from , to })
160- }
161- }
162- sort .Slice (edges , func (i , j int ) bool {
163- // Sort by from node first
164- if edges [i ][0 ] == edges [j ][0 ] {
165- return edges [i ][1 ] < edges [j ][1 ]
166- }
167- return edges [i ][0 ] < edges [j ][0 ]
168- })
169- return edges
170- }
171-
172- func (d * DirectedAcyclicGraph ) HasCycle () (bool , []string ) {
176+ func (d * DirectedAcyclicGraph ) hasCycle () (bool , []string ) {
173177 visited := make (map [string ]bool )
174178 recStack := make (map [string ]bool )
175179 var cyclePath []string
@@ -180,14 +184,14 @@ func (d *DirectedAcyclicGraph) HasCycle() (bool, []string) {
180184 recStack [node ] = true
181185 cyclePath = append (cyclePath , node )
182186
183- for neighbor := range d .Vertices [node ].Edges {
184- if ! visited [neighbor ] {
185- if dfs (neighbor ) {
187+ for dependency := range d .Vertices [node ].DependsOn {
188+ if ! visited [dependency ] {
189+ if dfs (dependency ) {
186190 return true
187191 }
188- } else if recStack [neighbor ] {
192+ } else if recStack [dependency ] {
189193 // Found a cycle, add the closing node to complete the cycle
190- cyclePath = append (cyclePath , neighbor )
194+ cyclePath = append (cyclePath , dependency )
191195 return true
192196 }
193197 }
0 commit comments