add shrink-votes cmd

main.go

@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"slices"
 	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/peterbourgon/ff/v3/ffcli"
@@ -51,7 +52,7 @@ func main() {
 			tallyCmd(), accountsCmd(), genesisCmd(), autoStakingCmd(),
 			distributionCmd(), top20Cmd(), propJSONCmd(),
 			signTxCmd(), vestingCmd(), depositThrottlingCmd(),
-			tallyGenesisCmd(),
+			tallyGenesisCmd(), shrinkVotesCmd(),
 		},
 		Exec: func(ctx context.Context, args []string) error {
 			return flag.ErrHelp
@@ -63,6 +64,24 @@ func main() {
 	}
 }
 
+func shrinkVotesCmd() *ffcli.Command {
+	return &ffcli.Command{
+		Name:       "shrink-votes",
+		ShortUsage: "govbox shrink-votes <genesis.json> <high>",
+		ShortHelp:  "Outputs a genesis where only the first <high> votes are kept from <genesis.json>",
+		Exec: func(ctx context.Context, args []string) error {
+			if len(args) != 2 {
+				return flag.ErrHelp
+			}
+			high, err := strconv.Atoi(args[1])
+			if err != nil {
+				return err
+			}
+			return shrinkVotes(ctx, args[0], high)
+		},
+	}
+}
+
 func tallyGenesisCmd() *ffcli.Command {
 	fs := flag.NewFlagSet("tallyGenesis", flag.ContinueOnError)
 	numVals := fs.Int("numVals", 1, "number of validators")
@@ -73,7 +92,7 @@ func tallyGenesisCmd() *ffcli.Command {
 	return &ffcli.Command{
 		Name:       "tally-genesis",
 		ShortUsage: "govbox tally-genesis <genesis.json>",
-		ShortHelp: `Generate a genesis with validators, delegators, governors, delegations, votes and one proposal.
+		ShortHelp: `Outputs a genesis with validators, delegators, governors, delegations, votes and one proposal.
 Used to evaluate the performance of the governance tally.`,
 		FlagSet: fs,
 		Exec: func(ctx context.Context, args []string) error {

shrink_votes.go

@@ -0,0 +1,49 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+
+	govtypes "github.com/atomone-hub/atomone/x/gov/types/v1"
+
+	tmjson "github.com/cometbft/cometbft/libs/json"
+	tmtypes "github.com/cometbft/cometbft/types"
+)
+
+func shrinkVotes(_ context.Context, genesisFile string, high int) error {
+	// Read input genesis and update it
+	bz, err := os.ReadFile(genesisFile)
+	if err != nil {
+		return fmt.Errorf("readfile %s: %w", genesisFile, err)
+	}
+	var genesisState tmtypes.GenesisDoc
+	if err := tmjson.Unmarshal(bz, &genesisState); err != nil {
+		return fmt.Errorf("unmarshal genesis doc: %w", err)
+	}
+	var appState map[string]json.RawMessage
+	if err := json.Unmarshal(genesisState.AppState, &appState); err != nil {
+		return fmt.Errorf("unmarshal appstate: %w", err)
+	}
+	var govGen govtypes.GenesisState
+	if err := cdc.UnmarshalJSON(appState["gov"], &govGen); err != nil {
+		return fmt.Errorf("umarshal gov genesis: %w", err)
+	}
+	fmt.Fprintf(os.Stderr, "Shrinking votes from %d to %d\n", len(govGen.Votes), high)
+	govGen.Votes = govGen.Votes[:high]
+	appState["gov"], err = cdc.MarshalJSON(&govGen)
+	if err != nil {
+		return fmt.Errorf("marshal gov genesis: %w", err)
+	}
+	genesisState.AppState, err = json.MarshalIndent(appState, "", "  ")
+	if err != nil {
+		return err
+	}
+	bz, err = tmjson.MarshalIndent(genesisState, "", "  ")
+	if err != nil {
+		return err
+	}
+	fmt.Println(string(bz))
+	return nil
+}