smegmesh/pkg/graph/graph.go

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

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

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

type GraphType string
type Shape string

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

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

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

type Node struct {
	Name  string
	Shape Shape
}

type Edge interface {
	Dottable
}

type DirectedEdge struct {
	Label string
	From  *Node
	To    *Node
}

type UndirectedEdge struct {
	Label string
	From  *Node
	To    *Node
}

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

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

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

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

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

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 *Graph) GetDOT() (string, error) {
	var result strings.Builder

	_, err := result.WriteString(fmt.Sprintf("%s {\n", g.Type))

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

	_, err = result.WriteString("node [colorscheme=set312];\n")

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

	nodes := lib.MapValues(g.nodes)

	err = writeContituents(&result, nodes...)

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

	err = writeContituents(&result, g.edges...)

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

	_, err = result.WriteString("}")

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

	return result.String(), nil
}

func (g *Graph) constructEdge(label string, from *Node, to *Node) Edge {
	switch g.Type {
	case DIGRAPH:
		return &DirectedEdge{Label: label, From: from, To: to}
	default:
		return &UndirectedEdge{Label: label, From: from, To: to}
	}
}

// AddEdge: adds an edge between two nodes in the graph
func (g *Graph) AddEdge(label string, from string, to string) error {
	fromNode, exists := g.nodes[from]

	if !exists {
		return errors.New(fmt.Sprintf("Node %s does not exist", from))
	}

	toNode, exists := g.nodes[to]

	if !exists {
		return errors.New(fmt.Sprintf("Node %s does not exist", to))
	}

	g.edges = append(g.edges, g.constructEdge(label, fromNode, toNode))
	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[shape=%s, style=\"filled\", fillcolor=%d] %s;\n",
		n.Shape, n.hash(), n.Name), nil
}

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

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