Added LSM data to init/export genesis

docs/core/proto-docs.md

@@ -493,6 +493,7 @@
 - [cosmos/staking/v1beta1/genesis.proto](#cosmos/staking/v1beta1/genesis.proto)
     - [GenesisState](#cosmos.staking.v1beta1.GenesisState)
     - [LastValidatorPower](#cosmos.staking.v1beta1.LastValidatorPower)
+    - [TokenizeShareLock](#cosmos.staking.v1beta1.TokenizeShareLock)
 
 - [cosmos/staking/v1beta1/query.proto](#cosmos/staking/v1beta1/query.proto)
     - [QueryAllTokenizeShareRecordsRequest](#cosmos.staking.v1beta1.QueryAllTokenizeShareRecordsRequest)
@@ -7083,6 +7084,8 @@ GenesisState defines the staking module's genesis state.
 | `exported` | [bool](#bool) |  |  |
 | `tokenize_share_records` | [TokenizeShareRecord](#cosmos.staking.v1beta1.TokenizeShareRecord) | repeated | store tokenize share records to provide reward to record owners |
 | `last_tokenize_share_record_id` | [uint64](#uint64) |  | last tokenize share record id, used for next share record id calculation |
+| `total_liquid_staked_tokens` | [bytes](#bytes) |  | total number of liquid staked tokens at genesis |
+| `tokenize_share_locks` | [TokenizeShareLock](#cosmos.staking.v1beta1.TokenizeShareLock) | repeated | tokenize shares locks at genesis |
 
 
 
@@ -7104,6 +7107,23 @@ LastValidatorPower required for validator set update logic.
 
 
 
+
+<a name="cosmos.staking.v1beta1.TokenizeShareLock"></a>
+
+### TokenizeShareLock
+TokenizeSharesLock required for specifying account locks at genesis
+
+
+| Field | Type | Label | Description |
+| ----- | ---- | ----- | ----------- |
+| `address` | [string](#string) |  | Address of the account that is locked |
+| `status` | [string](#string) |  | Status of the lock (LOCKED or LOCK_EXPIRING) |
+| `completion_time` | [google.protobuf.Timestamp](#google.protobuf.Timestamp) |  | Completion time if the lock is expiring |
+
+
+
+
+
  <!-- end messages -->
 
  <!-- end enums -->

proto/cosmos/staking/v1beta1/genesis.proto

@@ -5,6 +5,7 @@ option go_package = "github.com/cosmos/cosmos-sdk/x/staking/types";
 
 import "gogoproto/gogo.proto";
 import "cosmos/staking/v1beta1/staking.proto";
+import "google/protobuf/timestamp.proto";
 
 // GenesisState defines the staking module's genesis state.
 message GenesisState {
@@ -44,6 +45,27 @@ message GenesisState {
 
   // last tokenize share record id, used for next share record id calculation
   uint64 last_tokenize_share_record_id = 10;
+
+  // total number of liquid staked tokens at genesis
+  bytes total_liquid_staked_tokens = 11 [
+    (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
+    (gogoproto.moretags)   = "yaml:\"total_liquid_staked_tokens\"",
+    (gogoproto.nullable)   = false
+  ];
+
+  // tokenize shares locks at genesis
+  repeated TokenizeShareLock tokenize_share_locks = 12 [(gogoproto.nullable) = false];
+}
+
+// TokenizeSharesLock required for specifying account locks at genesis
+message TokenizeShareLock {
+  // Address of the account that is locked
+  string address = 1;
+  // Status of the lock (LOCKED or LOCK_EXPIRING)
+  string status = 2;
+  // Completion time if the lock is expiring
+  google.protobuf.Timestamp completion_time = 3
+      [(gogoproto.nullable) = false, (gogoproto.stdtime) = true, (gogoproto.moretags) = "yaml:\"completion_time\""];
 }
 
 // LastValidatorPower required for validator set update logic.

x/staking/genesis.go

@@ -150,6 +150,48 @@ func InitGenesis(
 		}
 	}
 
+	// Set the total liquid staked tokens
+	keeper.SetTotalLiquidStakedTokens(ctx, data.TotalLiquidStakedTokens)
+
+	// Set each tokenize share record, as well as the last tokenize share record ID
+	latestId := uint64(0)
+	for _, tokenizeShareRecord := range data.TokenizeShareRecords {
+		if err := keeper.AddTokenizeShareRecord(ctx, tokenizeShareRecord); err != nil {
+			panic(err)
+		}
+		if tokenizeShareRecord.Id > latestId {
+			latestId = tokenizeShareRecord.Id
+		}
+	}
+	if data.LastTokenizeShareRecordId < latestId {
+		panic("Tokenize share record specified with ID greater than the latest ID")
+	}
+	keeper.SetLastTokenizeShareRecordID(ctx, data.LastTokenizeShareRecordId)
+
+	// Set the tokenize shares locks for accounts that have disabled tokenizing shares
+	// The lock can either be in status LOCKED or LOCK_EXPIRING
+	// If it is in status LOCK_EXPIRING, a the unlocking must also be queued
+	for _, tokenizeShareLock := range data.TokenizeShareLocks {
+		address := sdk.MustAccAddressFromBech32(tokenizeShareLock.Address)
+
+		switch tokenizeShareLock.Status {
+		case types.TOKENIZE_SHARE_LOCK_STATUS_LOCKED.String():
+			keeper.AddTokenizeSharesLock(ctx, address)
+
+		case types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING.String():
+			completionTime := tokenizeShareLock.CompletionTime
+
+			authorizations := keeper.GetPendingTokenizeShareAuthorizations(ctx, completionTime)
+			authorizations.Addresses = append(authorizations.Addresses, address.String())
+
+			keeper.SetPendingTokenizeShareAuthorizations(ctx, completionTime, authorizations)
+			keeper.SetTokenizeSharesUnlockTime(ctx, address, completionTime)
+
+		default:
+			panic(fmt.Sprintf("Unsupported tokenize share lock status %s", tokenizeShareLock.Status))
+		}
+	}
+
 	return res
 }
 
@@ -179,14 +221,18 @@ func ExportGenesis(ctx sdk.Context, keeper keeper.Keeper) *types.GenesisState {
 	})
 
 	return &types.GenesisState{
-		Params:               keeper.GetParams(ctx),
-		LastTotalPower:       keeper.GetLastTotalPower(ctx),
-		LastValidatorPowers:  lastValidatorPowers,
-		Validators:           keeper.GetAllValidators(ctx),
-		Delegations:          keeper.GetAllDelegations(ctx),
-		UnbondingDelegations: unbondingDelegations,
-		Redelegations:        redelegations,
-		Exported:             true,
+		Params:                    keeper.GetParams(ctx),
+		LastTotalPower:            keeper.GetLastTotalPower(ctx),
+		LastValidatorPowers:       lastValidatorPowers,
+		Validators:                keeper.GetAllValidators(ctx),
+		Delegations:               keeper.GetAllDelegations(ctx),
+		UnbondingDelegations:      unbondingDelegations,
+		Redelegations:             redelegations,
+		Exported:                  true,
+		TokenizeShareRecords:      keeper.GetAllTokenizeShareRecords(ctx),
+		LastTokenizeShareRecordId: keeper.GetLastTokenizeShareRecordID(ctx),
+		TotalLiquidStakedTokens:   keeper.GetTotalLiquidStakedTokens(ctx),
+		TokenizeShareLocks:        keeper.GetAllTokenizeSharesLocks(ctx),
 	}
 }
 

x/staking/genesis_test.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"log"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -242,3 +243,49 @@ func TestValidateGenesis(t *testing.T) {
 		})
 	}
 }
+
+func TestInitExportLiquidStakingGenesis(t *testing.T) {
+	app := simapp.Setup(false)
+	ctx := app.NewContext(false, tmproto.Header{})
+
+	addresses := simapp.AddTestAddrs(app, ctx, 2, sdk.OneInt())
+	address1, address2 := addresses[0], addresses[1]
+
+	// Mock out a genesis state
+	inGenesisState := types.GenesisState{
+		Params: types.DefaultParams(),
+		TokenizeShareRecords: []types.TokenizeShareRecord{
+			{Id: 1, Owner: address1.String(), ModuleAccount: "module1", Validator: "val1"},
+			{Id: 2, Owner: address2.String(), ModuleAccount: "module2", Validator: "val2"},
+		},
+		LastTokenizeShareRecordId: 2,
+		TotalLiquidStakedTokens:   sdk.NewInt(1_000_000),
+		TokenizeShareLocks: []types.TokenizeShareLock{
+			{
+				Address: address1.String(),
+				Status:  types.TOKENIZE_SHARE_LOCK_STATUS_LOCKED.String(),
+			},
+			{
+				Address:        address2.String(),
+				Status:         types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING.String(),
+				CompletionTime: time.Date(2023, 1, 1, 1, 0, 0, 0, time.UTC),
+			},
+		},
+	}
+
+	// Call init and then export genesis - confirming the same state is returned
+	staking.InitGenesis(ctx, app.StakingKeeper, app.AccountKeeper, app.BankKeeper, &inGenesisState)
+	outGenesisState := *staking.ExportGenesis(ctx, app.StakingKeeper)
+
+	require.ElementsMatch(t, inGenesisState.TokenizeShareRecords, outGenesisState.TokenizeShareRecords,
+		"tokenize share records")
+
+	require.Equal(t, inGenesisState.LastTokenizeShareRecordId, outGenesisState.LastTokenizeShareRecordId,
+		"last tokenize share record ID")
+
+	require.Equal(t, inGenesisState.TotalLiquidStakedTokens.Int64(), outGenesisState.TotalLiquidStakedTokens.Int64(),
+		"total liquid staked")
+
+	require.ElementsMatch(t, inGenesisState.TokenizeShareLocks, outGenesisState.TokenizeShareLocks,
+		"tokenize share locks")
+}

x/staking/keeper/liquid_stake.go

@@ -203,6 +203,40 @@ func (k Keeper) GetTokenizeSharesLock(ctx sdk.Context, address sdk.AccAddress) (
 	return types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING, unlockTime
 }
 
+// Returns all tokenize share locks
+func (k Keeper) GetAllTokenizeSharesLocks(ctx sdk.Context) (tokenizeShareLocks []types.TokenizeShareLock) {
+	store := ctx.KVStore(k.storeKey)
+
+	iterator := sdk.KVStorePrefixIterator(store, types.TokenizeSharesLockPrefix)
+	defer iterator.Close()
+
+	for ; iterator.Valid(); iterator.Next() {
+		addressBz := iterator.Key()[2:] // remove prefix bytes and address length
+		unlockTime, err := sdk.ParseTimeBytes(iterator.Value())
+		if err != nil {
+			panic(err)
+		}
+
+		var status types.TokenizeShareLockStatus
+		if unlockTime.IsZero() {
+			status = types.TOKENIZE_SHARE_LOCK_STATUS_LOCKED
+		} else {
+			status = types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING
+		}
+
+		bechPrefix := sdk.GetConfig().GetBech32AccountAddrPrefix()
+		lock := types.TokenizeShareLock{
+			Address:        sdk.MustBech32ifyAddressBytes(bechPrefix, addressBz),
+			Status:         status.String(),
+			CompletionTime: unlockTime,
+		}
+
+		tokenizeShareLocks = append(tokenizeShareLocks, lock)
+	}
+
+	return tokenizeShareLocks
+}
+
 // Stores a list of addresses pending tokenize share unlocking at the same time
 func (k Keeper) SetPendingTokenizeShareAuthorizations(ctx sdk.Context, completionTime time.Time, authorizations types.PendingTokenizeShareAuthorizations) {
 	store := ctx.KVStore(k.storeKey)
@@ -251,7 +285,7 @@ func (k Keeper) RemoveExpiredTokenizeShareLocks(ctx sdk.Context, blockTime time.
 
 	// iterators all time slices from time 0 until the current block time
 	prefixEnd := sdk.InclusiveEndBytes(types.GetTokenizeShareAuthorizationTimeKey(blockTime))
-	iterator := store.Iterator(types.TokenizeSharesUnlockQueueKey, prefixEnd)
+	iterator := store.Iterator(types.TokenizeSharesUnlockQueuePrefix, prefixEnd)
 	defer iterator.Close()
 
 	unlockedAddresses = []string{}

x/staking/keeper/liquid_stake_test.go

@@ -751,6 +751,51 @@ func TestTokenizeSharesLock(t *testing.T) {
 	require.Equal(t, unlocked, status.String(), "addressB unlocked at end")
 }
 
+// Tests GetAllTokenizeSharesLocks
+func TestGetAllTokenizeSharesLocks(t *testing.T) {
+	_, app, ctx := createTestInput()
+
+	addresses := simapp.AddTestAddrs(app, ctx, 4, sdk.NewInt(1))
+
+	// Set 2 locked accounts, and two accounts with a lock expiring
+	app.StakingKeeper.AddTokenizeSharesLock(ctx, addresses[0])
+	app.StakingKeeper.AddTokenizeSharesLock(ctx, addresses[1])
+
+	unlockTime1 := time.Date(2023, 1, 1, 1, 0, 0, 0, time.UTC)
+	unlockTime2 := time.Date(2023, 1, 2, 1, 0, 0, 0, time.UTC)
+	app.StakingKeeper.SetTokenizeSharesUnlockTime(ctx, addresses[2], unlockTime1)
+	app.StakingKeeper.SetTokenizeSharesUnlockTime(ctx, addresses[3], unlockTime2)
+
+	// Defined expected locks after GetAll
+	expectedLocks := map[string]types.TokenizeShareLock{
+		addresses[0].String(): {
+			Status: types.TOKENIZE_SHARE_LOCK_STATUS_LOCKED.String(),
+		},
+		addresses[1].String(): {
+			Status: types.TOKENIZE_SHARE_LOCK_STATUS_LOCKED.String(),
+		},
+		addresses[2].String(): {
+			Status:         types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING.String(),
+			CompletionTime: unlockTime1,
+		},
+		addresses[3].String(): {
+			Status:         types.TOKENIZE_SHARE_LOCK_STATUS_LOCK_EXPIRING.String(),
+			CompletionTime: unlockTime2,
+		},
+	}
+
+	// Check output from GetAll
+	actualLocks := app.StakingKeeper.GetAllTokenizeSharesLocks(ctx)
+	require.Len(t, actualLocks, len(expectedLocks), "number of locks")
+
+	for i, actual := range actualLocks {
+		expected, ok := expectedLocks[actual.Address]
+		require.True(t, ok, "address %s not expected", actual.Address)
+		require.Equal(t, expected.Status, actual.Status, "tokenize share lock #%d status", i)
+		require.Equal(t, expected.CompletionTime, actual.CompletionTime, "tokenize share lock #%d completion time", i)
+	}
+}
+
 // Test Get/SetPendingTokenizeShareAuthorizations
 func TestPendingTokenizeShareAuthorizations(t *testing.T) {
 	_, app, ctx := createTestInput()