internal/dag/graph.go

package dag

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"sort"
)

// Graph is used to represent a dependency graph.
type Graph struct {
	vertices  *Set
	edges     *Set
	downEdges map[interface{}]*Set
	upEdges   map[interface{}]*Set

	// JSON encoder for recording debug information
	debug *encoder
}

// Subgrapher allows a Vertex to be a Graph itself, by returning a Grapher.
type Subgrapher interface {
	Subgraph() Grapher
}

// A Grapher is any type that returns a Grapher, mainly used to identify
// dag.Graph and dag.AcyclicGraph.  In the case of Graph and AcyclicGraph, they
// return themselves.
type Grapher interface {
	DirectedGraph() Grapher
}

// Vertex of the graph.
type Vertex interface{}

// NamedVertex is an optional interface that can be implemented by Vertex
// to give it a human-friendly name that is used for outputting the graph.
type NamedVertex interface {
	Vertex
	Name() string
}

func (g *Graph) DirectedGraph() Grapher {
	return g
}

// Vertices returns the list of all the vertices in the graph.
func (g *Graph) Vertices() []Vertex {
	list := g.vertices.List()
	result := make([]Vertex, len(list))
	for i, v := range list {
		result[i] = v.(Vertex)
	}

	return result
}

// Edges returns the list of all the edges in the graph.
func (g *Graph) Edges() []Edge {
	list := g.edges.List()
	result := make([]Edge, len(list))
	for i, v := range list {
		result[i] = v.(Edge)
	}

	return result
}

// EdgesFrom returns the list of edges from the given source.
func (g *Graph) EdgesFrom(v Vertex) []Edge {
	var result []Edge
	from := hashcode(v)
	for _, e := range g.Edges() {
		if hashcode(e.Source()) == from {
			result = append(result, e)
		}
	}

	return result
}

// EdgesTo returns the list of edges to the given target.
func (g *Graph) EdgesTo(v Vertex) []Edge {
	var result []Edge
	search := hashcode(v)
	for _, e := range g.Edges() {
		if hashcode(e.Target()) == search {
			result = append(result, e)
		}
	}

	return result
}

// HasVertex checks if the given Vertex is present in the graph.
func (g *Graph) HasVertex(v Vertex) bool {
	return g.vertices.Include(v)
}

// HasEdge checks if the given Edge is present in the graph.
func (g *Graph) HasEdge(e Edge) bool {
	return g.edges.Include(e)
}

// Add adds a vertex to the graph. This is safe to call multiple time with
// the same Vertex.
func (g *Graph) Add(v Vertex) Vertex {
	g.init()
	g.vertices.Add(v)
	g.debug.Add(v)
	return v
}

// Remove removes a vertex from the graph. This will also remove any
// edges with this vertex as a source or target.
func (g *Graph) Remove(v Vertex) Vertex {
	// Delete the vertex itself
	g.vertices.Delete(v)
	g.debug.Remove(v)

	// Delete the edges to non-existent things
	for _, target := range g.DownEdges(v).List() {
		g.RemoveEdge(BasicEdge(v, target))
	}
	for _, source := range g.UpEdges(v).List() {
		g.RemoveEdge(BasicEdge(source, v))
	}

	return nil
}

// Replace replaces the original Vertex with replacement. If the original
// does not exist within the graph, then false is returned. Otherwise, true
// is returned.
func (g *Graph) Replace(original, replacement Vertex) bool {
	// If we don't have the original, we can't do anything
	if !g.vertices.Include(original) {
		return false
	}

	defer g.debug.BeginOperation("Replace", "").End("")

	// If they're the same, then don't do anything
	if original == replacement {
		return true
	}

	// Add our new vertex, then copy all the edges
	g.Add(replacement)
	for _, target := range g.DownEdges(original).List() {
		g.Connect(BasicEdge(replacement, target))
	}
	for _, source := range g.UpEdges(original).List() {
		g.Connect(BasicEdge(source, replacement))
	}

	// Remove our old vertex, which will also remove all the edges
	g.Remove(original)

	return true
}

// RemoveEdge removes an edge from the graph.
func (g *Graph) RemoveEdge(edge Edge) {
	g.init()
	g.debug.RemoveEdge(edge)

	// Delete the edge from the set
	g.edges.Delete(edge)

	// Delete the up/down edges
	if s, ok := g.downEdges[hashcode(edge.Source())]; ok {
		s.Delete(edge.Target())
	}
	if s, ok := g.upEdges[hashcode(edge.Target())]; ok {
		s.Delete(edge.Source())
	}
}

// DownEdges returns the outward edges from the source Vertex v.
func (g *Graph) DownEdges(v Vertex) *Set {
	g.init()
	return g.downEdges[hashcode(v)]
}

// UpEdges returns the inward edges to the destination Vertex v.
func (g *Graph) UpEdges(v Vertex) *Set {
	g.init()
	return g.upEdges[hashcode(v)]
}

// Connect adds an edge with the given source and target. This is safe to
// call multiple times with the same value. Note that the same value is
// verified through pointer equality of the vertices, not through the
// value of the edge itself.
func (g *Graph) Connect(edge Edge) {
	g.init()
	g.debug.Connect(edge)

	source := edge.Source()
	target := edge.Target()
	sourceCode := hashcode(source)
	targetCode := hashcode(target)

	// Do we have this already? If so, don't add it again.
	if s, ok := g.downEdges[sourceCode]; ok && s.Include(target) {
		return
	}

	// Add the edge to the set
	g.edges.Add(edge)

	// Add the down edge
	s, ok := g.downEdges[sourceCode]
	if !ok {
		s = new(Set)
		g.downEdges[sourceCode] = s
	}
	s.Add(target)

	// Add the up edge
	s, ok = g.upEdges[targetCode]
	if !ok {
		s = new(Set)
		g.upEdges[targetCode] = s
	}
	s.Add(source)
}

