A trie structure, storing []byte elements, implemented in Go

trie.go

package trie

import (
	"container/vector"
	"sort"
)

// A 'set' structure, that can hold []byte objects.
// For any one []byte instance, it is either in the set or not.
type Trie struct {
	// How many instances of this []byte are in the set?
	// This value can be 0 if the []byte is just a prefix to a longer element.
	count uint

	children map[byte]*Trie
}

func New() *Trie {
	t := new(Trie)
	t.children = make(map[byte]*Trie)
	return t
}

// Is `s` is a member of the set?
// Returns the number of instances of `s` in the set.
func (self *Trie) Contains(s string) uint {
	if len(s) == 0 {
		return self.count
	}
	child, ok := self.children[s[0]]
	if !ok {
		return 0
	}
	return child.Contains(s[1:])
}

// Add `s` to the set.
// If `s` is already a member, this does nothing.
func (self *Trie) Add(s string) {
	if len(s) == 0 {
		self.count++
		return
	}
	child, ok := self.children[s[0]]
	if !ok {
		child = New()
		self.children[s[0]] = child
	}
	child.Add(s[1:])
}

// Remove `s` from the set.
// If there are no instances of `s`, this does nothing.
// Return value: whether the set is empty
// (i.e., has no instances and has no children)
func (self *Trie) Delete(s string) (empty bool) {
	if len(s) == 0 {
		self.count--
		return int(self.count)+len(self.children) <= 0
	}
	child, ok := self.children[s[0]]
	if ok && child.Delete(s[1:]) {
		// the child reports it is empty, hence can be deleted
		self.children[s[0]] = nil, false
	}
	return false
}

// The following ugliness is a wrapper for bytes,
// providing the necessary Less() method in order to sort them.
type lessByte struct {
	val byte
}

func (self *lessByte) Less(y interface{}) bool { return self.val < y.(*lessByte).val }
func byteToLess(b byte) *lessByte {
	// Wrap the byte.
	l := new(lessByte)
	l.val = b
	return l
}

// Retrieve all members of the set, in order.
func (self *Trie) Members(prefix string) *vector.StringVector {
	els := new(vector.StringVector)

	var i uint = 0
	for ; i < self.count; i++ {
		els.Push(prefix)
	}

	v := new(vector.Vector)
	for k, _ := range self.children {
		v.Push(byteToLess(k))
	}

	sort.Sort(v)

	for k := range v.Iter() {
		els.AppendVector(self.children[k.(*lessByte).val].Members(prefix + string(k.(*lessByte).val)))
	}

	return els
}

func Sort(s []string) *vector.StringVector {
	t := New()
	for _, s := range s {
		t.Add(s)
	}
	return t.Members("")
}