trout

Paddy 2015-12-13 Parent:bf38b050b6c4

trout/trie.go

Publish under MIT license. Add a license so this package can officially be used by any project, basically.

paddy@0	1 package trout
paddy@0	2
paddy@0	3 import (
paddy@0	4 "net/http"
paddy@0	5 "sort"
paddy@0	6 "sync"
paddy@0	7 )
paddy@0	8
paddy@0	9 type trie struct {
paddy@0	10 branch *branch
paddy@0	11 sync.RWMutex
paddy@0	12 }
paddy@0	13
paddy@0	14 func (t *trie) match(input []string) ([]int, bool) {
paddy@0	15 if t.branch == nil {
paddy@0	16 t.branch = &branch{}
paddy@0	17 }
paddy@0	18 b := t.branch
paddy@0	19 path := []int{}
paddy@0	20 offset := 0
paddy@0	21 num := len(input)
paddy@0	22 for i := 0; i < num; i++ {
paddy@0	23 // if we're on a nil branch, we're at the end of our line
paddy@0	24 if b == nil {
paddy@0	25 return path, false
paddy@0	26 }
paddy@0	27 offset = pickNextBranch(b, offset, input[i])
paddy@0	28 if offset == -1 {
paddy@0	29 if len(path) == 0 {
paddy@0	30 // can't find it, bail
paddy@0	31 return path, false
paddy@0	32 }
paddy@0	33 // no match, backup
paddy@0	34 path, offset = backup(path)
paddy@0	35 offset = offset + 1 // we want the next index from the previously matched one
paddy@0	36 b = b.parent // we need to be picking a branch from our parent, again
paddy@0	37 i = i - 2 // back up to the choice before the one that got us here
paddy@0	38 } else {
paddy@0	39 path = append(path, offset)
paddy@0	40 b = b.children[offset]
paddy@0	41 offset = 0
paddy@0	42 }
paddy@0	43 }
paddy@0	44 return path, true
paddy@0	45 }
paddy@0	46
paddy@0	47 type branch struct {
paddy@0	48 parent *branch
paddy@0	49 children []*branch
paddy@0	50 key string
paddy@0	51 isParam bool
paddy@0	52 methods map[string]http.Handler
paddy@0	53 }
paddy@0	54
paddy@0	55 func (b *branch) Less(i, j int) bool {
paddy@0	56 if b.children[i].key == b.children[j].key {
paddy@0	57 return b.children[j].isParam && !b.children[i].isParam
paddy@0	58 }
paddy@0	59 return b.children[i].key < b.children[j].key
paddy@0	60 }
paddy@0	61
paddy@0	62 func (b *branch) Len() int {
paddy@0	63 return len(b.children)
paddy@0	64 }
paddy@0	65
paddy@0	66 func (b *branch) Swap(i, j int) {
paddy@0	67 b.children[i], b.children[j] = b.children[j], b.children[i]
paddy@0	68 }
paddy@0	69
paddy@0	70 func (b branch) addChild(key string, param bool) branch {
paddy@0	71 child := &branch{key: key, parent: b, isParam: param}
paddy@0	72 if b.children == nil {
paddy@0	73 b.children = []*branch{child}
paddy@0	74 return child
paddy@0	75 }
paddy@0	76 b.children = append(b.children, child)
paddy@0	77 sort.Sort(b)
paddy@0	78 return child
paddy@0	79 }
paddy@0	80
paddy@0	81 func (b *branch) check(input string) bool {
paddy@0	82 if b.isParam && b.key != "" {
paddy@0	83 return true
paddy@0	84 }
paddy@0	85 if b.key == input {
paddy@0	86 return true
paddy@0	87 }
paddy@0	88 return false
paddy@0	89 }
paddy@0	90
paddy@0	91 func (b *branch) setHandler(method string, handler http.Handler) {
paddy@1	92 if b.methods == nil {
paddy@1	93 b.methods = map[string]http.Handler{}
paddy@1	94 }
paddy@0	95 b.methods[method] = handler
paddy@0	96 }
paddy@0	97
paddy@0	98 func pickNextBranch(b *branch, offset int, input string) int {
paddy@0	99 count := len(b.children)
paddy@0	100 for i := offset; i < count; i++ {
paddy@0	101 if b.children[i].check(input) {
paddy@0	102 return i
paddy@0	103 }
paddy@0	104 }
paddy@0	105 return -1
paddy@0	106 }
paddy@0	107
paddy@0	108 func backup(path []int) ([]int, int) {
paddy@0	109 last := len(path) - 1
paddy@0	110 pos := path[last]
paddy@0	111 path = path[:last]
paddy@0	112 return path, pos
paddy@0	113 }