Modified the way unstakeresult struct works

Nuklai · Sep 30, 2024 · 474b47a · 474b47a
1 parent 461b704
commit 474b47a
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 46 deletions.
diff --git a/actions/undelegate_user_stake.go b/actions/undelegate_user_stake.go
@@ -12,7 +12,6 @@ import (
 
 	"github.com/ava-labs/hypersdk/chain"
 	"github.com/ava-labs/hypersdk/codec"
-	"github.com/ava-labs/hypersdk/consts"
 	"github.com/ava-labs/hypersdk/state"
 
 	nconsts "github.com/nuklai/nuklaivm/consts"
@@ -85,13 +84,14 @@ func (u *UndelegateUserStake) Execute(
 	}
 
 	return &UndelegateUserStakeResult{
-		StakeStartBlock:      stakeStartBlock,
-		StakeEndBlock:        stakeEndBlock,
-		StakedAmount:         stakedAmount,
-		RewardAmount:         rewardAmount,
-		BalanceBeforeUnstake: balance - rewardAmount - stakedAmount,
-		BalanceAfterUnstake:  balance,
-		DistributedTo:        ownerAddress,
+		UnstakeResult: &UnstakeResult{
+			StakeStartBlock:      stakeStartBlock,
+			StakeEndBlock:        stakeEndBlock,
+			UnstakedAmount:       stakedAmount,
+			RewardAmount:         rewardAmount,
+			BalanceBeforeUnstake: balance - rewardAmount - stakedAmount,
+			BalanceAfterUnstake:  balance,
+		},
 	}, nil
 }
 
@@ -132,41 +132,21 @@ var (
 )
 
 type UndelegateUserStakeResult struct {
-	StakeStartBlock      uint64        `serialize:"true" json:"stake_start_block"`
-	StakeEndBlock        uint64        `serialize:"true" json:"stake_end_block"`
-	StakedAmount         uint64        `serialize:"true" json:"staked_amount"`
-	RewardAmount         uint64        `serialize:"true" json:"reward_amount"`
-	BalanceBeforeUnstake uint64        `serialize:"true" json:"balance_before_unstake"`
-	BalanceAfterUnstake  uint64        `serialize:"true" json:"balance_after_unstake"`
-	DistributedTo        codec.Address `serialize:"true" json:"distributed_to"`
+	*UnstakeResult
 }
 
 func (*UndelegateUserStakeResult) GetTypeID() uint8 {
 	return nconsts.UndelegateUserStakeID
 }
 
