wiki/graph.go

package main

import (
	"encoding/json"
	"os"
	"path/filepath"
	"strings"
)

type graphBuilder struct {
	refs    Refs
	nodeMap NodeMap
	nodes   []Node
	edges   []Edge
}

type NodeMap map[string]int

func NewGraphBuilder(mp PagesRepository) (*graphBuilder, error) {
	refs := make(Refs)

	f, err := os.Open(filepath.Join(mp.(*FilePages).dirname, "backrefs.json"))
	if err != nil {
		return nil, err
	}

	defer f.Close()

	err = json.NewDecoder(f).Decode(&refs)
	if err != nil {
		return nil, err
	}

	nodeMap := prepareNodeMap(refs)
	return &graphBuilder{
		refs:    refs,
		nodeMap: nodeMap,
		nodes:   nil,
		edges:   nil,
	}, nil
}

func (gb *graphBuilder) prepareGraph() error {
	gb.nodes = prepareNodes(gb.nodeMap)
	gb.edges = prepareEdges(gb.refs, gb.nodeMap)
	return nil
}

func (gb *graphBuilder) RemoveNode(name string) {
	delete(gb.nodeMap, name)
}

func (gb *graphBuilder) RemoveNodeWithSuffix(suffix string) {
	for k, _ := range gb.nodeMap {
		if strings.HasSuffix(k, suffix) {
			delete(gb.nodeMap, k)
		}
	}
}

func prepareEdges(refs Refs, nodeMap NodeMap) []Edge {
	var edges []Edge
	edgeSet := make(map[Edge]bool)
	for key, references := range refs {
		if toID, e := nodeMap[key]; e {
			for _, item := range references {
				if fromID, e := nodeMap[item.Name]; e {
					edge := Edge{
						From: fromID,
						To:   toID,
					}
					if _, e := edgeSet[edge]; !e {
						edgeSet[edge] = true
						edges = append(edges, edge)
					}
				}
			}
		}
	}
	return edges
}

func prepareNodes(nodeMap NodeMap) []Node {
	var nodes []Node
	for name, id := range nodeMap {
		nodes = append(nodes, Node{
			Id:    id,
			Label: name,
		})
	}
	return nodes
}

func prepareNodeMap(refs Refs) NodeMap {
	nodeCount := 1
	nodeMap := make(NodeMap)

	for key, references := range refs {
		if _, e := nodeMap[key]; !e {
			nodeMap[key] = nodeCount
			nodeCount += 1
		}
		for _, item := range references {
			if _, e := nodeMap[item.Name]; !e {
				nodeMap[item.Name] = nodeCount
				nodeCount += 1
			}
		}
	}

	return nodeMap
}