prometheus/vendor/github.com/cockroachdb/cmux/patricia.go
2017-07-06 14:38:20 +02:00

180 lines
3.2 KiB
Go

// Copyright 2016 The CMux Authors. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
package cmux
import (
"bytes"
"io"
)
// patriciaTree is a simple patricia tree that handles []byte instead of string
// and cannot be changed after instantiation.
type patriciaTree struct {
root *ptNode
maxDepth int // max depth of the tree.
}
func newPatriciaTree(bs ...[]byte) *patriciaTree {
max := 0
for _, b := range bs {
if max < len(b) {
max = len(b)
}
}
return &patriciaTree{
root: newNode(bs),
maxDepth: max + 1,
}
}
func newPatriciaTreeString(strs ...string) *patriciaTree {
b := make([][]byte, len(strs))
for i, s := range strs {
b[i] = []byte(s)
}
return newPatriciaTree(b...)
}
func (t *patriciaTree) matchPrefix(r io.Reader) bool {
buf := make([]byte, t.maxDepth)
n, _ := io.ReadFull(r, buf)
return t.root.match(buf[:n], true)
}
func (t *patriciaTree) match(r io.Reader) bool {
buf := make([]byte, t.maxDepth)
n, _ := io.ReadFull(r, buf)
return t.root.match(buf[:n], false)
}
type ptNode struct {
prefix []byte
next map[byte]*ptNode
terminal bool
}
func newNode(strs [][]byte) *ptNode {
if len(strs) == 0 {
return &ptNode{
prefix: []byte{},
terminal: true,
}
}
if len(strs) == 1 {
return &ptNode{
prefix: strs[0],
terminal: true,
}
}
p, strs := splitPrefix(strs)
n := &ptNode{
prefix: p,
}
nexts := make(map[byte][][]byte)
for _, s := range strs {
if len(s) == 0 {
n.terminal = true
continue
}
nexts[s[0]] = append(nexts[s[0]], s[1:])
}
n.next = make(map[byte]*ptNode)
for first, rests := range nexts {
n.next[first] = newNode(rests)
}
return n
}
func splitPrefix(bss [][]byte) (prefix []byte, rest [][]byte) {
if len(bss) == 0 || len(bss[0]) == 0 {
return prefix, bss
}
if len(bss) == 1 {
return bss[0], [][]byte{{}}
}
for i := 0; ; i++ {
var cur byte
eq := true
for j, b := range bss {
if len(b) <= i {
eq = false
break
}
if j == 0 {
cur = b[i]
continue
}
if cur != b[i] {
eq = false
break
}
}
if !eq {
break
}
prefix = append(prefix, cur)
}
rest = make([][]byte, 0, len(bss))
for _, b := range bss {
rest = append(rest, b[len(prefix):])
}
return prefix, rest
}
func (n *ptNode) match(b []byte, prefix bool) bool {
l := len(n.prefix)
if l > 0 {
if l > len(b) {
l = len(b)
}
if !bytes.Equal(b[:l], n.prefix) {
return false
}
}
if n.terminal && (prefix || len(n.prefix) == len(b)) {
return true
}
if l >= len(b) {
return false
}
nextN, ok := n.next[b[l]]
if !ok {
return false
}
if l == len(b) {
b = b[l:l]
} else {
b = b[l+1:]
}
return nextN.match(b, prefix)
}