github.com/matrixorigin/matrixone@v0.7.0/pkg/txn/storage/mem/kv_tree.go (about)

     1  // Copyright 2021 - 2022 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package mem
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"github.com/google/btree"
    21  )
    22  
    23  type treeItem struct {
    24  	key   []byte
    25  	value []byte
    26  }
    27  
    28  func (item treeItem) Less(other btree.Item) bool {
    29  	left := item.key
    30  	right := other.(treeItem).key
    31  	return bytes.Compare(right, left) > 0
    32  }
    33  
    34  // KV mem kv based on btree. Just used to test transaction.
    35  type KV struct {
    36  	tree *btree.BTree
    37  	tmp  treeItem
    38  }
    39  
    40  // NewKV returns a Mem KV
    41  func NewKV() *KV {
    42  	return &KV{
    43  		tree: btree.New(32),
    44  	}
    45  }
    46  
    47  // Set set key value to KV
    48  func (kv *KV) Set(key, value []byte) {
    49  	kv.tree.ReplaceOrInsert(treeItem{
    50  		key:   key,
    51  		value: value,
    52  	})
    53  }
    54  
    55  // Delete delete key from KV
    56  func (kv *KV) Delete(key []byte) bool {
    57  	item := treeItem{key: key}
    58  	return nil != kv.tree.Delete(item)
    59  }
    60  
    61  // Get returns the value of the key
    62  func (kv *KV) Get(key []byte) ([]byte, bool) {
    63  	kv.tmp.key = key
    64  	var result treeItem
    65  	kv.tree.AscendGreaterOrEqual(kv.tmp, func(i btree.Item) bool {
    66  		result = i.(treeItem)
    67  		return false
    68  	})
    69  	if bytes.Equal(result.key, key) {
    70  		return result.value, true
    71  	}
    72  	return nil, false
    73  }
    74  
    75  // Len return the count of keys
    76  func (kv *KV) Len() int {
    77  	return kv.tree.Len()
    78  }
    79  
    80  // AscendRange iter in [start, end)
    81  func (kv *KV) AscendRange(start, end []byte, fn func(key, value []byte) bool) {
    82  	kv.tmp.key = start
    83  	kv.tree.AscendGreaterOrEqual(kv.tmp, func(i btree.Item) bool {
    84  		target := i.(treeItem)
    85  		if bytes.Compare(target.key, end) < 0 {
    86  			return fn(target.key, target.value)
    87  		}
    88  		return false
    89  	})
    90  }