refactor: delegate series

Author: notJoon

contract/r/gnoswap/gov/staker/api_delegation.gno

@@ -21,7 +21,7 @@ func GetTotalVoteWeight() uint64 {
 
 // GetTotalDelegated returns the total amount of xGNS delegated.
 func GetTotalDelegated() uint64 {
-	return totalDelegated
+	return state.TotalDelegated()
 }
 
 // GetTotalLockedAmount returns the total amount of locked GNS.
@@ -31,7 +31,7 @@ func GetTotalLockedAmount() uint64 {
 
 // GetTotalDelegatedFrom returns the total amount of xGNS delegated by given address.
 func GetTotalDelegatedFrom(from std.Address) uint64 {
-	amount, exist := delegatorAmount.Get(from.String())
+	amount, exist := state.DelegatorAmount().Get(from.String())
 	if !exist {
 		return 0
 	}
@@ -40,7 +40,7 @@ func GetTotalDelegatedFrom(from std.Address) uint64 {
 
 // GetTotalDelegatedTo returns the total amount of xGNS delegated to given address.
 func GetTotalDelegatedTo(to std.Address) uint64 {
-	amount, exist := delegatedTo.Get(to.String())
+	amount, exist := state.DelegatedTo().Get(to.String())
 	if !exist {
 		return 0
 	}
@@ -49,7 +49,7 @@ func GetTotalDelegatedTo(to std.Address) uint64 {
 
 // GetDelegationAmountFromTo returns the amount of xGNS delegated by given address to given address.
 func GetDelegationAmountFromTo(from, to std.Address) uint64 {
-	toAmount, exist := delegatedFromTo.Get(from.String())
+	toAmount, exist := state.DelegatedFromTo().Get(from.String())
 	if !exist {
 		return 0
 	}

contract/r/gnoswap/gov/staker/delegate_undelegate.gno

@@ -8,83 +8,131 @@ import (
 	"gno.land/p/demo/ufmt"
 )
 
-var (
-	totalDelegated = uint64(0)
+// stakerState encapsulates all staking-related data.
+type stakerState struct {
+	totalDelegated            uint64
+	delegatorAmount           *avl.Tree // caller -> amount.
+	delegatedFromTo           *avl.Tree // caller -> (inner tree: to -> amount).
+	delegatedTo               *avl.Tree // to -> amount.
+	delegationHistory         *avl.Tree // caller -> []DelegationHistory.
+	delegationSnapShotHistory *avl.Tree // to -> []DelegationSnapShotHistory.
+}
 
-	delegatorAmount = avl.NewTree() // caller => amount
-	delegatedFromTo = avl.NewTree() // caller => to => amount
-	delegatedTo     = avl.NewTree() // to => amount
-)
- 
-func delegate(to std.Address, amount uint64) {
-	caller := std.PrevRealm().Addr().String()
-	toStr := to.String()
+func (s stakerState) TotalDelegated() uint64               { return s.totalDelegated }
+func (s stakerState) DelegatorAmount() *avl.Tree           { return s.delegatorAmount }
+func (s stakerState) DelegatedFromTo() *avl.Tree           { return s.delegatedFromTo }
+func (s stakerState) DelegatedTo() *avl.Tree               { return s.delegatedTo }
+func (s stakerState) DelegationHistory() *avl.Tree         { return s.delegationHistory }
+func (s stakerState) DelegationSnapShotHistory() *avl.Tree { return s.delegationSnapShotHistory }
+
+// Global state instance.
+var state = newStakerState()
+
+// newStakerState creates and returns a new StakerState.
+func newStakerState() *stakerState {
+	return &stakerState{
+		totalDelegated:            0,
+		delegatorAmount:           avl.NewTree(),
+		delegatedFromTo:           avl.NewTree(),
+		delegatedTo:               avl.NewTree(),
+		delegationHistory:         avl.NewTree(),
+		delegationSnapShotHistory: avl.NewTree(),
+	}
+}
 
-	// initialize the internal tree for callers to `delegatedFromTo`
-	innerTree := getOrCreateInnerTree(delegatedFromTo, caller)
+// getOrCreateInnerTree returns the inner tree for the given caller,
+// creating it if it does not exist.
+func (s *stakerState) getOrCreateInnerTree(caller string) *avl.Tree {
+	if inner, exists := s.delegatedFromTo.Get(caller); exists {
+		return inner.(*avl.Tree)
+	}
+	innerTree := avl.NewTree()
+	s.delegatedFromTo.Set(caller, innerTree)
+	return innerTree
+}
+
+// updateUint64 updates the value at key in the given tree by delta.
+// If add is true, it increments the value; otherwise, it decrements.
+func (s *stakerState) updateUint64(tree *avl.Tree, key string, delta uint64, add bool) {
+	var current uint64 = 0
+	if value, exists := tree.Get(key); exists {
+		current = value.(uint64)
+	}
+	if add {
+		current += delta
+	} else {
+		if current < delta {
+			panic(ufmt.Sprintf("not enough value for key %s", key))
+		}
+		current -= delta
+	}
+	tree.Set(key, current)
+}
 
-	totalDelegated += amount
+// Delegate processes a delegation from the caller to the given address.
+func (s *stakerState) Delegate(to std.Address, amount uint64) {
+	caller := std.PrevRealm().Addr().String()
+	toStr := to.String()
 
-	// update delegator amount
-	updateUint64InTree(delegatorAmount, caller, amount, true)
+	innerTree := s.getOrCreateInnerTree(caller)
 
-	// update delegatedFromTo's inner tree
-	updateUint64InTree(innerTree, toStr, amount, true)
+	s.totalDelegated += amount
 
-	// update delegatedTo
-	updateUint64InTree(delegatedTo, toStr, amount, true)
+	s.updateUint64(s.delegatorAmount, caller, amount, true)
+	s.updateUint64(innerTree, toStr, amount, true)
+	s.updateUint64(s.delegatedTo, toStr, amount, true)
 
 	timeStamp := uint64(time.Now().Unix())
-	// update delegation history
-	delegation := DelegationHistory{
-		to:        to,
-		amount:    amount,
-		timestamp: timeStamp,
-		height:    uint64(std.GetHeight()),
-		add:       true, // if true, delegation
-	}
 
-	history := make([]DelegationHistory, 0)
-	if value, exists := delegationHistory.Get(caller); exists {
+	// Update delegation history.
+	history := []DelegationHistory{}
+	if value, exists := s.delegationHistory.Get(caller); exists {
 		history = value.([]DelegationHistory)
 	}
+	delegation := newDelegationHistory(to, amount, timeStamp, uint64(std.GetHeight()), true)
 	history = append(history, delegation)
-	delegationHistory.Set(caller, history)
+	s.delegationHistory.Set(caller, history)
 
-	// update delegation stat history
-	updateAmount := uint64(0)
-	snapShotHistory := make([]DelegationSnapShotHistory, 0)
-	if value, exists := delegationSnapShotHistory.Get(toStr); exists {
+	// Update snapshot history.
+	updateAmount := amount
+	snapShotHistory := []DelegationSnapShotHistory{}
+	if value, exists := s.DelegationSnapShotHistory().Get(toStr); exists {
 		snapShotHistory = value.([]DelegationSnapShotHistory)
 		lastStat := snapShotHistory[len(snapShotHistory)-1]
 		updateAmount = lastStat.amount + amount
-	} else {
-		updateAmount = amount
 	}
+	snapShot := newDelegationSnapShotHistory(to, updateAmount, uint64(std.GetHeight()), timeStamp)
+	snapShotHistory = append(snapShotHistory, snapShot)
+	s.delegationSnapShotHistory.Set(toStr, snapShotHistory)
+}
 
-	snapShotHistory = append(snapShotHistory, DelegationSnapShotHistory{
-		to:           to,
-		amount:       updateAmount,
-		updatedBlock: uint64(std.GetHeight()),
-		updatedAt:    timeStamp,
-	})
-	delegationSnapShotHistory.Set(toStr, snapShotHistory)
+// getUint64OrPanic retrieves a uint64 value from the tree by key,
+// panicking with an error message if the key does not exist.
+func (s *stakerState) getUint64OrPanic(tree *avl.Tree, key, errMsg string) uint64 {
+	if value, exists := tree.Get(key); exists {
+		return value.(uint64)
+	}
+	panic(addDetailToError(
+		errDataNotFound,
+		ufmt.Sprintf("%s: %s", key, errMsg),
+	))
 }
 
-func undelegate(to std.Address, amount uint64) {
+// Undelegate processes an undelegation from the caller for the given address.
+func (s *stakerState) Undelegate(to std.Address, amount uint64) {
 	caller := std.PrevRealm().Addr().String()
 	toStr := to.String()
 
-	// check caller's delegatedFromTo
-	innerTree, exists := delegatedFromTo.Get(caller)
+	innerTreeValue, exists := s.DelegatedFromTo().Get(caller)
 	if !exists {
 		panic(addDetailToError(
 			errNoDelegatedAmount,
 			ufmt.Sprintf("caller(%s) has no delegated amount", caller),
 		))
 	}
-	// check caller's delegatedFromTo's inner tree
-	delegatedAmountValue, exists := innerTree.(*avl.Tree).Get(toStr)
+	innerTree := innerTreeValue.(*avl.Tree)
+
+	delegatedAmountValue, exists := innerTree.Get(toStr)
 	if !exists {
 		panic(addDetailToError(
 			errNoDelegatedTarget,
@@ -95,86 +143,74 @@ func undelegate(to std.Address, amount uint64) {
 	if delegatedAmount < amount {
 		panic(addDetailToError(
 			errNotEnoughDelegated,
-			ufmt.Sprintf("caller(%s) has only %d delegated amount(request: %d) to %s", caller, delegatedAmount, amount, to),
+			ufmt.Sprintf("caller(%s) has only %d delegated amount (requested: %d) to %s", caller, delegatedAmount, amount, to),
 		))
 	}
-	innerTree.(*avl.Tree).Set(toStr, delegatedAmount-amount)
-	delegatedFromTo.Set(caller, innerTree)
+	innerTree.Set(toStr, delegatedAmount-amount)
+	s.DelegatedFromTo().Set(caller, innerTree)
 
-	// update total delegated amount
-	totalDelegated -= amount
+	s.totalDelegated -= amount
 
-	currentAmount := uint64(0)
-	if value, exists := delegatorAmount.Get(caller); exists {
-		currentAmount = value.(uint64)
-	} else {
-		panic(addDetailToError(
-			errDataNotFound,
-			ufmt.Sprintf("caller(%s) has no delegated amount", caller),
-		))
-	}
-	if currentAmount < amount {
+	currentDelegatorAmount := s.getUint64OrPanic(s.delegatorAmount, caller, "has no delegated amount")
+	if currentDelegatorAmount < amount {
 		panic(addDetailToError(
 			errNotEnoughDelegated,
-			ufmt.Sprintf("caller(%s) has only %d delegated amount(request: %d)", caller, currentAmount, amount),
+			ufmt.Sprintf("caller(%s) has only %d delegated amount (requested: %d)", caller, currentDelegatorAmount, amount),
 		))
 	}
-	delegatorAmount.Set(caller, currentAmount-amount)
+	s.delegatorAmount.Set(caller, currentDelegatorAmount-amount)
 
-	currentToAmount := uint64(0)
-	if value, exists := delegatedTo.Get(toStr); exists {
-		currentToAmount = value.(uint64)
-	} else {
-		panic(addDetailToError(
-			errDataNotFound,
-			ufmt.Sprintf("to(%s) has no delegated amount", toStr),
-		))
-	}
-	if currentToAmount < amount {
+	currentDelegatedTo := s.getUint64OrPanic(s.delegatedTo, toStr, "target has no delegated amount")
+	if currentDelegatedTo < amount {
 		panic(addDetailToError(
 			errNotEnoughDelegated,
-			ufmt.Sprintf("to(%s) has only %d delegated amount(request: %d)", toStr, currentToAmount, amount),
+			ufmt.Sprintf("to(%s) has only %d delegated amount (requested: %d)", toStr, currentDelegatedTo, amount),
 		))
 	}
-	delegatedTo.Set(toStr, currentToAmount-amount)
-
-	// update delegation history
-	delegation := DelegationHistory{
-		to:        to,
-		amount:    amount,
-		timestamp: uint64(time.Now().Unix()),
-		height:    uint64(std.GetHeight()),
-		add:       false,
-	}
-	var history []DelegationHistory
-	if value, exists := delegationHistory.Get(caller); exists {
+	s.delegatedTo.Set(toStr, currentDelegatedTo-amount)
+
+	// Update delegation history.
+	history := []DelegationHistory{}
+	if value, exists := s.delegationHistory.Get(caller); exists {
 		history = value.([]DelegationHistory)
 	}
+	delegation := newDelegationHistory(to, amount, uint64(time.Now().Unix()), uint64(std.GetHeight()), false)
 	history = append(history, delegation)
-	delegationHistory.Set(caller, history)
+	s.delegationHistory.Set(caller, history)
 
-	// update delegation stat history
-	statValue, exists := delegationSnapShotHistory.Get(toStr)
+	// Update snapshot history.
+	statValue, exists := s.DelegationSnapShotHistory().Get(toStr)
 	if !exists {
 		panic(addDetailToError(
 			errDataNotFound,
 			ufmt.Sprintf("caller(%s) has no delegation stat history", caller),
 		))
 	}
-
-	stat := statValue.([]DelegationSnapShotHistory)
+	stats := statValue.([]DelegationSnapShotHistory)
 	remainingAmount := amount
-	for i := 0; i < len(stat); i++ {
-		if stat[i].amount > 0 {
-			if stat[i].amount < remainingAmount {
-				remainingAmount -= stat[i].amount
-				stat = append(stat[:i], stat[i+1:]...)
+	for i := 0; i < len(stats); i++ {
+		if stats[i].amount > 0 {
+			if stats[i].amount < remainingAmount {
+				remainingAmount -= stats[i].amount
+				stats = append(stats[:i], stats[i+1:]...)
+				i--
 			} else {
-				stat[i].amount -= remainingAmount
-				stat[i].updatedAt = uint64(time.Now().Unix())
+				stats[i].amount -= remainingAmount
+				stats[i].updatedAt = uint64(time.Now().Unix())
+				remainingAmount = 0
 				break
 			}
 		}
 	}
-	delegationSnapShotHistory.Set(to.String(), stat)
+	s.delegationSnapShotHistory.Set(toStr, stats)
+}
+
+// delegate is the external API for delegation.
+func delegate(to std.Address, amount uint64) {
+	state.Delegate(to, amount)
+}
+
+// undelegate is the external API for undelegation.
+func undelegate(to std.Address, amount uint64) {
+	state.Undelegate(to, amount)
 }