Fixed some issues with emission balancer and updated the readme

cmd/nuklai-cli/cmd/emission.go

@@ -30,7 +30,7 @@ var emissionInfoCmd = &cobra.Command{
 		ncli := nclients[0]
 
 		// Get emission info
-		_, _, _, err = handler.GetEmissionInfo(ctx, ncli)
+		_, _, _, _, _, err = handler.GetEmissionInfo(ctx, ncli)
 		if err != nil {
 			return err
 		}

cmd/nuklai-cli/cmd/genesis.go

@@ -5,13 +5,16 @@ package cmd
 
 import (
 	"encoding/json"
+	"errors"
 	"os"
 
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 
 	"github.com/ava-labs/hypersdk/chain"
+	"github.com/ava-labs/hypersdk/codec"
 
+	nconsts "github.com/nuklai/nuklaivm/consts"
 	"github.com/nuklai/nuklaivm/genesis"
 )
 
@@ -23,10 +26,10 @@ var genesisCmd = &cobra.Command{
 }
 
 var genGenesisCmd = &cobra.Command{
-	Use:   "generate [custom allocates file] [options]",
+	Use:   "generate [custom allocates file] [emission balancer file] [options]",
 	Short: "Creates a new genesis in the default location",
 	PreRunE: func(cmd *cobra.Command, args []string) error {
-		if len(args) != 1 {
+		if len(args) != 2 {
 			return ErrInvalidArgs
 		}
 		return nil
@@ -58,6 +61,7 @@ var genGenesisCmd = &cobra.Command{
 			g.MinBlockGap = minBlockGap
 		}
 
+		// Read custom allocations file
 		a, err := os.ReadFile(args[0])
 		if err != nil {
 			return err
@@ -68,6 +72,31 @@ var genGenesisCmd = &cobra.Command{
 		}
 		g.CustomAllocation = allocs
 
+		// Read emission balancer file
+		eb, err := os.ReadFile(args[1])
+		if err != nil {
+			return err
+		}
+		emissionBalancer := genesis.EmissionBalancer{}
+		if err := json.Unmarshal(eb, &emissionBalancer); err != nil {
+			return err
+		}
+		totalSupply := uint64(0)
+		for _, alloc := range allocs {
+			if _, err := codec.ParseAddressBech32(nconsts.HRP, alloc.Address); err != nil {
+				return err
+			}
+			totalSupply += alloc.Balance
+		}
+		if _, err := codec.ParseAddressBech32(nconsts.HRP, emissionBalancer.EmissionAddress); err != nil {
+			return err
+		}
+		if emissionBalancer.TotalSupply > 0 && totalSupply != emissionBalancer.TotalSupply {
+			return errors.New("total supply in emission balancer file does not match the sum of balances in custom allocates file")
+		}
+		emissionBalancer.TotalSupply = totalSupply
+		g.EmissionBalancer = emissionBalancer
+
 		b, err := json.Marshal(g)
 		if err != nil {
 			return err

cmd/nuklai-cli/cmd/handler.go

@@ -5,6 +5,7 @@ package cmd
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/ava-labs/avalanchego/ids"
 	"github.com/ava-labs/hypersdk/chain"
@@ -87,15 +88,18 @@ func (*Handler) GetAssetInfo(
 	if err != nil {
 		return nil, 0, 0, ids.Empty, err
 	}
+	balanceToShow := fmt.Sprintf("%s %s", hutils.FormatBalance(balance, decimals), symbol)
+	if assetID != ids.Empty {
+		balanceToShow = fmt.Sprintf("%s %s", hutils.FormatBalance(balance, decimals), assetID)
+	}
 	if balance == 0 {
-		hutils.Outf("{{red}}balance:{{/}} 0 %s\n", assetID)
+		hutils.Outf("{{red}}balance:{{/}} %s\n", balanceToShow)
 		hutils.Outf("{{red}}please send funds to %s{{/}}\n", saddr)
 		hutils.Outf("{{red}}exiting...{{/}}\n")
 	} else {
 		hutils.Outf(
-			"{{yellow}}balance:{{/}} %s %s\n",
-			hutils.FormatBalance(balance, decimals),
-			symbol,
+			"{{yellow}}balance:{{/}} %s\n",
+			balanceToShow,
 		)
 	}
 	return symbol, decimals, balance, sourceChainID, nil
@@ -177,19 +181,26 @@ func (h *Handler) DefaultNuklaiVMJSONRPCClient(checkAllChains bool) ([]*nrpc.JSO
 func (*Handler) GetEmissionInfo(
 	ctx context.Context,
 	cli *nrpc.JSONRPCClient,
-) (uint64, uint64, uint64, error) {
-	totalSupply, maxSupply, rewardsPerBlock, err := cli.EmissionInfo(ctx)
+) (uint64, uint64, uint64, string, uint64, error) {
+	totalSupply, maxSupply, rewardsPerBlock, emissionAccount, err := cli.EmissionInfo(ctx)
+	if err != nil {
+		return 0, 0, 0, "", 0, err
+	}
+
+	emissionAddress, err := codec.AddressBech32(nconsts.HRP, emissionAccount.Address)
 	if err != nil {
-		return 0, 0, 0, err
+		return 0, 0, 0, "", 0, err
 	}
 
 	hutils.Outf(
-		"{{yellow}}emission info: {{/}}\nTotalSupply=%d MaxSupply=%d RewardsPerBlock=%d\n",
+		"{{yellow}}emission info: {{/}}\nTotalSupply=%d MaxSupply=%d RewardsPerBlock=%d EmissionAddress=%s EmissionBalance=%d\n",
 		totalSupply,
 		maxSupply,
 		rewardsPerBlock,
+		emissionAddress,
+		emissionAccount.Balance,
 	)
-	return totalSupply, maxSupply, rewardsPerBlock, err
+	return totalSupply, maxSupply, rewardsPerBlock, emissionAddress, emissionAccount.Balance, err
 }
 
 func (*Handler) GetAllValidators(

controller/controller.go

@@ -18,7 +18,6 @@ import (
 	"github.com/ava-labs/hypersdk/codec"
 	"github.com/ava-labs/hypersdk/gossiper"
 	hrpc "github.com/ava-labs/hypersdk/rpc"
-	"github.com/ava-labs/hypersdk/state"
 	hstorage "github.com/ava-labs/hypersdk/storage"
 	"github.com/ava-labs/hypersdk/vm"
 	"go.uber.org/zap"
@@ -153,7 +152,11 @@ func (c *Controller) Initialize(
 		// We only get the validators in non-test mode
 		currentValidators, _ = inner.CurrentValidators(context.TODO())
 	}
-	c.emission = emission.New(c, c.genesis.MaxSupply, c.genesis.RewardsPerBlock, currentValidators)
+	emissionAddr, err := codec.ParseAddressBech32(nconsts.HRP, c.genesis.EmissionBalancer.EmissionAddress)
+	if err != nil {
+		return nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, err
+	}
+	c.emission = emission.New(c, c.genesis.EmissionBalancer.TotalSupply, c.genesis.EmissionBalancer.MaxSupply, c.genesis.EmissionBalancer.RewardsPerBlock, emissionAddr, currentValidators)
 
 	return c.config, c.genesis, build, gossip, blockDB, stateDB, apis, nconsts.ActionRegistry, nconsts.AuthRegistry, auth.Engines(), nil
 }
@@ -234,44 +237,20 @@ func (c *Controller) Accepted(ctx context.Context, blk *chain.StatelessBlock) er
 		}
 	}
 
-	// Retrieve the vm state
-	stateDB, err := c.inner.State()
-	if err != nil {
-		return err
-	}
-	// Retrieve the state.Mutable to write to
-	mu := state.NewSimpleMutable(stateDB)
-	emissionAddr, err := codec.ParseAddressBech32(nconsts.HRP, c.genesis.EmissionAddress)
-	if err != nil {
-		return err // This should never happen
-	}
-
 	// Distribute fees
 	if totalFee > 0 {
-		if feesForEmission := c.emission.FeesToDistribute(totalFee); feesForEmission > 0 {
-			// Give 50% fees to Emission
-			if err := storage.AddBalance(ctx, mu, emissionAddr, ids.Empty, feesForEmission, true); err != nil {
-				c.inner.Logger().Error("failed to distribute fees to emission address", zap.Error(err))
-			}
-			if err := mu.Commit(ctx); err != nil {
-				c.inner.Logger().Error("failed to commit fees to emission address", zap.Error(err))
-			}
-			c.inner.Logger().Info("distributed fees to Emission and Validators", zap.Uint64("current block height", c.inner.LastAcceptedBlock().Height()), zap.Uint64("total fee", totalFee), zap.Uint64("total supply", c.emission.GetTotalSupply()), zap.Uint64("max supply", c.emission.GetMaxSupply()))
+		c.emission.DistributeFees(totalFee)
+		emissionAddress, err := codec.AddressBech32(nconsts.HRP, c.emission.GetEmissionAccount().Address)
+		if err != nil {
+			return err // This should never happen
 		}
+		c.inner.Logger().Info("distributed fees to Emission and Validators", zap.Uint64("current block height", c.inner.LastAcceptedBlock().Height()), zap.Uint64("total fee", totalFee), zap.Uint64("total supply", c.emission.GetTotalSupply()), zap.Uint64("max supply", c.emission.GetMaxSupply()), zap.Uint64("rewards per block", c.emission.GetRewardsPerBlock()), zap.String("emission address", emissionAddress), zap.Uint64("emission address balance", c.emission.GetEmissionAccount().Balance))
 	}
 
 	// Mint new NAI if needed
-	mintNewNAI, mintToEmissionAddr := c.emission.MintNewNAI()
+	mintNewNAI := c.emission.MintNewNAI()
 	if mintNewNAI > 0 {
 		c.emission.AddToTotalSupply(mintNewNAI)
-		if mintToEmissionAddr {
-			if err := storage.AddBalance(ctx, mu, emissionAddr, ids.Empty, mintNewNAI, true); err != nil {
-				c.inner.Logger().Error("failed to mint new NAI to emission address", zap.Error(err))
-			}
-			if err := mu.Commit(ctx); err != nil {
-				c.inner.Logger().Error("failed to commit new NAI to emission address", zap.Error(err))
-			}
-		}
 		c.inner.Logger().Info("minted new NAI", zap.Uint64("current block height", c.inner.LastAcceptedBlock().Height()), zap.Uint64("newly minted NAI", mintNewNAI), zap.Uint64("total supply", c.emission.GetTotalSupply()), zap.Uint64("max supply", c.emission.GetMaxSupply()))
 		c.metrics.mintNAI.Add(float64(mintNewNAI))
 	}

controller/resolutions.go

@@ -58,8 +58,8 @@ func (c *Controller) GetLoanFromState(
 	return storage.GetLoanFromState(ctx, c.inner.ReadState, asset, destination)
 }
 
-func (c *Controller) GetEmissionInfo() (uint64, uint64, uint64, error) {
-	return c.emission.GetTotalSupply(), c.emission.GetMaxSupply(), c.emission.GetRewardsPerBlock(), nil
+func (c *Controller) GetEmissionInfo() (uint64, uint64, uint64, *emission.EmissionAccount, error) {
+	return c.emission.GetTotalSupply(), c.emission.GetMaxSupply(), c.emission.GetRewardsPerBlock(), c.emission.GetEmissionAccount(), nil
 }
 
 func (c *Controller) GetAllValidators() ([]*emission.Validator, error) {

emission/emission.go

@@ -46,12 +46,18 @@ type Validator struct {
 	StakedReward  uint64                `json:"stakedReward"`
 }
 
+type EmissionAccount struct {
+	Address codec.Address `json:"address"`
+	Balance uint64        `json:"balance"`
+}
+
 type Emission struct {
 	c Controller
 
 	totalSupply     uint64
 	maxSupply       uint64
 	rewardsPerBlock uint64
+	emissionAccount EmissionAccount
 
 	validators     map[ids.NodeID]*Validator // the key is NodeID
 	minStakeAmount uint64
@@ -60,7 +66,7 @@ type Emission struct {
 }
 
 // New initializes the Emission with the maximum supply
-func New(c Controller, maxSupply, rewardsPerBlock uint64, currentValidators map[ids.NodeID]*validators.GetValidatorOutput) *Emission {
+func New(c Controller, totalSupply, maxSupply, rewardsPerBlock uint64, emissionAddress codec.Address, currentValidators map[ids.NodeID]*validators.GetValidatorOutput) *Emission {
 	once.Do(func() {
 		c.Logger().Info("setting maxSupply and rewardsPerBlock for emission")
 
@@ -76,10 +82,17 @@ func New(c Controller, maxSupply, rewardsPerBlock uint64, currentValidators map[
 			validators[nodeID] = newValidator
 		}
 
+		emissionAccount := &EmissionAccount{
+			Address: emissionAddress,
+			Balance: 0,
+		}
+
 		emission = &Emission{
 			c:               c,
+			totalSupply:     totalSupply,
 			maxSupply:       maxSupply,
 			rewardsPerBlock: rewardsPerBlock,
+			emissionAccount: *emissionAccount,
 			validators:      validators,
 			minStakeAmount:  100,
 		}
@@ -119,6 +132,11 @@ func (e *Emission) GetRewardsPerBlock() uint64 {
 	return e.rewardsPerBlock
 }
 
+func (e *Emission) GetEmissionAccount() *EmissionAccount {
+	e.c.Logger().Info("fetching Emission Account")
+	return &e.emissionAccount
+}
+
 // StakeValidator stakes the validator
 func (e *Emission) StakeToValidator(txID ids.ID, actor codec.Address, currentValidators map[ids.NodeID]*validators.GetValidatorOutput, startLockUp uint64, nodeIDByte []byte, stakedAmount uint64) error {
 	e.lock.Lock()
@@ -249,7 +267,7 @@ func (e *Emission) getAllValidators() []*Validator {
 }
 
 // MintNewNAI mints new tokens and distributes them to validators
-func (e *Emission) MintNewNAI() (uint64, bool) {
+func (e *Emission) MintNewNAI() uint64 {
 	e.lock.Lock()
 	defer e.lock.Unlock()
 
@@ -258,13 +276,15 @@ func (e *Emission) MintNewNAI() (uint64, bool) {
 		mintNewNAI = e.maxSupply - e.totalSupply // Adjust to not exceed max supply
 	}
 	if mintNewNAI == 0 {
-		return 0, false // Nothing to mint
+		return 0 // Nothing to mint
 	}
 
 	totalStaked := e.totalStaked()
 	// No validators to distribute rewards to if totalStaked is 0
+	// So, give all the rewards to Emission Address
 	if totalStaked == 0 {
-		return mintNewNAI, true
+		e.emissionAccount.Balance += mintNewNAI
+		return mintNewNAI
 	}
 
 	// Distribute rewards based on stake proportion
@@ -273,17 +293,24 @@ func (e *Emission) MintNewNAI() (uint64, bool) {
 			v.StakedReward += mintNewNAI * v.StakedAmount / totalStaked
 		}
 	}
-	return mintNewNAI, false
+	return mintNewNAI
 }
 
-func (e *Emission) FeesToDistribute(fee uint64) uint64 {
+func (e *Emission) DistributeFees(fee uint64) {
 	e.lock.Lock()
 	defer e.lock.Unlock()
 
+	if e.totalSupply+fee > e.maxSupply {
+		fee = e.maxSupply - e.totalSupply // Adjust to not exceed max supply
+	}
+
+	// Give 50% fees to Emission
 	feesForEmission := fee / 2
-	feesForValidators := fee - feesForEmission
+	e.emissionAccount.Balance += feesForEmission
 
 	// Give remaining to Validators
+	feesForValidators := fee - feesForEmission
+
 	totalStaked := e.totalStaked()
 	if totalStaked > 0 {
 		// Distribute rewards based on stake proportion
@@ -293,8 +320,6 @@ func (e *Emission) FeesToDistribute(fee uint64) uint64 {
 			}
 		}
 	}
-
-	return feesForEmission
 }
 
 // totalStaked calculates the total amount staked by all validators