github.com/cilium/statedb@v0.3.2/part/tree.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright Authors of Cilium
     3  
     4  package part
     5  
     6  // Tree is a persistent (immutable) adaptive radix tree. It supports
     7  // map-like operations on values keyed by []byte and additionally
     8  // prefix searching and lower bound searching. Each node in the tree
     9  // has an associated channel that is closed when that node is mutated.
    10  // This allows watching any part of the tree (any prefix) for changes.
    11  type Tree[T any] struct {
    12  	opts *options
    13  	root *header[T]
    14  	size int // the number of objects in the tree
    15  }
    16  
    17  // New constructs a new tree.
    18  func New[T any](opts ...Option) *Tree[T] {
    19  	var o options
    20  	for _, opt := range opts {
    21  		opt(&o)
    22  	}
    23  	return &Tree[T]{
    24  		root: newNode4[T](),
    25  		size: 0,
    26  		opts: &o,
    27  	}
    28  }
    29  
    30  type Option func(*options)
    31  
    32  // RootOnlyWatch sets the tree to only have a watch channel on the root
    33  // node. This improves the speed at the cost of having a much more coarse
    34  // grained notifications.
    35  func RootOnlyWatch(o *options) { o.rootOnlyWatch = true }
    36  
    37  // Txn constructs a new transaction against the tree. Transactions
    38  // enable efficient large mutations of the tree by caching cloned
    39  // nodes.
    40  func (t *Tree[T]) Txn() *Txn[T] {
    41  	txn := &Txn[T]{
    42  		Tree:    *t,
    43  		watches: make(map[chan struct{}]struct{}),
    44  	}
    45  	return txn
    46  }
    47  
    48  // Len returns the number of objects in the tree.
    49  func (t *Tree[T]) Len() int {
    50  	return t.size
    51  }
    52  
    53  // Get fetches the value associated with the given key.
    54  // Returns the value, a watch channel (which is closed on
    55  // modification to the key) and boolean which is true if
    56  // value was found.
    57  func (t *Tree[T]) Get(key []byte) (T, <-chan struct{}, bool) {
    58  	value, watch, ok := search(t.root, key)
    59  	if t.opts.rootOnlyWatch {
    60  		watch = t.root.watch
    61  	}
    62  	return value, watch, ok
    63  }
    64  
    65  // Prefix returns an iterator for all objects that starts with the
    66  // given prefix, and a channel that closes when any objects matching
    67  // the given prefix are upserted or deleted.
    68  func (t *Tree[T]) Prefix(prefix []byte) (*Iterator[T], <-chan struct{}) {
    69  	iter, watch := prefixSearch(t.root, prefix)
    70  	if t.opts.rootOnlyWatch {
    71  		watch = t.root.watch
    72  	}
    73  	return iter, watch
    74  }
    75  
    76  // RootWatch returns a watch channel for the root of the tree.
    77  // Since this is the channel associated with the root, this closes
    78  // when there are any changes to the tree.
    79  func (t *Tree[T]) RootWatch() <-chan struct{} {
    80  	return t.root.watch
    81  }
    82  
    83  // LowerBound returns an iterator for all keys that have a value
    84  // equal to or higher than 'key'.
    85  func (t *Tree[T]) LowerBound(key []byte) *Iterator[T] {
    86  	return lowerbound(t.root, key)
    87  }
    88  
    89  // Insert inserts the key into the tree with the given value.
    90  // Returns the old value if it exists and a new tree.
    91  func (t *Tree[T]) Insert(key []byte, value T) (old T, hadOld bool, tree *Tree[T]) {
    92  	txn := t.Txn()
    93  	old, hadOld = txn.Insert(key, value)
    94  	tree = txn.Commit()
    95  	return
    96  }
    97  
    98  // Modify a value in the tree. If the key does not exist the modify
    99  // function is called with the zero value for T. It is up to the
   100  // caller to not mutate the value in-place and to return a clone.
   101  // Returns the old value if it exists.
   102  func (t *Tree[T]) Modify(key []byte, mod func(T) T) (old T, hadOld bool, tree *Tree[T]) {
   103  	txn := t.Txn()
   104  	old, hadOld = txn.Modify(key, mod)
   105  	tree = txn.Commit()
   106  	return
   107  }
   108  
   109  // Delete the given key from the tree.
   110  // Returns the old value if it exists and the new tree.
   111  func (t *Tree[T]) Delete(key []byte) (old T, hadOld bool, tree *Tree[T]) {
   112  	txn := t.Txn()
   113  	old, hadOld = txn.Delete(key)
   114  	tree = txn.Commit()
   115  	return
   116  }
   117  
   118  // Iterator returns an iterator for all objects.
   119  func (t *Tree[T]) Iterator() *Iterator[T] {
   120  	return newIterator[T](t.root)
   121  }
   122  
   123  // PrintTree to the standard output. For debugging.
   124  func (t *Tree[T]) PrintTree() {
   125  	t.root.printTree(0)
   126  }
   127  
   128  type options struct {
   129  	rootOnlyWatch bool
   130  }