github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/pingcap/tidb/kv/rbtree_buffer.go (about)

     1  package kv
     2  
     3  import (
     4  	"bytes"
     5  
     6  	"github.com/insionng/yougam/libraries/juju/errors"
     7  	"github.com/insionng/yougam/libraries/petar/GoLLRB/llrb"
     8  )
     9  
    10  type keyItem Key
    11  
    12  func (k keyItem) Less(item llrb.Item) bool {
    13  	switch x := item.(type) {
    14  	case keyItem:
    15  		return bytes.Compare(k, x) < 0
    16  	case *pairItem:
    17  		return bytes.Compare(k, x.key) < 0
    18  	}
    19  	return true
    20  }
    21  
    22  type pairItem struct {
    23  	key   Key
    24  	value []byte
    25  }
    26  
    27  func (pair *pairItem) Less(item llrb.Item) bool {
    28  	switch x := item.(type) {
    29  	case keyItem:
    30  		return bytes.Compare(pair.key, x) < 0
    31  	case *pairItem:
    32  		return bytes.Compare(pair.key, x.key) < 0
    33  	}
    34  	return true
    35  }
    36  
    37  type rbTreeBuffer struct {
    38  	tree *llrb.LLRB
    39  }
    40  
    41  type rbTreeIter struct {
    42  	tree *llrb.LLRB
    43  	seek Key
    44  	pair *pairItem
    45  }
    46  
    47  // NewRBTreeBuffer creates a new rbTreeBuffer.
    48  func NewRBTreeBuffer() MemBuffer {
    49  	return &rbTreeBuffer{tree: llrb.New()}
    50  }
    51  
    52  // Seek creates an Iterator.
    53  func (m *rbTreeBuffer) Seek(k Key) (Iterator, error) {
    54  	it := &rbTreeIter{tree: m.tree, seek: k}
    55  	it.Next()
    56  	return it, nil
    57  }
    58  
    59  func (m *rbTreeBuffer) SeekReverse(k Key) (Iterator, error) {
    60  	// TODO: implement Prev.
    61  	return nil, ErrNotImplemented
    62  }
    63  
    64  // Get returns the value associated with key.
    65  func (m *rbTreeBuffer) Get(k Key) ([]byte, error) {
    66  	pair := m.tree.Get(keyItem(k))
    67  	if pair == nil {
    68  		return nil, ErrNotExist
    69  	}
    70  	return pair.(*pairItem).value, nil
    71  }
    72  
    73  // Set associates key with value.
    74  func (m *rbTreeBuffer) Set(k Key, v []byte) error {
    75  	if len(v) == 0 {
    76  		return errors.Trace(ErrCannotSetNilValue)
    77  	}
    78  	m.tree.ReplaceOrInsert(&pairItem{key: k, value: v})
    79  	return nil
    80  }
    81  
    82  // Delete removes the entry from buffer with provided key.
    83  func (m *rbTreeBuffer) Delete(k Key) error {
    84  	m.tree.ReplaceOrInsert(&pairItem{key: k, value: nil})
    85  	return nil
    86  }
    87  
    88  // Release reset the buffer.
    89  func (m *rbTreeBuffer) Release() {
    90  	m.tree = llrb.New()
    91  }
    92  
    93  // Next implements the Iterator Next.
    94  func (i *rbTreeIter) Next() error {
    95  	i.pair = nil
    96  	i.tree.AscendGreaterOrEqual(keyItem(i.seek), func(item llrb.Item) bool {
    97  		i.pair = item.(*pairItem)
    98  		i.seek = i.pair.key.Next()
    99  		return false
   100  	})
   101  	return nil
   102  }
   103  
   104  // Valid implements the Iterator Valid.
   105  func (i *rbTreeIter) Valid() bool {
   106  	return i.pair != nil
   107  }
   108  
   109  // Key implements the Iterator Key.
   110  func (i *rbTreeIter) Key() Key {
   111  	return i.pair.key
   112  }
   113  
   114  // Value implements the Iterator Value.
   115  func (i *rbTreeIter) Value() []byte {
   116  	return i.pair.value
   117  }
   118  
   119  // Close Implements the Iterator Close.
   120  func (i *rbTreeIter) Close() {
   121  }