github.com/peterbourgon/diskv@v2.0.1+incompatible/index.go (about)

     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  }