diff --git a/internal/graph/encode_dot.go b/internal/graph/encode_dot.go new file mode 100644 index 0000000..1c8f530 --- /dev/null +++ b/internal/graph/encode_dot.go @@ -0,0 +1,103 @@ +package graph + +import ( + "fmt" + "io" + "sort" +) + +// dotBuilder wraps an io.Writer and writes a graph in DOT format +type dotBuilder struct { + io.Writer +} + +func newDotBuilder(w io.Writer) *dotBuilder { + builder := dotBuilder{w} + if err := builder.init(); err != nil { + panic(err) + } + return &builder +} + +// init initializes a graph +func (b *dotBuilder) init() error { + _, err := fmt.Fprintln(b, `digraph G {`) + if err != nil { + return err + } + + _, err = fmt.Fprintln(b, `node [fontname="Helvetica,Arial,sans-serif"]`) + if err != nil { + return err + } + + return nil +} + +// finish closes the opening curly bracket of the current graph +func (b *dotBuilder) finish() error { + _, err := fmt.Fprintln(b, "}") + return err +} + +// addNode adds a node to the graph +func (b *dotBuilder) addNode(id int, label string) error { + attr := fmt.Sprintf(`label=%q`, label) + _, err := fmt.Fprintf(b, "n%d [%s]\n", id, attr) + return err +} + +// addEdge adds an edge to the graph +func (b *dotBuilder) addEdge(from, to int) error { + _, err := fmt.Fprintf(b, "n%d -> n%d\n", from, to) + return err +} + +// EncodeDOT encodes a graph in DOT format to enable visualization of the graph +func EncodeDOT[V Vertex](g *Graph[V], w io.Writer, sortVertices bool) error { + builder := newDotBuilder(w) + + vertexIds := make([]string, 0, len(g.verticesById)) + for id := range g.verticesById { + vertexIds = append(vertexIds, id) + } + if sortVertices { + // Sort the vertices so that the ordering of the output is deterministic, + // mainly for testing purposes. + sort.Strings(vertexIds) + } + + nodeIdsByVertex := make(map[string]int) + for i, id := range vertexIds { + err := builder.addNode(i, id) + if err != nil { + return fmt.Errorf("addNode(%d, %s): %w", i, id, err) + } + nodeIdsByVertex[id] = i + } + + for _, source := range vertexIds { + adjacentEdgesMap := g.edges[source] + + adjacentVerticesIds := make([]string, 0, len(adjacentEdgesMap)) + for tid := range adjacentEdgesMap { + adjacentVerticesIds = append(adjacentVerticesIds, tid) + } + if sortVertices { + sort.Strings(adjacentVerticesIds) + } + + for _, target := range adjacentVerticesIds { + if adjacentEdgesMap[target] { + err := builder.addEdge(nodeIdsByVertex[source], nodeIdsByVertex[target]) + if err != nil { + return fmt.Errorf( + "addEdge(%d, %d): %w", nodeIdsByVertex[source], nodeIdsByVertex[target], err, + ) + } + } + } + } + + return builder.finish() +} diff --git a/internal/graph/encode_dot_test.go b/internal/graph/encode_dot_test.go new file mode 100644 index 0000000..83282d3 --- /dev/null +++ b/internal/graph/encode_dot_test.go @@ -0,0 +1,72 @@ +package graph + +import ( + "bytes" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestDOTEncoding(t *testing.T) { + for _, tc := range []struct { + name string + adjList map[string][]string + expected string + }{ + { + name: "empty graph", + adjList: map[string][]string{}, + expected: `digraph G { +node [fontname="Helvetica,Arial,sans-serif"] +} +`, + }, + { + name: "multiple vertices", + adjList: map[string][]string{ + "v_0": {"v_1", "v_2"}, + "v_1": {"v_2"}, + "v_2": {"v_3"}, + "v_3": {}, + "v_4": {"v_0", "v_3"}, + }, + expected: `digraph G { +node [fontname="Helvetica,Arial,sans-serif"] +n0 [label="v_0"] +n1 [label="v_1"] +n2 [label="v_2"] +n3 [label="v_3"] +n4 [label="v_4"] +n0 -> n1 +n0 -> n2 +n1 -> n2 +n2 -> n3 +n4 -> n0 +n4 -> n3 +} +`, + }, + } { + t.Run(tc.name, func(t *testing.T) { + g := NewGraph[vertex]() + + // add vertices + for id := range tc.adjList { + g.AddVertex(NewV(id)) + } + + // add edges + for id, neighbors := range tc.adjList { + for _, neighborId := range neighbors { + assert.NoError(t, g.AddEdge(id, neighborId)) + } + } + + // encode to DOT + buf := bytes.Buffer{} + // sort vertices to ensure deterministic ordering of nodes and edges + assert.NoError(t, EncodeDOT(g, &buf, true)) + assert.Equal(t, tc.expected, buf.String()) + }) + } +}