pathfinder.go

package main

import (
	"database/sql"
	"fmt"
	"github.com/Sereal/Sereal/Go/sereal"
	_ "github.com/go-sql-driver/mysql"
	"io/ioutil"
	"os"
	"strings"
)

const DATA_FILE = "graph.dat"

func checkErr(e error) {
	if e != nil {
		panic(e)
	}
}

type Vertex struct {
	Title     string
	Id        int
	Neighbors []*Vertex
}

func (from *Vertex) AddNeighbour(to *Vertex) {
	if to != nil {
		from.Neighbors = append(from.Neighbors, to)
	}
}

func NewVertex(id int, title string) *Vertex {
	neighbors := make([]*Vertex, 0)
	return &Vertex{Id: id, Title: title, Neighbors: neighbors}
}

func constructPath(visited map[int]*Vertex, target *Vertex, from *Vertex) []string {
	path := make([]string, 0)

	path = append(path, target.Title)

	curr, present := visited[target.Id]
	for ; present && curr != from; curr, present = visited[curr.Id] {
		path = append(path, curr.Title)
	}

	path = append(path, from.Title)

	for i := len(path)/2 - 1; i >= 0; i-- {
		j := len(path) - 1 - i
		path[i], path[j] = path[j], path[i]
	}
	return path
}

func (from *Vertex) BFS(to *Vertex) {
	visited := make(map[int]*Vertex)
	queue := make([]*Vertex, 0, 1)

	queue = append(queue, from)

	for v := queue[0]; len(queue) > 0; v = queue[0] {
		if v == to {
			fmt.Println(strings.Join(constructPath(visited, to, from), " -> "))
			return
		}

		for _, n := range v.Neighbors {
			if _, seen := visited[n.Id]; !seen {
				queue = append(queue, n)
				visited[n.Id] = v
			}
		}
		queue = queue[1:]
	}
}

func buildVertices(db *sql.DB, vertices map[string]*Vertex, idToTitle map[int]string) {
	rows, err := db.Query("SELECT page_id, page_title FROM page WHERE page_namespace = 0")
	checkErr(err)

	for rows.Next() {
		var id int
		var title string

		err := rows.Scan(&id, &title)
		checkErr(err)

		idToTitle[id] = title
		vertices[title] = NewVertex(id, title)
	}

	fmt.Println("Done building vertices")
}

func buildEdges(db *sql.DB, vertices map[string]*Vertex, idToTitle map[int]string) {
	rows, err := db.Query("SELECT pl_from, pl_title, pl_namespace, pl_from_namespace FROM pagelinks")
	checkErr(err)

	for rows.Next() {
		var fromId int
		var toTitle string
		var toNamespace int
		var fromNamespace int

		err := rows.Scan(&fromId, &toTitle, &toNamespace, &fromNamespace)
		checkErr(err)
		if fromNamespace == 0 && toNamespace == 0 {
			if fromTitle, found := idToTitle[fromId]; found {
				vertices[fromTitle].AddNeighbour(vertices[toTitle])
			}
		}
	}
	fmt.Println("Done building edges")
}

func buildGraph() map[string]*Vertex {
	if _, err := os.Stat(DATA_FILE); err == nil {
		var vertices map[string]*Vertex

		data, err := ioutil.ReadFile(DATA_FILE)
		checkErr(err)

		err = sereal.Unmarshal(data, &vertices)
		checkErr(err)

		return vertices
	} else {
		db, err := sql.Open("mysql", "root:@/wikilinks")
		defer db.Close()

		vertices := make(map[string]*Vertex)
		idToTitle := make(map[int]string)
		buildVertices(db, vertices, idToTitle)

		buildEdges(db, vertices, idToTitle)

		data, err := sereal.Marshal(vertices)
		checkErr(err)

		err = ioutil.WriteFile(DATA_FILE, data, 0600)
		checkErr(err)

		return vertices
	}
}

func findPath(graph map[string]*Vertex, from string, to string) {
	f, fromFound := graph[from]
	t, toFound := graph[to]

	if fromFound && toFound {
		f.BFS(t)
	}
}

func main() {
	graph := buildGraph()
	findPath(graph, os.Args[1], os.Args[2])
}