github.com/nutsdb/nutsdb@v1.0.4/btree.go (about) 1 // Copyright 2023 The nutsdb Author. All rights reserved. 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 nutsdb 16 17 import ( 18 "bytes" 19 "errors" 20 "regexp" 21 22 "github.com/tidwall/btree" 23 ) 24 25 // ErrKeyNotFound is returned when the key is not in the b tree. 26 var ErrKeyNotFound = errors.New("key not found") 27 28 type Item struct { 29 key []byte 30 record *Record 31 } 32 33 type BTree struct { 34 btree *btree.BTreeG[*Item] 35 } 36 37 func NewBTree() *BTree { 38 return &BTree{ 39 btree: btree.NewBTreeG[*Item](func(a, b *Item) bool { 40 return bytes.Compare(a.key, b.key) == -1 41 }), 42 } 43 } 44 45 func (bt *BTree) Find(key []byte) (*Record, bool) { 46 item, ok := bt.btree.Get(&Item{key: key}) 47 if ok { 48 return item.record, ok 49 } 50 return nil, ok 51 } 52 53 func (bt *BTree) Insert(record *Record) bool { 54 _, replaced := bt.btree.Set(&Item{key: record.Key, record: record}) 55 return replaced 56 } 57 58 func (bt *BTree) InsertRecord(key []byte, record *Record) bool { 59 _, replaced := bt.btree.Set(&Item{key: key, record: record}) 60 return replaced 61 } 62 63 func (bt *BTree) Delete(key []byte) bool { 64 _, deleted := bt.btree.Delete(&Item{key: key}) 65 return deleted 66 } 67 68 func (bt *BTree) All() []*Record { 69 items := bt.btree.Items() 70 71 records := make([]*Record, len(items)) 72 for i, item := range items { 73 records[i] = item.record 74 } 75 76 return records 77 } 78 79 func (bt *BTree) AllItems() []*Item { 80 items := bt.btree.Items() 81 return items 82 } 83 84 func (bt *BTree) Range(start, end []byte) []*Record { 85 records := make([]*Record, 0) 86 87 bt.btree.Ascend(&Item{key: start}, func(item *Item) bool { 88 if bytes.Compare(item.key, end) > 0 { 89 return false 90 } 91 records = append(records, item.record) 92 return true 93 }) 94 95 return records 96 } 97 98 func (bt *BTree) PrefixScan(prefix []byte, offset, limitNum int) []*Record { 99 records := make([]*Record, 0) 100 101 bt.btree.Ascend(&Item{key: prefix}, func(item *Item) bool { 102 if !bytes.HasPrefix(item.key, prefix) { 103 return false 104 } 105 106 if offset > 0 { 107 offset-- 108 return true 109 } 110 111 records = append(records, item.record) 112 113 limitNum-- 114 return limitNum != 0 115 }) 116 117 return records 118 } 119 120 func (bt *BTree) PrefixSearchScan(prefix []byte, reg string, offset, limitNum int) []*Record { 121 records := make([]*Record, 0) 122 123 rgx := regexp.MustCompile(reg) 124 125 bt.btree.Ascend(&Item{key: prefix}, func(item *Item) bool { 126 if !bytes.HasPrefix(item.key, prefix) { 127 return false 128 } 129 130 if offset > 0 { 131 offset-- 132 return true 133 } 134 135 if !rgx.Match(bytes.TrimPrefix(item.key, prefix)) { 136 return true 137 } 138 139 records = append(records, item.record) 140 141 limitNum-- 142 return limitNum != 0 143 }) 144 145 return records 146 } 147 148 func (bt *BTree) Count() int { 149 return bt.btree.Len() 150 } 151 152 func (bt *BTree) PopMin() (*Item, bool) { 153 return bt.btree.PopMin() 154 } 155 156 func (bt *BTree) PopMax() (*Item, bool) { 157 return bt.btree.PopMax() 158 } 159 160 func (bt *BTree) Min() (*Item, bool) { 161 return bt.btree.Min() 162 } 163 164 func (bt *BTree) Max() (*Item, bool) { 165 return bt.btree.Max() 166 }