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  }