120 lines
2.9 KiB
Go
120 lines
2.9 KiB
Go
package truncindex
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/tchap/go-patricia/patricia"
|
|
)
|
|
|
|
var (
|
|
ErrEmptyPrefix = errors.New("Prefix can't be empty")
|
|
ErrAmbiguousPrefix = errors.New("Multiple IDs found with provided prefix")
|
|
)
|
|
|
|
// TruncIndex allows the retrieval of string identifiers by any of their unique prefixes.
|
|
// This is used to retrieve image and container IDs by more convenient shorthand prefixes.
|
|
type TruncIndex struct {
|
|
sync.RWMutex
|
|
trie *patricia.Trie
|
|
ids map[string]struct{}
|
|
}
|
|
|
|
// NewTruncIndex creates a new TruncIndex and initializes with a list of IDs
|
|
func NewTruncIndex(ids []string) (idx *TruncIndex) {
|
|
idx = &TruncIndex{
|
|
ids: make(map[string]struct{}),
|
|
|
|
// Change patricia max prefix per node length,
|
|
// because our len(ID) always 64
|
|
trie: patricia.NewTrie(patricia.MaxPrefixPerNode(64)),
|
|
}
|
|
for _, id := range ids {
|
|
idx.addID(id)
|
|
}
|
|
return
|
|
}
|
|
|
|
func (idx *TruncIndex) addID(id string) error {
|
|
if strings.Contains(id, " ") {
|
|
return fmt.Errorf("illegal character: ' '")
|
|
}
|
|
if id == "" {
|
|
return ErrEmptyPrefix
|
|
}
|
|
if _, exists := idx.ids[id]; exists {
|
|
return fmt.Errorf("id already exists: '%s'", id)
|
|
}
|
|
idx.ids[id] = struct{}{}
|
|
if inserted := idx.trie.Insert(patricia.Prefix(id), struct{}{}); !inserted {
|
|
return fmt.Errorf("failed to insert id: %s", id)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Add adds a new ID to the TruncIndex
|
|
func (idx *TruncIndex) Add(id string) error {
|
|
idx.Lock()
|
|
defer idx.Unlock()
|
|
if err := idx.addID(id); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Delete removes an ID from the TruncIndex. If there are multiple IDs
|
|
// with the given prefix, an error is thrown.
|
|
func (idx *TruncIndex) Delete(id string) error {
|
|
idx.Lock()
|
|
defer idx.Unlock()
|
|
if _, exists := idx.ids[id]; !exists || id == "" {
|
|
return fmt.Errorf("no such id: '%s'", id)
|
|
}
|
|
delete(idx.ids, id)
|
|
if deleted := idx.trie.Delete(patricia.Prefix(id)); !deleted {
|
|
return fmt.Errorf("no such id: '%s'", id)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Get retrieves an ID from the TruncIndex. If there are multiple IDs
|
|
// with the given prefix, an error is thrown.
|
|
func (idx *TruncIndex) Get(s string) (string, error) {
|
|
if s == "" {
|
|
return "", ErrEmptyPrefix
|
|
}
|
|
var (
|
|
id string
|
|
)
|
|
subTreeVisitFunc := func(prefix patricia.Prefix, item patricia.Item) error {
|
|
if id != "" {
|
|
// we haven't found the ID if there are two or more IDs
|
|
id = ""
|
|
return ErrAmbiguousPrefix
|
|
}
|
|
id = string(prefix)
|
|
return nil
|
|
}
|
|
|
|
idx.RLock()
|
|
defer idx.RUnlock()
|
|
if err := idx.trie.VisitSubtree(patricia.Prefix(s), subTreeVisitFunc); err != nil {
|
|
return "", err
|
|
}
|
|
if id != "" {
|
|
return id, nil
|
|
}
|
|
return "", fmt.Errorf("no such id: %s", s)
|
|
}
|
|
|
|
// Iterates over all stored IDs, and passes each of them to the given handler
|
|
func (idx *TruncIndex) Iterate(handler func(id string)) {
|
|
idx.RLock()
|
|
defer idx.RUnlock()
|
|
idx.trie.Visit(func(prefix patricia.Prefix, item patricia.Item) error {
|
|
handler(string(prefix))
|
|
return nil
|
|
})
|
|
}
|