1package dag 2 3import ( 4 "bytes" 5 "encoding/json" 6 "fmt" 7 "io" 8 "sort" 9) 10 11// Graph is used to represent a dependency graph. 12type Graph struct { 13 vertices *Set 14 edges *Set 15 downEdges map[interface{}]*Set 16 upEdges map[interface{}]*Set 17 18 // JSON encoder for recording debug information 19 debug *encoder 20} 21 22// Subgrapher allows a Vertex to be a Graph itself, by returning a Grapher. 23type Subgrapher interface { 24 Subgraph() Grapher 25} 26 27// A Grapher is any type that returns a Grapher, mainly used to identify 28// dag.Graph and dag.AcyclicGraph. In the case of Graph and AcyclicGraph, they 29// return themselves. 30type Grapher interface { 31 DirectedGraph() Grapher 32} 33 34// Vertex of the graph. 35type Vertex interface{} 36 37// NamedVertex is an optional interface that can be implemented by Vertex 38// to give it a human-friendly name that is used for outputting the graph. 39type NamedVertex interface { 40 Vertex 41 Name() string 42} 43 44func (g *Graph) DirectedGraph() Grapher { 45 return g 46} 47 48// Vertices returns the list of all the vertices in the graph. 49func (g *Graph) Vertices() []Vertex { 50 list := g.vertices.List() 51 result := make([]Vertex, len(list)) 52 for i, v := range list { 53 result[i] = v.(Vertex) 54 } 55 56 return result 57} 58 59// Edges returns the list of all the edges in the graph. 60func (g *Graph) Edges() []Edge { 61 list := g.edges.List() 62 result := make([]Edge, len(list)) 63 for i, v := range list { 64 result[i] = v.(Edge) 65 } 66 67 return result 68} 69 70// EdgesFrom returns the list of edges from the given source. 71func (g *Graph) EdgesFrom(v Vertex) []Edge { 72 var result []Edge 73 from := hashcode(v) 74 for _, e := range g.Edges() { 75 if hashcode(e.Source()) == from { 76 result = append(result, e) 77 } 78 } 79 80 return result 81} 82 83// EdgesTo returns the list of edges to the given target. 84func (g *Graph) EdgesTo(v Vertex) []Edge { 85 var result []Edge 86 search := hashcode(v) 87 for _, e := range g.Edges() { 88 if hashcode(e.Target()) == search { 89 result = append(result, e) 90 } 91 } 92 93 return result 94} 95 96// HasVertex checks if the given Vertex is present in the graph. 97func (g *Graph) HasVertex(v Vertex) bool { 98 return g.vertices.Include(v) 99} 100 101// HasEdge checks if the given Edge is present in the graph. 102func (g *Graph) HasEdge(e Edge) bool { 103 return g.edges.Include(e) 104} 105 106// Add adds a vertex to the graph. This is safe to call multiple time with 107// the same Vertex. 108func (g *Graph) Add(v Vertex) Vertex { 109 g.init() 110 g.vertices.Add(v) 111 g.debug.Add(v) 112 return v 113} 114 115// Remove removes a vertex from the graph. This will also remove any 116// edges with this vertex as a source or target. 117func (g *Graph) Remove(v Vertex) Vertex { 118 // Delete the vertex itself 119 g.vertices.Delete(v) 120 g.debug.Remove(v) 121 122 // Delete the edges to non-existent things 123 for _, target := range g.DownEdges(v).List() { 124 g.RemoveEdge(BasicEdge(v, target)) 125 } 126 for _, source := range g.UpEdges(v).List() { 127 g.RemoveEdge(BasicEdge(source, v)) 128 } 129 130 return nil 131} 132 133// Replace replaces the original Vertex with replacement. If the original 134// does not exist within the graph, then false is returned. Otherwise, true 135// is returned. 136func (g *Graph) Replace(original, replacement Vertex) bool { 137 // If we don't have the original, we can't do anything 138 if !g.vertices.Include(original) { 139 return false 140 } 141 142 defer g.debug.BeginOperation("Replace", "").End("") 143 144 // If they're the same, then don't do anything 145 if original == replacement { 146 return true 147 } 148 149 // Add our new vertex, then copy all the edges 150 g.Add(replacement) 151 for _, target := range g.DownEdges(original).List() { 152 g.Connect(BasicEdge(replacement, target)) 153 } 154 for _, source := range g.UpEdges(original).List() { 155 g.Connect(BasicEdge(source, replacement)) 156 } 157 158 // Remove our old vertex, which will also remove all the edges 159 g.Remove(original) 160 161 return true 162} 163 164// RemoveEdge removes an edge from the graph. 165func (g *Graph) RemoveEdge(edge Edge) { 166 g.init() 167 g.debug.RemoveEdge(edge) 168 169 // Delete the edge from the set 170 g.edges.Delete(edge) 171 172 // Delete the up/down edges 173 if s, ok := g.downEdges[hashcode(edge.Source())]; ok { 174 s.Delete(edge.Target()) 175 } 176 if s, ok := g.upEdges[hashcode(edge.Target())]; ok { 177 s.Delete(edge.Source()) 178 } 179} 180 181// DownEdges returns the outward edges from the source Vertex v. 182func (g *Graph) DownEdges(v Vertex) *Set { 183 g.init() 184 return g.downEdges[hashcode(v)] 185} 186 187// UpEdges returns the inward edges to the destination Vertex v. 188func (g *Graph) UpEdges(v Vertex) *Set { 189 g.init() 190 return g.upEdges[hashcode(v)] 191} 192 193// Connect adds an edge with the given source and target. This is safe to 194// call multiple times with the same value. Note that the same value is 195// verified through pointer equality of the vertices, not through the 196// value of the edge itself. 197func (g *Graph) Connect(edge Edge) { 198 g.init() 199 g.debug.Connect(edge) 200 201 source := edge.Source() 202 target := edge.Target() 203 sourceCode := hashcode(source) 204 targetCode := hashcode(target) 205 206 // Do we have this already? If so, don't add it again. 207 if s, ok := g.downEdges[sourceCode]; ok && s.Include(target) { 208 return 209 } 210 211 // Add the edge to the set 212 g.edges.Add(edge) 213 214 // Add the down edge 215 s, ok := g.downEdges[sourceCode] 216 if !ok { 217 s = new(Set) 218 g.downEdges[sourceCode] = s 219 } 220 s.Add(target) 221 222 // Add the up edge 223 s, ok = g.upEdges[targetCode] 224 if !ok { 225 s = new(Set) 226 g.upEdges[targetCode] = s 227 } 228 s.Add(source) 229} 230 231// String outputs some human-friendly output for the graph structure. 232func (g *Graph) StringWithNodeTypes() string { 233 var buf bytes.Buffer 234 235 // Build the list of node names and a mapping so that we can more 236 // easily alphabetize the output to remain deterministic. 237 vertices := g.Vertices() 238 names := make([]string, 0, len(vertices)) 239 mapping := make(map[string]Vertex, len(vertices)) 240 for _, v := range vertices { 241 name := VertexName(v) 242 names = append(names, name) 243 mapping[name] = v 244 } 245 sort.Strings(names) 246 247 // Write each node in order... 248 for _, name := range names { 249 v := mapping[name] 250 targets := g.downEdges[hashcode(v)] 251 252 buf.WriteString(fmt.Sprintf("%s - %T\n", name, v)) 253 254 // Alphabetize dependencies 255 deps := make([]string, 0, targets.Len()) 256 targetNodes := make(map[string]Vertex) 257 for _, target := range targets.List() { 258 dep := VertexName(target) 259 deps = append(deps, dep) 260 targetNodes[dep] = target 261 } 262 sort.Strings(deps) 263 264 // Write dependencies 265 for _, d := range deps { 266 buf.WriteString(fmt.Sprintf(" %s - %T\n", d, targetNodes[d])) 267 } 268 } 269 270 return buf.String() 271} 272 273// String outputs some human-friendly output for the graph structure. 274func (g *Graph) String() string { 275 var buf bytes.Buffer 276 277 // Build the list of node names and a mapping so that we can more 278 // easily alphabetize the output to remain deterministic. 279 vertices := g.Vertices() 280 names := make([]string, 0, len(vertices)) 281 mapping := make(map[string]Vertex, len(vertices)) 282 for _, v := range vertices { 283 name := VertexName(v) 284 names = append(names, name) 285 mapping[name] = v 286 } 287 sort.Strings(names) 288 289 // Write each node in order... 290 for _, name := range names { 291 v := mapping[name] 292 targets := g.downEdges[hashcode(v)] 293 294 buf.WriteString(fmt.Sprintf("%s\n", name)) 295 296 // Alphabetize dependencies 297 deps := make([]string, 0, targets.Len()) 298 for _, target := range targets.List() { 299 deps = append(deps, VertexName(target)) 300 } 301 sort.Strings(deps) 302 303 // Write dependencies 304 for _, d := range deps { 305 buf.WriteString(fmt.Sprintf(" %s\n", d)) 306 } 307 } 308 309 return buf.String() 310} 311 312func (g *Graph) init() { 313 if g.vertices == nil { 314 g.vertices = new(Set) 315 } 316 if g.edges == nil { 317 g.edges = new(Set) 318 } 319 if g.downEdges == nil { 320 g.downEdges = make(map[interface{}]*Set) 321 } 322 if g.upEdges == nil { 323 g.upEdges = make(map[interface{}]*Set) 324 } 325} 326 327// Dot returns a dot-formatted representation of the Graph. 328func (g *Graph) Dot(opts *DotOpts) []byte { 329 return newMarshalGraph("", g).Dot(opts) 330} 331 332// MarshalJSON returns a JSON representation of the entire Graph. 333func (g *Graph) MarshalJSON() ([]byte, error) { 334 dg := newMarshalGraph("root", g) 335 return json.MarshalIndent(dg, "", " ") 336} 337 338// SetDebugWriter sets the io.Writer where the Graph will record debug 339// information. After this is set, the graph will immediately encode itself to 340// the stream, and continue to record all subsequent operations. 341func (g *Graph) SetDebugWriter(w io.Writer) { 342 g.debug = &encoder{w: w} 343 g.debug.Encode(newMarshalGraph("root", g)) 344} 345 346// DebugVertexInfo encodes arbitrary information about a vertex in the graph 347// debug logs. 348func (g *Graph) DebugVertexInfo(v Vertex, info string) { 349 va := newVertexInfo(typeVertexInfo, v, info) 350 g.debug.Encode(va) 351} 352 353// DebugEdgeInfo encodes arbitrary information about an edge in the graph debug 354// logs. 355func (g *Graph) DebugEdgeInfo(e Edge, info string) { 356 ea := newEdgeInfo(typeEdgeInfo, e, info) 357 g.debug.Encode(ea) 358} 359 360// DebugVisitInfo records a visit to a Vertex during a walk operation. 361func (g *Graph) DebugVisitInfo(v Vertex, info string) { 362 vi := newVertexInfo(typeVisitInfo, v, info) 363 g.debug.Encode(vi) 364} 365 366// DebugOperation marks the start of a set of graph transformations in 367// the debug log, and returns a DebugOperationEnd func, which marks the end of 368// the operation in the log. Additional information can be added to the log via 369// the info parameter. 370// 371// The returned func's End method allows this method to be called from a single 372// defer statement: 373// defer g.DebugOperationBegin("OpName", "operating").End("") 374// 375// The returned function must be called to properly close the logical operation 376// in the logs. 377func (g *Graph) DebugOperation(operation string, info string) DebugOperationEnd { 378 return g.debug.BeginOperation(operation, info) 379} 380 381// VertexName returns the name of a vertex. 382func VertexName(raw Vertex) string { 383 switch v := raw.(type) { 384 case NamedVertex: 385 return v.Name() 386 case fmt.Stringer: 387 return fmt.Sprintf("%s", v) 388 default: 389 return fmt.Sprintf("%v", v) 390 } 391} 392