feat: UnitTree, RewardCacheTree Size Support

gnoswap-labs · Dec 31, 2024 · 3f3895f · 3f3895f
1 parent 85ef73b
commit 3f3895f
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/staker/reward_calculation_types.gno b/staker/reward_calculation_types.gno
@@ -66,6 +66,10 @@ func (self *UintTree) ReverseIterate(start, end uint64, fn func(key uint64, valu
 	})
 }
 
+func (self *UintTree) Size() uint64 {
+	return uint64(self.tree.Size())
+}
+
 type RewardCacheTree struct {
 	tree *avl.Tree
 }
@@ -98,6 +102,10 @@ func (self *RewardCacheTree) Iterate(start, end uint64, fn func(key uint64, valu
 	})
 }
 
+func (self *RewardCacheTree) Size() uint64 {
+	return uint64(self.tree.Size())
+}
+
 func (self *RewardCacheTree) CurrentReward(currentHeight uint64) interface{} {
 	result, ok := self.tree.Get(EncodeUint(currentHeight))
 	if ok {
@@ -124,4 +132,4 @@ func (self *RewardCacheTree) RewardPerInterval(startHeight, endHeight uint64, f
 	if endHeight > currentHeight {
 		f(endHeight-currentHeight, currentPoolReward)
 	}
-}
+}