github.com/rosedblabs/rosedb/v2@v2.3.7-0.20240423093736-a89ea823e5b9/index/btree.go (about)

     1  package index
     2  
     3  import (
     4  	"bytes"
     5  	"sync"
     6  
     7  	"github.com/google/btree"
     8  	"github.com/rosedblabs/wal"
     9  )
    10  
    11  // MemoryBTree is a memory based btree implementation of the Index interface
    12  // It is a wrapper around the google/btree package: github.com/google/btree
    13  type MemoryBTree struct {
    14  	tree *btree.BTree
    15  	lock *sync.RWMutex
    16  }
    17  
    18  type item struct {
    19  	key []byte
    20  	pos *wal.ChunkPosition
    21  }
    22  
    23  func newBTree() *MemoryBTree {
    24  	return &MemoryBTree{
    25  		tree: btree.New(32),
    26  		lock: new(sync.RWMutex),
    27  	}
    28  }
    29  
    30  func (it *item) Less(bi btree.Item) bool {
    31  	if bi == nil {
    32  		return false
    33  	}
    34  	return bytes.Compare(it.key, bi.(*item).key) < 0
    35  }
    36  
    37  func (mt *MemoryBTree) Put(key []byte, position *wal.ChunkPosition) *wal.ChunkPosition {
    38  	mt.lock.Lock()
    39  	defer mt.lock.Unlock()
    40  
    41  	oldValue := mt.tree.ReplaceOrInsert(&item{key: key, pos: position})
    42  	if oldValue != nil {
    43  		return oldValue.(*item).pos
    44  	}
    45  	return nil
    46  }
    47  
    48  func (mt *MemoryBTree) Get(key []byte) *wal.ChunkPosition {
    49  	mt.lock.RLock()
    50  	defer mt.lock.RUnlock()
    51  	value := mt.tree.Get(&item{key: key})
    52  	if value != nil {
    53  		return value.(*item).pos
    54  	}
    55  	return nil
    56  }
    57  
    58  func (mt *MemoryBTree) Delete(key []byte) (*wal.ChunkPosition, bool) {
    59  	mt.lock.Lock()
    60  	defer mt.lock.Unlock()
    61  
    62  	value := mt.tree.Delete(&item{key: key})
    63  	if value != nil {
    64  		return value.(*item).pos, true
    65  	}
    66  	return nil, false
    67  }
    68  
    69  func (mt *MemoryBTree) Size() int {
    70  	return mt.tree.Len()
    71  }
    72  
    73  func (mt *MemoryBTree) Ascend(handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
    74  	mt.lock.RLock()
    75  	defer mt.lock.RUnlock()
    76  
    77  	mt.tree.Ascend(func(i btree.Item) bool {
    78  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
    79  		if err != nil {
    80  			return false
    81  		}
    82  		return cont
    83  	})
    84  }
    85  
    86  func (mt *MemoryBTree) Descend(handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
    87  	mt.lock.RLock()
    88  	defer mt.lock.RUnlock()
    89  
    90  	mt.tree.Descend(func(i btree.Item) bool {
    91  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
    92  		if err != nil {
    93  			return false
    94  		}
    95  		return cont
    96  	})
    97  }
    98  
    99  func (mt *MemoryBTree) AscendRange(startKey, endKey []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
   100  	mt.lock.RLock()
   101  	defer mt.lock.RUnlock()
   102  
   103  	mt.tree.AscendRange(&item{key: startKey}, &item{key: endKey}, func(i btree.Item) bool {
   104  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
   105  		if err != nil {
   106  			return false
   107  		}
   108  		return cont
   109  	})
   110  }
   111  
   112  func (mt *MemoryBTree) DescendRange(startKey, endKey []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
   113  	mt.lock.RLock()
   114  	defer mt.lock.RUnlock()
   115  
   116  	mt.tree.DescendRange(&item{key: startKey}, &item{key: endKey}, func(i btree.Item) bool {
   117  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
   118  		if err != nil {
   119  			return false
   120  		}
   121  		return cont
   122  	})
   123  }
   124  
   125  func (mt *MemoryBTree) AscendGreaterOrEqual(key []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
   126  	mt.lock.RLock()
   127  	defer mt.lock.RUnlock()
   128  
   129  	mt.tree.AscendGreaterOrEqual(&item{key: key}, func(i btree.Item) bool {
   130  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
   131  		if err != nil {
   132  			return false
   133  		}
   134  		return cont
   135  	})
   136  }
   137  
   138  func (mt *MemoryBTree) DescendLessOrEqual(key []byte, handleFn func(key []byte, position *wal.ChunkPosition) (bool, error)) {
   139  	mt.lock.RLock()
   140  	defer mt.lock.RUnlock()
   141  
   142  	mt.tree.DescendLessOrEqual(&item{key: key}, func(i btree.Item) bool {
   143  		cont, err := handleFn(i.(*item).key, i.(*item).pos)
   144  		if err != nil {
   145  			return false
   146  		}
   147  		return cont
   148  	})
   149  }