-func (*UndelegateUserStakeResult) Size() int {
-	return 6*consts.Uint64Len + codec.AddressLen
+func (r *UndelegateUserStakeResult) Size() int {
+	return r.UnstakeResult.Size()
 }
 
 func (r *UndelegateUserStakeResult) Marshal(p *codec.Packer) {
-	p.PackUint64(r.StakeStartBlock)
-	p.PackUint64(r.StakeEndBlock)
-	p.PackUint64(r.StakedAmount)
-	p.PackUint64(r.RewardAmount)
-	p.PackUint64(r.BalanceBeforeUnstake)
-	p.PackUint64(r.BalanceAfterUnstake)
-	p.PackAddress(r.DistributedTo)
+	r.UnstakeResult.Marshal(p)
 }
 
 func UnmarshalUndelegateUserStakeResult(p *codec.Packer) (codec.Typed, error) {
-	var result UndelegateUserStakeResult
-	result.StakeStartBlock = p.UnpackUint64(true)
-	result.StakeEndBlock = p.UnpackUint64(true)
-	result.StakedAmount = p.UnpackUint64(true)
-	result.RewardAmount = p.UnpackUint64(false)
-	result.BalanceBeforeUnstake = p.UnpackUint64(false)
-	result.BalanceAfterUnstake = p.UnpackUint64(true)
-	p.UnpackAddress(&result.DistributedTo)
-	return &result, p.Err()
+	return UnmarshalWithdrawValidatorStakeResult(p)
 }
diff --git a/actions/withdraw_validator_stake.go b/actions/withdraw_validator_stake.go
@@ -91,7 +91,7 @@ func (u *WithdrawValidatorStake) Execute(
 		return nil, err
 	}
 
-	return &WithdrawValidatorStakeResult{
+	return &UnstakeResult{
 		UnstakedAmount: stakedAmount,
 		RewardAmount:   rewardAmount,
 	}, nil
@@ -129,31 +129,46 @@ func UnmarshalWithdrawValidatorStake(p *codec.Packer) (chain.Action, error) {
 }
 
 var (
-	_ codec.Typed     = (*WithdrawValidatorStakeResult)(nil)
-	_ chain.Marshaler = (*WithdrawValidatorStakeResult)(nil)
+	_ codec.Typed     = (*UnstakeResult)(nil)
+	_ chain.Marshaler = (*UnstakeResult)(nil)
 )
 
-type WithdrawValidatorStakeResult struct {
-	UnstakedAmount uint64 `serialize:"true" json:"unstaked_amount"`
-	RewardAmount   uint64 `serialize:"true" json:"reward_amount"`
+type UnstakeResult struct {
+	StakeStartBlock      uint64        `serialize:"true" json:"stake_start_block"`
+	StakeEndBlock        uint64        `serialize:"true" json:"stake_end_block"`
+	UnstakedAmount       uint64        `serialize:"true" json:"unstaked_amount"`
+	RewardAmount         uint64        `serialize:"true" json:"reward_amount"`
+	BalanceBeforeUnstake uint64        `serialize:"true" json:"balance_before_unstake"`
+	BalanceAfterUnstake  uint64        `serialize:"true" json:"balance_after_unstake"`
+	DistributedTo        codec.Address `serialize:"true" json:"distributed_to"`
 }
 
-func (*WithdrawValidatorStakeResult) GetTypeID() uint8 {
+func (*UnstakeResult) GetTypeID() uint8 {
 	return nconsts.WithdrawValidatorStakeID
 }
 
-func (*WithdrawValidatorStakeResult) Size() int {
-	return 2 * consts.Uint64Len
+func (*UnstakeResult) Size() int {
+	return 6*consts.Uint64Len + codec.AddressLen
 }
 
-func (r *WithdrawValidatorStakeResult) Marshal(p *codec.Packer) {
+func (r *UnstakeResult) Marshal(p *codec.Packer) {
+	p.PackUint64(r.StakeStartBlock)
+	p.PackUint64(r.StakeEndBlock)
 	p.PackUint64(r.UnstakedAmount)
 	p.PackUint64(r.RewardAmount)
+	p.PackUint64(r.BalanceBeforeUnstake)
+	p.PackUint64(r.BalanceAfterUnstake)
+	p.PackAddress(r.DistributedTo)
 }
 
 func UnmarshalWithdrawValidatorStakeResult(p *codec.Packer) (codec.Typed, error) {
-	var result WithdrawValidatorStakeResult
-	result.UnstakedAmount = p.UnpackUint64(true)
+	var result UnstakeResult
+	result.StakeStartBlock = p.UnpackUint64(true)
+	result.StakeEndBlock = p.UnpackUint64(true)
+	result.UnstakedAmount = p.UnpackUint64(false)
 	result.RewardAmount = p.UnpackUint64(false)
+	result.BalanceBeforeUnstake = p.UnpackUint64(false)
+	result.BalanceAfterUnstake = p.UnpackUint64(true)
+	p.UnpackAddress(&result.DistributedTo)
 	return &result, p.Err()
 }
diff --git a/vm/vm.go b/vm/vm.go
@@ -84,7 +84,7 @@ func init() {
 		OutputParser.Register(&actions.BurnAssetFTResult{}, actions.UnmarshalBurnAssetFTResult),
 		OutputParser.Register(&actions.BurnAssetNFTResult{}, actions.UnmarshalBurnAssetNFTResult),
 		OutputParser.Register(&actions.RegisterValidatorStakeResult{}, actions.UnmarshalRegisterValidatorStakeResult),
-		OutputParser.Register(&actions.WithdrawValidatorStakeResult{}, actions.UnmarshalWithdrawValidatorStakeResult),
+		OutputParser.Register(&actions.UnstakeResult{}, actions.UnmarshalWithdrawValidatorStakeResult),
 		OutputParser.Register(&actions.ClaimValidatorStakeRewardsResult{}, actions.UnmarshalClaimValidatorStakeRewardsResult),
 		OutputParser.Register(&actions.DelegateUserStakeResult{}, actions.UnmarshalDelegateUserStakeResult),
 		OutputParser.Register(&actions.UndelegateUserStakeResult{}, actions.UnmarshalUndelegateUserStakeResult),