tree.gno
4.15 Kb · 126 lines
1package staker
2
3import (
4 "gno.land/p/gnoswap/gnsmath/v1"
5 "gno.land/p/gnoswap/utils/v1"
6 bptree "gno.land/p/nt/bptree/v0"
7 ufmt "gno.land/p/nt/ufmt/v0"
8)
9
10// UintTree is a wrapper around a BPTree for storing non-negative int64 keys as
11// ordered strings. Depending on the caller, keys represent Unix timestamps,
12// accrual epochs, or stake-event indexes.
13//
14// Since keys are int64 values, they are converted to uint64-compatible strings.
15//
16// Methods:
17// - Get: Retrieves a value associated with a non-negative int64 key.
18// - set: Stores a value with a non-negative int64 key.
19// - Has: Checks if a non-negative int64 key exists in the tree.
20// - remove: Removes a non-negative int64 key and its associated value.
21// - Iterate: Iterates over keys and values in a range.
22// - ReverseIterate: Iterates in reverse order over keys and values in a range.
23type UintTree struct {
24 tree *bptree.BPTree // non-negative int64 key -> any
25}
26
27// NewUintTree creates a new UintTree instance.
28//
29// Returns:
30// - *UintTree: new tree backed by an empty ordered BPTree
31func NewUintTree() *UintTree {
32 return &UintTree{
33 tree: bptree.NewBPTreeN(64),
34 }
35}
36
37// Get looks up a value by its non-negative int64 key.
38//
39// Parameters:
40// - key: non-negative int64 key encoded for the ordered tree lookup; a negative key panics
41//
42// Returns:
43// - any: value stored at key, or nil when no entry exists
44// - bool: true when key has an entry; false when absent
45func (self *UintTree) Get(key int64) (any, bool) {
46 v := self.tree.Get(encodeInt64(key))
47 if v == nil {
48 return nil, false
49 }
50 return v, true
51}
52
53// Set associates a value with a non-negative int64 key, replacing any existing
54// value at that key.
55//
56// Parameters:
57// - key: non-negative int64 key to encode and store; a negative key panics
58// - value: value to associate with key
59func (self *UintTree) Set(key int64, value any) {
60 self.tree.Set(encodeInt64(key), value)
61}
62
63// Has checks whether a non-negative int64 key is present in the tree.
64//
65// Parameters:
66// - key: non-negative int64 key to encode and test; a negative key panics
67//
68// Returns:
69// - bool: true when key is present; false when no entry is stored at key
70func (self *UintTree) Has(key int64) bool {
71 return self.tree.Has(encodeInt64(key))
72}
73
74// Remove deletes the entry associated with a non-negative int64 key, if one
75// exists.
76//
77// Parameters:
78// - key: non-negative int64 key to encode and remove; a negative key panics
79func (self *UintTree) Remove(key int64) {
80 self.tree.Remove(encodeInt64(key))
81}
82
83// Iterate visits entries in the half-open [start, end) key range in ascending
84// order.
85//
86// Parameters:
87// - start: non-negative lower-bound key for the iteration range; a negative bound panics
88// - end: non-negative upper-bound key for the iteration range; a negative bound panics
89// - fn: callback receiving each decoded key and value; return true to stop iteration, or false to continue
90func (self *UintTree) Iterate(start, end int64, fn func(key int64, value any) bool) {
91 self.tree.Iterate(encodeInt64(start), encodeInt64(end), func(key string, value any) bool {
92 return fn(decodeInt64(key), value)
93 })
94}
95
96// ReverseIterate visits entries in descending order over the half-open
97// [start, end) key range.
98//
99// Parameters:
100// - start: non-negative lower-bound key for the iteration range; a negative bound panics
101// - end: non-negative upper-bound key for the iteration range; a negative bound panics
102// - fn: callback receiving each decoded key and value; return true to stop iteration, or false to continue
103func (self *UintTree) ReverseIterate(start, end int64, fn func(key int64, value any) bool) {
104 self.tree.ReverseIterate(encodeInt64(start), encodeInt64(end), func(key string, value any) bool {
105 return fn(decodeInt64(key), value)
106 })
107}
108
109// Size returns the number of entries in the tree.
110//
111// Returns:
112// - int: number of key-value entries currently stored in the tree
113func (self *UintTree) Size() int {
114 return self.tree.Size()
115}
116
117func encodeInt64(num int64) string {
118 if num < 0 {
119 panic(ufmt.Sprintf("negative value not supported: %d", num))
120 }
121 return utils.EncodeUint64(uint64(num))
122}
123
124func decodeInt64(s string) int64 {
125 return gnsmath.SafeUint64ToInt64(utils.DecodeUint64(s))
126}