1package dag
2
3import (
4	"bytes"
5	"fmt"
6	"sort"
7	"strings"
8)
9
10// DotOpts are the options for generating a dot formatted Graph.
11type DotOpts struct {
12	// Allows some nodes to decide to only show themselves when the user has
13	// requested the "verbose" graph.
14	Verbose bool
15
16	// Highlight Cycles
17	DrawCycles bool
18
19	// How many levels to expand modules as we draw
20	MaxDepth int
21
22	// use this to keep the cluster_ naming convention from the previous dot writer
23	cluster bool
24}
25
26// GraphNodeDotter can be implemented by a node to cause it to be included
27// in the dot graph. The Dot method will be called which is expected to
28// return a representation of this node.
29type GraphNodeDotter interface {
30	// Dot is called to return the dot formatting for the node.
31	// The first parameter is the title of the node.
32	// The second parameter includes user-specified options that affect the dot
33	// graph. See GraphDotOpts below for details.
34	DotNode(string, *DotOpts) *DotNode
35}
36
37// DotNode provides a structure for Vertices to return in order to specify their
38// dot format.
39type DotNode struct {
40	Name  string
41	Attrs map[string]string
42}
43
44// Returns the DOT representation of this Graph.
45func (g *marshalGraph) Dot(opts *DotOpts) []byte {
46	if opts == nil {
47		opts = &DotOpts{
48			DrawCycles: true,
49			MaxDepth:   -1,
50			Verbose:    true,
51		}
52	}
53
54	var w indentWriter
55	w.WriteString("digraph {\n")
56	w.Indent()
57
58	// some dot defaults
59	w.WriteString(`compound = "true"` + "\n")
60	w.WriteString(`newrank = "true"` + "\n")
61
62	// the top level graph is written as the first subgraph
63	w.WriteString(`subgraph "root" {` + "\n")
64	g.writeBody(opts, &w)
65
66	// cluster isn't really used other than for naming purposes in some graphs
67	opts.cluster = opts.MaxDepth != 0
68	maxDepth := opts.MaxDepth
69	if maxDepth == 0 {
70		maxDepth = -1
71	}
72
73	for _, s := range g.Subgraphs {
74		g.writeSubgraph(s, opts, maxDepth, &w)
75	}
76
77	w.Unindent()
78	w.WriteString("}\n")
79	return w.Bytes()
80}
81
82func (v *marshalVertex) dot(g *marshalGraph, opts *DotOpts) []byte {
83	var buf bytes.Buffer
84	graphName := g.Name
85	if graphName == "" {
86		graphName = "root"
87	}
88
89	name := v.Name
90	attrs := v.Attrs
91	if v.graphNodeDotter != nil {
92		node := v.graphNodeDotter.DotNode(name, opts)
93		if node == nil {
94			return []byte{}
95		}
96
97		newAttrs := make(map[string]string)
98		for k, v := range attrs {
99			newAttrs[k] = v
100		}
101		for k, v := range node.Attrs {
102			newAttrs[k] = v
103		}
104
105		name = node.Name
106		attrs = newAttrs
107	}
108
109	buf.WriteString(fmt.Sprintf(`"[%s] %s"`, graphName, name))
110	writeAttrs(&buf, attrs)
111	buf.WriteByte('\n')
112
113	return buf.Bytes()
114}
115
116func (e *marshalEdge) dot(g *marshalGraph) string {
117	var buf bytes.Buffer
118	graphName := g.Name
119	if graphName == "" {
120		graphName = "root"
121	}
122
123	sourceName := g.vertexByID(e.Source).Name
124	targetName := g.vertexByID(e.Target).Name
125	s := fmt.Sprintf(`"[%s] %s" -> "[%s] %s"`, graphName, sourceName, graphName, targetName)
126	buf.WriteString(s)
127	writeAttrs(&buf, e.Attrs)
128
129	return buf.String()
130}
131
132func cycleDot(e *marshalEdge, g *marshalGraph) string {
133	return e.dot(g) + ` [color = "red", penwidth = "2.0"]`
134}
135
136// Write the subgraph body. The is recursive, and the depth argument is used to
137// record the current depth of iteration.
138func (g *marshalGraph) writeSubgraph(sg *marshalGraph, opts *DotOpts, depth int, w *indentWriter) {
139	if depth == 0 {
140		return
141	}
142	depth--
143
144	name := sg.Name
145	if opts.cluster {
146		// we prefix with cluster_ to match the old dot output
147		name = "cluster_" + name
148		sg.Attrs["label"] = sg.Name
149	}
150	w.WriteString(fmt.Sprintf("subgraph %q {\n", name))
151	sg.writeBody(opts, w)
152
153	for _, sg := range sg.Subgraphs {
154		g.writeSubgraph(sg, opts, depth, w)
155	}
156}
157
158func (g *marshalGraph) writeBody(opts *DotOpts, w *indentWriter) {
159	w.Indent()
160
161	for _, as := range attrStrings(g.Attrs) {
162		w.WriteString(as + "\n")
163	}
164
165	// list of Vertices that aren't to be included in the dot output
166	skip := map[string]bool{}
167
168	for _, v := range g.Vertices {
169		if v.graphNodeDotter == nil {
170			skip[v.ID] = true
171			continue
172		}
173
174		w.Write(v.dot(g, opts))
175	}
176
177	var dotEdges []string
178
179	if opts.DrawCycles {
180		for _, c := range g.Cycles {
181			if len(c) < 2 {
182				continue
183			}
184
185			for i, j := 0, 1; i < len(c); i, j = i+1, j+1 {
186				if j >= len(c) {
187					j = 0
188				}
189				src := c[i]
190				tgt := c[j]
191
192				if skip[src.ID] || skip[tgt.ID] {
193					continue
194				}
195
196				e := &marshalEdge{
197					Name:   fmt.Sprintf("%s|%s", src.Name, tgt.Name),
198					Source: src.ID,
199					Target: tgt.ID,
200					Attrs:  make(map[string]string),
201				}
202
203				dotEdges = append(dotEdges, cycleDot(e, g))
204				src = tgt
205			}
206		}
207	}
208
209	for _, e := range g.Edges {
210		dotEdges = append(dotEdges, e.dot(g))
211	}
212
213	// srot these again to match the old output
214	sort.Strings(dotEdges)
215
216	for _, e := range dotEdges {
217		w.WriteString(e + "\n")
218	}
219
220	w.Unindent()
221	w.WriteString("}\n")
222}
223
224func writeAttrs(buf *bytes.Buffer, attrs map[string]string) {
225	if len(attrs) > 0 {
226		buf.WriteString(" [")
227		buf.WriteString(strings.Join(attrStrings(attrs), ", "))
228		buf.WriteString("]")
229	}
230}
231
232func attrStrings(attrs map[string]string) []string {
233	strings := make([]string, 0, len(attrs))
234	for k, v := range attrs {
235		strings = append(strings, fmt.Sprintf("%s = %q", k, v))
236	}
237	sort.Strings(strings)
238	return strings
239}
240
241// Provide a bytes.Buffer like structure, which will indent when starting a
242// newline.
243type indentWriter struct {
244	bytes.Buffer
245	level int
246}
247
248func (w *indentWriter) indent() {
249	newline := []byte("\n")
250	if !bytes.HasSuffix(w.Bytes(), newline) {
251		return
252	}
253	for i := 0; i < w.level; i++ {
254		w.Buffer.WriteString("\t")
255	}
256}
257
258// Indent increases indentation by 1
259func (w *indentWriter) Indent() { w.level++ }
260
261// Unindent decreases indentation by 1
262func (w *indentWriter) Unindent() { w.level-- }
263
264// the following methods intercecpt the byte.Buffer writes and insert the
265// indentation when starting a new line.
266func (w *indentWriter) Write(b []byte) (int, error) {
267	w.indent()
268	return w.Buffer.Write(b)
269}
270
271func (w *indentWriter) WriteString(s string) (int, error) {
272	w.indent()
273	return w.Buffer.WriteString(s)
274}
275func (w *indentWriter) WriteByte(b byte) error {
276	w.indent()
277	return w.Buffer.WriteByte(b)
278}
279func (w *indentWriter) WriteRune(r rune) (int, error) {
280	w.indent()
281	return w.Buffer.WriteRune(r)
282}
283