Genesis export via code prefix

CosmWasm · Jul 20, 2020 · e6e4839 · e6e4839
1 parent 65e4e3a
commit e6e4839
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 16 deletions.
diff --git a/x/wasm/internal/keeper/genesis.go b/x/wasm/internal/keeper/genesis.go
@@ -57,23 +57,20 @@ func ExportGenesis(ctx sdk.Context, keeper Keeper) types.GenesisState {
 
 	genState.Params = keeper.GetParams(ctx)
 
-	maxCodeID := keeper.GetNextCodeID(ctx)
-	for i := uint64(1); i < maxCodeID; i++ {
-		if !keeper.containsCodeInfo(ctx, i) {
-			continue
-		}
-		bytecode, err := keeper.GetByteCode(ctx, i)
+	keeper.IterateCodeInfos(ctx, func(codeID uint64, info types.CodeInfo) bool {
+		bytecode, err := keeper.GetByteCode(ctx, codeID)
 		if err != nil {
 			panic(err)
 		}
 		genState.Codes = append(genState.Codes, types.Code{
-			CodeID:     i,
-			CodeInfo:   *keeper.GetCodeInfo(ctx, i),
+			CodeID:     codeID,
+			CodeInfo:   info,
 			CodesBytes: bytecode,
 		})
-	}
+		return false
+	})
 
-	keeper.ListContractInfo(ctx, func(addr sdk.AccAddress, contract types.ContractInfo) bool {
+	keeper.IterateContractInfo(ctx, func(addr sdk.AccAddress, contract types.ContractInfo) bool {
 		contractStateIterator := keeper.GetContractState(ctx, addr)
 		var state []types.Model
 		for ; contractStateIterator.Valid(); contractStateIterator.Next() {

diff --git a/x/wasm/internal/keeper/genesis_test.go b/x/wasm/internal/keeper/genesis_test.go
@@ -70,7 +70,7 @@ func TestGenesisExportImport(t *testing.T) {
 	require.NoError(t, err)
 
 	// reset contract history in source DB for comparision with dest DB
-	srcKeeper.ListContractInfo(srcCtx, func(address sdk.AccAddress, info wasmTypes.ContractInfo) bool {
+	srcKeeper.IterateContractInfo(srcCtx, func(address sdk.AccAddress, info wasmTypes.ContractInfo) bool {
 		info.ResetFromGenesis(srcCtx)
 		srcKeeper.setContractInfo(srcCtx, address, &info)
 		return false
@@ -91,7 +91,7 @@ func TestGenesisExportImport(t *testing.T) {
 		dstIT := dstCtx.KVStore(dstStoreKeys[j]).Iterator(nil, nil)
 
 		for i := 0; srcIT.Valid(); i++ {
-			require.True(t, dstIT.Valid(), "destination DB has less elements than source. Missing: %q", srcIT.Key())
+			require.True(t, dstIT.Valid(), "[%s] destination DB has less elements than source. Missing: %s", srcStoreKeys[j].Name(), srcIT.Key())
 			require.Equal(t, srcIT.Key(), dstIT.Key(), i)
 			require.Equal(t, srcIT.Value(), dstIT.Value(), "[%s] element (%d): %s", srcStoreKeys[j].Name(), i, srcIT.Key())
 			srcIT.Next()

diff --git a/x/wasm/internal/keeper/keeper.go b/x/wasm/internal/keeper/keeper.go
@@ -459,7 +459,7 @@ func (k Keeper) setContractInfo(ctx sdk.Context, contractAddress sdk.AccAddress,
 	store.Set(types.GetContractAddressKey(contractAddress), k.cdc.MustMarshalBinaryBare(contract))
 }
 
-func (k Keeper) ListContractInfo(ctx sdk.Context, cb func(sdk.AccAddress, types.ContractInfo) bool) {
+func (k Keeper) IterateContractInfo(ctx sdk.Context, cb func(sdk.AccAddress, types.ContractInfo) bool) {
 	prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.ContractKeyPrefix)
 	iter := prefixStore.Iterator(nil, nil)
 	for ; iter.Valid(); iter.Next() {
@@ -509,6 +509,19 @@ func (k Keeper) containsCodeInfo(ctx sdk.Context, codeID uint64) bool {
 	return store.Has(types.GetCodeKey(codeID))
 }
 
+func (k Keeper) IterateCodeInfos(ctx sdk.Context, cb func(uint64, types.CodeInfo) bool) {
+	prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), types.CodeKeyPrefix)
+	iter := prefixStore.Iterator(nil, nil)
+	for ; iter.Valid(); iter.Next() {
+		var c types.CodeInfo
+		k.cdc.MustUnmarshalBinaryBare(iter.Value(), &c)
+		// cb returns true to stop early
+		if cb(binary.BigEndian.Uint64(iter.Key()), c) {
+			return
+		}
+	}
+}
+
 func (k Keeper) GetByteCode(ctx sdk.Context, codeID uint64) ([]byte, error) {
 	store := ctx.KVStore(k.storeKey)
 	var codeInfo types.CodeInfo

diff --git a/x/wasm/internal/keeper/proposal_integration_test.go b/x/wasm/internal/keeper/proposal_integration_test.go
@@ -106,7 +106,7 @@ func TestInstantiateProposal(t *testing.T) {
 	assert.Equal(t, "testing", cInfo.Label)
 	expHistory := []types.ContractCodeHistoryEntry{{
 		Operation: types.InitContractCodeHistoryType,
-		CodeID:    src.Code,
+		CodeID:    src.CodeID,
 		Updated:   types.NewAbsoluteTxPosition(ctx),
 		Msg:       src.InitMsg,
 	}}
@@ -184,7 +184,7 @@ func TestMigrateProposal(t *testing.T) {
 		Updated:   types.NewAbsoluteTxPosition(ctx),
 	}, {
 		Operation: types.MigrateContractCodeHistoryType,
-		CodeID:    src.Code,
+		CodeID:    src.CodeID,
 		Updated:   types.NewAbsoluteTxPosition(ctx),
 		Msg:       src.MigrateMsg,
 	}}

diff --git a/x/wasm/internal/keeper/querier.go b/x/wasm/internal/keeper/querier.go
@@ -97,7 +97,7 @@ func queryContractListByCode(ctx sdk.Context, codeIDstr string, req abci.Request
 	}
 
 	var contracts []ContractInfoWithAddress
-	keeper.ListContractInfo(ctx, func(addr sdk.AccAddress, info types.ContractInfo) bool {
+	keeper.IterateContractInfo(ctx, func(addr sdk.AccAddress, info types.ContractInfo) bool {
 		if info.CodeID == codeID {
 			// and add the address
 			infoWithAddress := ContractInfoWithAddress{