// String outputs some human-friendly output for the graph structure.
func (g *Graph) StringWithNodeTypes() string {
	var buf bytes.Buffer

	// Build the list of node names and a mapping so that we can more
	// easily alphabetize the output to remain deterministic.
	vertices := g.Vertices()
	names := make([]string, 0, len(vertices))
	mapping := make(map[string]Vertex, len(vertices))
	for _, v := range vertices {
		name := VertexName(v)
		names = append(names, name)
		mapping[name] = v
	}
	sort.Strings(names)

	// Write each node in order...
	for _, name := range names {
		v := mapping[name]
		targets := g.downEdges[hashcode(v)]

		buf.WriteString(fmt.Sprintf("%s - %T\n", name, v))

		// Alphabetize dependencies
		deps := make([]string, 0, targets.Len())
		targetNodes := make(map[string]Vertex)
		for _, target := range targets.List() {
			dep := VertexName(target)
			deps = append(deps, dep)
			targetNodes[dep] = target
		}
		sort.Strings(deps)

		// Write dependencies
		for _, d := range deps {
			buf.WriteString(fmt.Sprintf("  %s - %T\n", d, targetNodes[d]))
		}
	}

	return buf.String()
}

// String outputs some human-friendly output for the graph structure.
func (g *Graph) String() string {
	var buf bytes.Buffer

	// Build the list of node names and a mapping so that we can more
	// easily alphabetize the output to remain deterministic.
	vertices := g.Vertices()
	names := make([]string, 0, len(vertices))
	mapping := make(map[string]Vertex, len(vertices))
	for _, v := range vertices {
		name := VertexName(v)
		names = append(names, name)
		mapping[name] = v
	}
	sort.Strings(names)

	// Write each node in order...
	for _, name := range names {
		v := mapping[name]
		targets := g.downEdges[hashcode(v)]

		buf.WriteString(fmt.Sprintf("%s\n", name))

		// Alphabetize dependencies
		deps := make([]string, 0, targets.Len())
		for _, target := range targets.List() {
			deps = append(deps, VertexName(target))
		}
		sort.Strings(deps)

		// Write dependencies
		for _, d := range deps {
			buf.WriteString(fmt.Sprintf("  %s\n", d))
		}
	}

	return buf.String()
}

func (g *Graph) init() {
	if g.vertices == nil {
		g.vertices = new(Set)
	}
	if g.edges == nil {
		g.edges = new(Set)
	}
	if g.downEdges == nil {
		g.downEdges = make(map[interface{}]*Set)
	}
	if g.upEdges == nil {
		g.upEdges = make(map[interface{}]*Set)
	}
}

// Dot returns a dot-formatted representation of the Graph.
func (g *Graph) Dot(opts *DotOpts) []byte {
	return newMarshalGraph("", g).Dot(opts)
}

// MarshalJSON returns a JSON representation of the entire Graph.
func (g *Graph) MarshalJSON() ([]byte, error) {
	dg := newMarshalGraph("root", g)
	return json.MarshalIndent(dg, "", "  ")
}

// SetDebugWriter sets the io.Writer where the Graph will record debug
// information. After this is set, the graph will immediately encode itself to
// the stream, and continue to record all subsequent operations.
func (g *Graph) SetDebugWriter(w io.Writer) {
	g.debug = &encoder{w: w}
	g.debug.Encode(newMarshalGraph("root", g))
}

// DebugVertexInfo encodes arbitrary information about a vertex in the graph
// debug logs.
func (g *Graph) DebugVertexInfo(v Vertex, info string) {
	va := newVertexInfo(typeVertexInfo, v, info)
	g.debug.Encode(va)
}

// DebugEdgeInfo encodes arbitrary information about an edge in the graph debug
// logs.
func (g *Graph) DebugEdgeInfo(e Edge, info string) {
	ea := newEdgeInfo(typeEdgeInfo, e, info)
	g.debug.Encode(ea)
}

// DebugVisitInfo records a visit to a Vertex during a walk operation.
func (g *Graph) DebugVisitInfo(v Vertex, info string) {
	vi := newVertexInfo(typeVisitInfo, v, info)
	g.debug.Encode(vi)
}

// DebugOperation marks the start of a set of graph transformations in
// the debug log, and returns a DebugOperationEnd func, which marks the end of
// the operation in the log. Additional information can be added to the log via
// the info parameter.
//
// The returned func's End method allows this method to be called from a single
// defer statement:
//     defer g.DebugOperationBegin("OpName", "operating").End("")
//
// The returned function must be called to properly close the logical operation
// in the logs.
func (g *Graph) DebugOperation(operation string, info string) DebugOperationEnd {
	return g.debug.BeginOperation(operation, info)
}

// VertexName returns the name of a vertex.
func VertexName(raw Vertex) string {
	switch v := raw.(type) {
	case NamedVertex:
		return v.Name()
	case fmt.Stringer:
		return fmt.Sprintf("%s", v)
	default:
		return fmt.Sprintf("%v", v)
	}
}