smegmesh/pkg/dot/dot.go

// Graph allows the definition of a DOT graph in golang
package graph

import (
	"fmt"
	"hash/fnv"
	"strings"

	"github.com/tim-beatham/wgmesh/pkg/lib"
)

type GraphType string
type Shape string

const (
	GRAPH   GraphType = "graph"
	DIGRAPH GraphType = "digraph"
)

const (
	CIRCLE  Shape = "circle"
	STAR    Shape = "star"
	HEXAGON Shape = "hexagon"
)

type Graph interface {
	Dottable
	GetType() GraphType
}

type Cluster struct {
	Type  GraphType
	Name  string
	Label string
	nodes map[string]*Node
	edges map[string]Edge
}

type RootGraph struct {
	Type     GraphType
	Label    string
	nodes    map[string]*Node
	clusters map[string]*Cluster
	edges    map[string]Edge
}

type Node struct {
	Name  string
	Label string
	Shape Shape
	Size  int
}

type Edge interface {
	Dottable
}

type DirectedEdge struct {
	Name  string
	Label string
	From  string
	To    string
}

type UndirectedEdge struct {
	Name  string
	Label string
	From  string
	To    string
}

// Dottable means an implementer can convert the struct to DOT representation
type Dottable interface {
	GetDOT() (string, error)
}

func NewGraph(label string, graphType GraphType) *RootGraph {
	return &RootGraph{Type: graphType, Label: label, clusters: map[string]*Cluster{}, nodes: make(map[string]*Node), edges: make(map[string]Edge)}
}

// PutNode: puts a node in the graph
func (g *RootGraph) PutNode(name, label string, size int, shape Shape) error {
	_, exists := g.nodes[name]

	if exists {
		// If exists no need to add the ndoe
		return nil
	}

	g.nodes[name] = &Node{Name: name, Label: label, Size: size, Shape: shape}
	return nil
}

func (g *RootGraph) PutCluster(graph *Cluster) {
	g.clusters[graph.Label] = graph
}

func writeContituents[D Dottable](result *strings.Builder, elements ...D) error {
	for _, node := range elements {
		dot, err := node.GetDOT()

		if err != nil {
			return err
		}

		_, err = result.WriteString(dot)

		if err != nil {
			return err
		}
	}
	return nil
}

func (g *RootGraph) GetDOT() (string, error) {
	var result strings.Builder

	result.WriteString(fmt.Sprintf("%s {\n", g.Type))
	result.WriteString("node [colorscheme=set312];\n")
	result.WriteString("layout = fdp;\n")
	nodes := lib.MapValues(g.nodes)
	edges := lib.MapValues(g.edges)
	writeContituents(&result, nodes...)
	writeContituents(&result, edges...)

	for _, cluster := range g.clusters {
		clusterDOT, err := cluster.GetDOT()

		if err != nil {
			return "", err
		}

		result.WriteString(clusterDOT)
	}

	result.WriteString("}")
	return result.String(), nil
}

// GetType implements Graph.
func (r *RootGraph) GetType() GraphType {
	return r.Type
}

func constructEdge(graph Graph, name, label, from, to string) Edge {
	switch graph.GetType() {
	case DIGRAPH:
		return &DirectedEdge{Name: name, Label: label, From: from, To: to}
	default:
		return &UndirectedEdge{Name: name, Label: label, From: from, To: to}
	}
}

// AddEdge: adds an edge between two nodes in the graph
func (g *RootGraph) AddEdge(name string, label string, from string, to string) error {
	g.edges[name] = constructEdge(g, name, label, from, to)
	return nil
}

const numColours = 12

func (n *Node) hash() int {
	h := fnv.New32a()
	h.Write([]byte(n.Name))
	return (int(h.Sum32()) % numColours) + 1
}

func (n *Node) GetDOT() (string, error) {
	return fmt.Sprintf("node[label=\"%s\",shape=%s, style=\"filled\", fillcolor=%d, width=%d, height=%d, fixedsize=true] \"%s\";\n",
		n.Label, n.Shape, n.hash(), n.Size, n.Size, n.Name), nil
}

func (e *DirectedEdge) GetDOT() (string, error) {
	return fmt.Sprintf("\"%s\" -> \"%s\" [label=\"%s\"];\n", e.From, e.To, e.Label), nil
}

func (e *UndirectedEdge) GetDOT() (string, error) {
	return fmt.Sprintf("\"%s\" -- \"%s\" [label=\"%s\"];\n", e.From, e.To, e.Label), nil
}

// AddEdge: adds an edge between two nodes in the graph
func (g *Cluster) AddEdge(name string, label string, from string, to string) error {
	g.edges[name] = constructEdge(g, name, label, from, to)
	return nil
}

// PutNode: puts a node in the graph
func (g *Cluster) PutNode(name, label string, size int, shape Shape) error {
	_, exists := g.nodes[name]

	if exists {
		// If exists no need to add the ndoe
		return nil
	}

	g.nodes[name] = &Node{Name: name, Label: label, Shape: shape, Size: size}
	return nil
}

func (g *Cluster) GetDOT() (string, error) {
	var builder strings.Builder

	builder.WriteString(fmt.Sprintf("subgraph \"cluster%s\" {\n", g.Label))
	builder.WriteString(fmt.Sprintf("label = \"%s\"\n", g.Label))
	nodes := lib.MapValues(g.nodes)
	edges := lib.MapValues(g.edges)
	writeContituents(&builder, nodes...)
	writeContituents(&builder, edges...)

	builder.WriteString("}\n")
	return builder.String(), nil
}

func (g *Cluster) GetType() GraphType {
	return g.Type
}

func NewSubGraph(name string, label string, graphType GraphType) *Cluster {
	return &Cluster{
		Label: name,
		Type:  graphType,
		Name:  name,
		nodes: make(map[string]*Node),
		edges: make(map[string]Edge),
	}
}