squid-rewriter/trie/trie.go

package trie

func commonPrefixLength(s1 string, s2 string) int {
	var i, l int

	l1 := len(s1)
	l2 := len(s2)
	if l1 < l2 {
		l = l1
	} else {
		l = l2
	}

	for i = 0; i < l && s1[i] == s2[i]; i++ {
	}

	return i
}

type Trie struct {
	value    interface{}
	children map[string]*Trie
}

func NewTrie() Trie {
	return Trie{children: make(map[string]*Trie)}
}

func (t *Trie) Count() int {
	r := 0
	if t.value != nil {
		r++
	}

	for _, subTrie := range t.children {
		r += subTrie.Count()
	}

	return r
}

func (t *Trie) Put(key string, value interface{}) {
	currentKey := key
	currentTrie := t

	for currentKey != "" {
		found := false
		for k, v := range currentTrie.children {
			keyLen := len(k)
			prefixLen := commonPrefixLength(currentKey, k)
			if prefixLen == keyLen {
				currentKey = currentKey[keyLen:]
				currentTrie = v
				found = true
				break
			}
			if prefixLen > 0 {
				commonPrefix := k[:prefixLen]
				oldSuffix := k[prefixLen:]
				intermediate := NewTrie()
				intermediate.children[oldSuffix] = v
				currentTrie.children[commonPrefix] = &intermediate
				delete(currentTrie.children, k)
				currentTrie = &intermediate
				currentKey = currentKey[prefixLen:]
				break
			}
		}
		if !found {
			break
		}
	}

	if currentKey == "" {
		currentTrie.value = value
	} else {
		n := NewTrie()
		n.value = value
		currentTrie.children[currentKey] = &n
	}
}

func (t *Trie) Get(key string) (interface{}, bool) {
	currentKey := key
	currentTrie := t

	for currentKey != "" {
		found := false
		for k, v := range currentTrie.children {
			keyLen := len(k)
			prefixLen := commonPrefixLength(currentKey, k)
			if prefixLen == keyLen {
				currentKey = currentKey[keyLen:]
				currentTrie = v
				found = true
				break
			}
			if prefixLen > 0 {
				return nil, false
			}
		}
		if !found {
			break
		}
	}

	if currentKey == "" {
		return currentTrie.value, currentTrie.value != nil
	}

	return nil, false
}

func (t *Trie) GetLongestPrefix(key string) (string, interface{}) {
	currentKey := key
	currentTrie := t
	matchedPrefix := ""
	longestPrefix := ""
	var matchedValue interface{}

	for currentKey != "" {
		found := false
		for k, v := range currentTrie.children {
			keyLen := len(k)
			prefixLen := commonPrefixLength(currentKey, k)
			if prefixLen == keyLen {
				currentKey = currentKey[keyLen:]
				currentTrie = v
				matchedPrefix += k
				if v.value != nil {
					longestPrefix = matchedPrefix
					matchedValue = v.value
				}
				found = true
				break
			}
			if prefixLen > 0 {
				return longestPrefix, matchedValue
			}
		}
		if !found {
			break
		}
	}

	return longestPrefix, matchedValue
}