cmd/bpa-operator/vendor/github.com/peterbourgon/diskv/index.go

   1 package diskv
   2
   3 import (
   4         "sync"
   5
   6         "github.com/google/btree"
   7 )
   8
   9 // Index is a generic interface for things that can
  10 // provide an ordered list of keys.
  11 type Index interface {
  12         Initialize(less LessFunction, keys <-chan string)
  13         Insert(key string)
  14         Delete(key string)
  15         Keys(from string, n int) []string
  16 }
  17
  18 // LessFunction is used to initialize an Index of keys in a specific order.
  19 type LessFunction func(string, string) bool
  20
  21 // btreeString is a custom data type that satisfies the BTree Less interface,
  22 // making the strings it wraps sortable by the BTree package.
  23 type btreeString struct {
  24         s string
  25         l LessFunction
  26 }
  27
  28 // Less satisfies the BTree.Less interface using the btreeString's LessFunction.
  29 func (s btreeString) Less(i btree.Item) bool {
  30         return s.l(s.s, i.(btreeString).s)
  31 }
  32
  33 // BTreeIndex is an implementation of the Index interface using google/btree.
  34 type BTreeIndex struct {
  35         sync.RWMutex
  36         LessFunction
  37         *btree.BTree
  38 }
  39
  40 // Initialize populates the BTree tree with data from the keys channel,
  41 // according to the passed less function. It's destructive to the BTreeIndex.
  42 func (i *BTreeIndex) Initialize(less LessFunction, keys <-chan string) {
  43         i.Lock()
  44         defer i.Unlock()
  45         i.LessFunction = less
  46         i.BTree = rebuild(less, keys)
  47 }
  48
  49 // Insert inserts the given key (only) into the BTree tree.
  50 func (i *BTreeIndex) Insert(key string) {
  51         i.Lock()
  52         defer i.Unlock()
  53         if i.BTree == nil || i.LessFunction == nil {
  54                 panic("uninitialized index")
  55         }
  56         i.BTree.ReplaceOrInsert(btreeString{s: key, l: i.LessFunction})
  57 }
  58
  59 // Delete removes the given key (only) from the BTree tree.
  60 func (i *BTreeIndex) Delete(key string) {
  61         i.Lock()
  62         defer i.Unlock()
  63         if i.BTree == nil || i.LessFunction == nil {
  64                 panic("uninitialized index")
  65         }
  66         i.BTree.Delete(btreeString{s: key, l: i.LessFunction})
  67 }
  68
  69 // Keys yields a maximum of n keys in order. If the passed 'from' key is empty,
  70 // Keys will return the first n keys. If the passed 'from' key is non-empty, the
  71 // first key in the returned slice will be the key that immediately follows the
  72 // passed key, in key order.
  73 func (i *BTreeIndex) Keys(from string, n int) []string {
  74         i.RLock()
  75         defer i.RUnlock()
  76
  77         if i.BTree == nil || i.LessFunction == nil {
  78                 panic("uninitialized index")
  79         }
  80
  81         if i.BTree.Len() <= 0 {
  82                 return []string{}
  83         }
  84
  85         btreeFrom := btreeString{s: from, l: i.LessFunction}
  86         skipFirst := true
  87         if len(from) <= 0 || !i.BTree.Has(btreeFrom) {
  88                 // no such key, so fabricate an always-smallest item
  89                 btreeFrom = btreeString{s: "", l: func(string, string) bool { return true }}
  90                 skipFirst = false
  91         }
  92
  93         keys := []string{}
  94         iterator := func(i btree.Item) bool {
  95                 keys = append(keys, i.(btreeString).s)
  96                 return len(keys) < n
  97         }
  98         i.BTree.AscendGreaterOrEqual(btreeFrom, iterator)
  99
 100         if skipFirst && len(keys) > 0 {
 101                 keys = keys[1:]
 102         }
 103
 104         return keys
 105 }
 106
 107 // rebuildIndex does the work of regenerating the index
 108 // with the given keys.
 109 func rebuild(less LessFunction, keys <-chan string) *btree.BTree {
 110         tree := btree.New(2)
 111         for key := range keys {
 112                 tree.ReplaceOrInsert(btreeString{s: key, l: less})
 113         }
 114         return tree
 115 }