golang-github-ctdk-go-trie-dev 0.0~git20161027.0.6443fbc-1 / usr / share / gocode / src / github.com / ctdk / go-trie / gtrie / gtrie.go

//go:generate msgp
//msgp:shim rune as:int32 using:int32/rune
// Copyright (c) 2013 Mathieu Turcotte
// Licensed under the MIT license.

// Package gtrie provides a trie implementation based on a minimal acyclic
// finite-state automaton.
package gtrie

import (
	"errors"
	"sort"
	"strings"
	"fmt"
)

type NodeId int

type nodeIdGen struct {
	id NodeId
}

func (g *nodeIdGen) next() (next NodeId) {
	next = g.id
	g.id++
	return
}

// Represents a transition in the acyclic finite-state automaton. Each
// transition has one label and leads to one node.
type Transition struct {
	Child *Node
	Label rune
}

// Represents a node in the acyclic finite-state automaton.
type Node struct {
	Id          NodeId
	Terminal    bool
	Transitions []Transition
}

// Checks whether the node has children.
func (n *Node) HasChildren() bool {
	return len(n.Transitions) > 0
}

// Checks whether the node has a child for the given letter.
func (n *Node) HasChild(letter rune) bool {
	return n.GetChild(letter) != nil
}

func (n *Node) ChildKeys() []string {
	if !n.HasChildren() || n.Terminal {
		return nil
	}
	strs := make([]string, 0)
	for _, c := range n.Transitions {
		letter := c.Label
		child_strings := c.Child.ChildKeys()
		if child_strings != nil {
			for _, cs := range child_strings {
				strs = append(strs, fmt.Sprintf("%c%s", letter, cs))
			}
		} else {
			strs = append(strs, fmt.Sprintf("%c", letter))
		}
	}
	return strs
}

// Retrieves the child for the given letter. Returns nil if there is no child
// for this letter.
func (n *Node) GetChild(letter rune) (child *Node) {
	transitions := n.Transitions
	finder := func(i int) bool { return transitions[i].Label >= letter }
	// It is possible to use a binary search here because we know, by
	// construction, that the transitions are sorted by their labels.
	index := sort.Search(len(transitions), finder)
	if index < len(transitions) && transitions[index].Label == letter {
		child = transitions[index].Child
	}
	return
}

// Whether the node recognizes the given suffix. A suffix is accepted if there
// exists a path from the current node to a final node labeled with the suffix
// elements.
func (n *Node) Accepts(suffix string) bool {
	letters := []rune(suffix)
	current := n
	for i := 0; current != nil && i < len(letters); i++ {
		current = current.GetChild(letters[i])
	}
	return current != nil && current.Terminal
}

// Whether the given prefix is found in the tree. Most useful reading off the
// root node.
func (n *Node) HasPrefix(prefix string) (*Node, error) {
	letters := []rune(prefix)
	current := n
	for i := 0; current != nil && i < len(letters); i++ {
		current = current.GetChild(letters[i])
	}
	if current == nil {
		err := fmt.Errorf("%s not found", prefix)
		return nil, err
	}
	return current, nil
}

// Gets the number of nodes in the given automaton.
func Size(node *Node) int {
	ids := make(map[NodeId]bool)
	queue := []*Node{node}
	for len(queue) > 0 {
		node = queue[0]
		queue = queue[1:]
		ids[node.Id] = true
		for _, t := range node.Transitions {
			queue = append(queue, t.Child)
		}
	}
	return len(ids)
}

func newNode(idGen *nodeIdGen) *Node {
	return &Node{Id: idGen.next()}
}

func addTransition(node *Node, child *Node, letter rune) {
	node.Transitions = append(node.Transitions, Transition{child, letter})
}

func addChild(node *Node, letter rune, idGen *nodeIdGen) (child *Node) {
	child = node.GetChild(letter)
	if child == nil {
		child = newNode(idGen)
		addTransition(node, child, letter)
	}
	return
}

func getLastChild(node *Node) *Node {
	t := node.Transitions
	return t[len(t)-1].Child
}

func setLastChild(node *Node, last *Node) {
	t := node.Transitions
	t[len(t)-1].Child = last
}

type eqClass struct {
	terminal bool
	children string
}

// Obtains the equivalence class for this node, knowing that two nodes p and
// q belongs to the same class if and only if:
//  1. they are either both final or both nonfinal; and
//  2. they have the same number of outgoing transitions; and
//  3. corresponding outgoing transitions have the same labels; and
//  4. corresponding transitions lead to the same states.
func getEquivalenceClass(node *Node) (class eqClass) {
	children := []string{}
	for _, t := range node.Transitions {
		children = append(children, string(t.Label)+":"+string(t.Child.Id))
	}
	class.children = strings.Join(children, ";")
	class.terminal = node.Terminal
	return
}

type registry struct {
	// Mapping from equivalence class to node.
	eqv map[eqClass]*Node
	// Set of nodes that are registered.
	nodes map[*Node]bool
}

func newRegistery() (reg *registry) {
	reg = new(registry)
	reg.eqv = make(map[eqClass]*Node)
	reg.nodes = make(map[*Node]bool)
	return
}

func (r *registry) find(class eqClass) *Node {
	return r.eqv[class]
}

func (r *registry) register(class eqClass, node *Node) {
	r.eqv[class] = node
	r.nodes[node] = true
}

func (r *registry) registered(node *Node) bool {
	return r.nodes[node]
}

// Creates an acyclic finite-state automaton from a sorted list of words and
// returns the root node. Words can contain any unicode chararcters. An error
// will be returned if the list of words is not lexicographically sorted.
func Create(words []string) (automaton *Node, err error) {
	reg := newRegistery()
	idGen := new(nodeIdGen)
	automaton = newNode(idGen)

	if !sort.StringsAreSorted(words) {
		err = errors.New("the words are not sorted")
		return
	}

	for _, word := range words {
		insertWord(word, automaton, reg, idGen)
	}

	replaceOrRegister(automaton, reg)
	return
}

func insertWord(word string, automaton *Node, reg *registry, idGen *nodeIdGen) {
	letters := []rune(word)
	var last *Node

	if len(letters) == 0 {
		return
	}

	// Find last common state.
	for current := automaton; current != nil && len(letters) > 0; {
		last = current
		current = last.GetChild(letters[0])
		if current != nil {
			letters = letters[1:]
		}
	}

	// Minimize.
	if last.HasChildren() {
		replaceOrRegister(last, reg)
	}

	// Add suffix.
	for len(letters) > 0 {
		last = addChild(last, letters[0], idGen)
		letters = letters[1:]
	}

	last.Terminal = true
}

func replaceOrRegister(node *Node, reg *registry) {
	var child = getLastChild(node)

	if reg.registered(child) {
		return
	}

	if child.HasChildren() {
		replaceOrRegister(child, reg)
	}

	class := getEquivalenceClass(child)

	if eq := reg.find(class); eq != nil {
		setLastChild(node, eq)
	} else {
		reg.register(class, child)
	}
}