refactor nft mergration function

modules/nft/keeper/migrations.go

@@ -18,5 +18,5 @@ func NewMigrator(k Keeper) Migrator {
 
 // Migrate1to2 migrates from version 1 to 2.
 func (m Migrator) Migrate1to2(ctx sdk.Context) error {
-	return v2.Migrate(ctx, m.k.storeKey, m.k.cdc, m.k.Logger(ctx), m.k)
+	return v2.Migrate(ctx, m.k.storeKey, m.k.cdc, m.k.Logger(ctx), m.k.SaveDenom)
 }

modules/nft/migrations/v2/expected_keeper.go

@@ -2,26 +2,14 @@ package v2
 
 import sdk "github.com/cosmos/cosmos-sdk/types"
 
-type NFTKeeper interface {
-	SaveDenom(ctx sdk.Context, id,
-		name,
-		schema,
-		symbol string,
-		creator sdk.AccAddress,
-		mintRestricted,
-		updateRestricted bool,
-		description,
-		uri,
-		uriHash,
-		data string,
-	) error
-
-	SaveNFT(ctx sdk.Context, denomID,
-		tokenID,
-		tokenNm,
-		tokenURI,
-		tokenUriHash,
-		tokenData string,
-		receiver sdk.AccAddress,
-	) error
-}
+// SaveDenom save the denom of class
+type SaveDenom func(ctx sdk.Context, id, name, schema,
+	symbol string,
+	creator sdk.AccAddress,
+	mintRestricted,
+	updateRestricted bool,
+	description,
+	uri,
+	uriHash,
+	data string,
+) error

modules/nft/migrations/v2/keeper.go

@@ -0,0 +1,164 @@
+package v2
+
+import (
+	"reflect"
+	"unsafe"
+
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/store/prefix"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
+	"github.com/cosmos/cosmos-sdk/x/nft"
+	nftkeeper "github.com/cosmos/cosmos-sdk/x/nft/keeper"
+
+	"github.com/irisnet/irismod/modules/nft/types"
+)
+
+type keeper struct {
+	storeKey storetypes.StoreKey // Unexposed key to access store from sdk.Context
+	cdc      codec.Codec
+}
+
+func (k keeper) saveNFT(ctx sdk.Context, denomID,
+	tokenID,
+	tokenNm,
+	tokenURI,
+	tokenURIHash,
+	tokenData string,
+	receiver sdk.AccAddress,
+) error {
+	nftMetadata := &types.NFTMetadata{
+		Name: tokenNm,
+		Data: tokenData,
+	}
+	data, err := codectypes.NewAnyWithValue(nftMetadata)
+	if err != nil {
+		return err
+	}
+
+	token := nft.NFT{
+		ClassId: denomID,
+		Id:      tokenID,
+		Uri:     tokenURI,
+		UriHash: tokenURIHash,
+		Data:    data,
+	}
+	k.setNFT(ctx, token)
+	k.setOwner(ctx, token.ClassId, token.Id, receiver)
+	k.incrTotalSupply(ctx, token.ClassId)
+	return nil
+}
+
+func (k keeper) setNFT(ctx sdk.Context, token nft.NFT) {
+	nftStore := k.getNFTStore(ctx, token.ClassId)
+	bz := k.cdc.MustMarshal(&token)
+	nftStore.Set([]byte(token.Id), bz)
+}
+
+func (k keeper) setOwner(ctx sdk.Context, classID, nftID string, owner sdk.AccAddress) {
+	store := ctx.KVStore(k.storeKey)
+	store.Set(ownerStoreKey(classID, nftID), owner.Bytes())
+
+	ownerStore := k.getClassStoreByOwner(ctx, owner, classID)
+	ownerStore.Set([]byte(nftID), nftkeeper.Placeholder)
+}
+
+func (k keeper) incrTotalSupply(ctx sdk.Context, classID string) {
+	supply := k.GetTotalSupply(ctx, classID) + 1
+	k.updateTotalSupply(ctx, classID, supply)
+}
+
+// GetTotalSupply returns the number of all nfts under the specified classID
+func (k keeper) GetTotalSupply(ctx sdk.Context, classID string) uint64 {
+	store := ctx.KVStore(k.storeKey)
+	bz := store.Get(classTotalSupply(classID))
+	return sdk.BigEndianToUint64(bz)
+}
+
+func (k keeper) updateTotalSupply(ctx sdk.Context, classID string, supply uint64) {
+	store := ctx.KVStore(k.storeKey)
+	supplyKey := classTotalSupply(classID)
+	store.Set(supplyKey, sdk.Uint64ToBigEndian(supply))
+}
+
+func (k keeper) getClassStoreByOwner(ctx sdk.Context, owner sdk.AccAddress, classID string) prefix.Store {
+	store := ctx.KVStore(k.storeKey)
+	key := nftOfClassByOwnerStoreKey(owner, classID)
+	return prefix.NewStore(store, key)
+}
+
+func (k keeper) getNFTStore(ctx sdk.Context, classID string) prefix.Store {
+	store := ctx.KVStore(k.storeKey)
+	return prefix.NewStore(store, nftStoreKey(classID))
+}
+
+// classTotalSupply returns the byte representation of the ClassTotalSupply
+func classTotalSupply(classID string) []byte {
+	key := make([]byte, len(nftkeeper.ClassTotalSupply)+len(classID))
+	copy(key, nftkeeper.ClassTotalSupply)
+	copy(key[len(nftkeeper.ClassTotalSupply):], classID)
+	return key
+}
+
+// nftStoreKey returns the byte representation of the nft
+func nftStoreKey(classID string) []byte {
+	key := make([]byte, len(nftkeeper.NFTKey)+len(classID)+len(nftkeeper.Delimiter))
+	copy(key, nftkeeper.NFTKey)
+	copy(key[len(nftkeeper.NFTKey):], classID)
+	copy(key[len(nftkeeper.NFTKey)+len(classID):], nftkeeper.Delimiter)
+	return key
+}
+
+// ownerStoreKey returns the byte representation of the nft owner
+// Items are stored with the following key: values
+// 0x04<classID><Delimiter(1 Byte)><nftID>
+func ownerStoreKey(classID, nftID string) []byte {
+	// key is of format:
+	classIDBz := UnsafeStrToBytes(classID)
+	nftIDBz := UnsafeStrToBytes(nftID)
+
+	key := make([]byte, len(nftkeeper.OwnerKey)+len(classIDBz)+len(nftkeeper.Delimiter)+len(nftIDBz))
+	copy(key, nftkeeper.OwnerKey)
+	copy(key[len(nftkeeper.OwnerKey):], classIDBz)
+	copy(key[len(nftkeeper.OwnerKey)+len(classIDBz):], nftkeeper.Delimiter)
+	copy(key[len(nftkeeper.OwnerKey)+len(classIDBz)+len(nftkeeper.Delimiter):], nftIDBz)
+	return key
+}
+
+// nftOfClassByOwnerStoreKey returns the byte representation of the nft owner
+// Items are stored with the following key: values
+// 0x03<owner><Delimiter(1 Byte)><classID><Delimiter(1 Byte)>
+func nftOfClassByOwnerStoreKey(owner sdk.AccAddress, classID string) []byte {
+	owner = address.MustLengthPrefix(owner)
+	classIDBz := UnsafeStrToBytes(classID)
+
+	key := make([]byte, len(nftkeeper.NFTOfClassByOwnerKey)+len(owner)+len(nftkeeper.Delimiter)+len(classIDBz)+len(nftkeeper.Delimiter))
+	copy(key, nftkeeper.NFTOfClassByOwnerKey)
+	copy(key[len(nftkeeper.NFTOfClassByOwnerKey):], owner)
+	copy(key[len(nftkeeper.NFTOfClassByOwnerKey)+len(owner):], nftkeeper.Delimiter)
+	copy(key[len(nftkeeper.NFTOfClassByOwnerKey)+len(owner)+len(nftkeeper.Delimiter):], classIDBz)
+	copy(key[len(nftkeeper.NFTOfClassByOwnerKey)+len(owner)+len(nftkeeper.Delimiter)+len(classIDBz):], nftkeeper.Delimiter)
+	return key
+}
+
+// UnsafeStrToBytes uses unsafe to convert string into byte array. Returned bytes
+// must not be altered after this function is called as it will cause a segmentation fault.
+func UnsafeStrToBytes(s string) []byte {
+	var buf []byte
+	sHdr := (*reflect.StringHeader)(unsafe.Pointer(&s))
+	bufHdr := (*reflect.SliceHeader)(unsafe.Pointer(&buf))
+	bufHdr.Data = sHdr.Data
+	bufHdr.Cap = sHdr.Len
+	bufHdr.Len = sHdr.Len
+	return buf
+}
+
+// UnsafeBytesToStr is meant to make a zero allocation conversion
+// from []byte -> string to speed up operations, it is not meant
+// to be used generally, but for a specific pattern to delete keys
+// from a map.
+func UnsafeBytesToStr(b []byte) string {
+	return *(*string)(unsafe.Pointer(&b))
+}

modules/nft/migrations/v2/store.go

@@ -12,36 +12,29 @@ import (
 	"github.com/irisnet/irismod/modules/nft/types"
 )
 
+// Migrate is used to migrate nft data from irismod/nft to x/nft
 func Migrate(ctx sdk.Context,
 	storeKey storetypes.StoreKey,
 	cdc codec.Codec,
 	logger log.Logger,
-	k NFTKeeper,
+	saveDenom SaveDenom,
 ) error {
 	logger.Info("migrate store data from version 1 to 2")
 	startTime := time.Now()
-	denoms, err := migrateDenoms(ctx, storeKey, cdc, k)
-	if err != nil {
-		return err
-	}
-	logger.Info("migrate denoms success", "denomNum", len(denoms))
 
-	if err := migrateTokens(ctx, storeKey, cdc, logger, denoms, k); err != nil {
-		return err
-	}
-	logger.Info("migrate store data success", "consume", time.Since(startTime).String())
-	return nil
-}
-
-func migrateDenoms(ctx sdk.Context,
-	storeKey storetypes.StoreKey,
-	cdc codec.Codec,
-	k NFTKeeper,
-) (denoms []string, err error) {
 	store := ctx.KVStore(storeKey)
 	iterator := sdk.KVStorePrefixIterator(store, KeyDenom(""))
 	defer iterator.Close()
 
+	k := keeper{
+		storeKey: storeKey,
+		cdc:      cdc,
+	}
+
+	var (
+		denomNum int64
+		tokenNum int64
+	)
 	for ; iterator.Valid(); iterator.Next() {
 		var denom types.Denom
 		cdc.MustUnmarshal(iterator.Value(), &denom)
@@ -53,10 +46,10 @@ func migrateDenoms(ctx sdk.Context,
 
 		creator, err := sdk.AccAddressFromBech32(denom.Creator)
 		if err != nil {
-			return denoms, err
+			return err
 		}
 
-		if err := k.SaveDenom(ctx, denom.Id,
+		if err := saveDenom(ctx, denom.Id,
 			denom.Name,
 			denom.Schema,
 			denom.Symbol,
@@ -68,59 +61,66 @@ func migrateDenoms(ctx sdk.Context,
 			denom.UriHash,
 			denom.Data,
 		); err != nil {
-			return denoms, err
+			return err
 		}
-		denoms = append(denoms, denom.Id)
+
+		tokenInDenom, err := migrateToken(ctx, k, logger, denom.Id)
+		if err != nil {
+			return err
+		}
+		denomNum++
+		tokenNum += tokenInDenom
 
 	}
-	return denoms, nil
+	logger.Info("migrate store data success",
+		"denomTotalNum", denomNum,
+		"tokenTotalNum", tokenNum,
+		"consume", time.Since(startTime).String(),
+	)
+	return nil
 }
-
-func migrateTokens(ctx sdk.Context,
-	storeKey storetypes.StoreKey,
-	cdc codec.Codec,
+func migrateToken(
+	ctx sdk.Context,
+	k keeper,
 	logger log.Logger,
-	denoms []string,
-	k NFTKeeper,
-) error {
-	store := ctx.KVStore(storeKey)
-
+	denomID string,
+) (int64, error) {
 	var iterator sdk.Iterator
 	defer func() {
 		if iterator != nil {
 			_ = iterator.Close()
 		}
 	}()
 
+	store := ctx.KVStore(k.storeKey)
+
 	total := int64(0)
-	for _, denomID := range denoms {
-		iterator = sdk.KVStorePrefixIterator(store, KeyNFT(denomID, ""))
-		for ; iterator.Valid(); iterator.Next() {
-			var baseNFT types.BaseNFT
-			cdc.MustUnmarshal(iterator.Value(), &baseNFT)
-
-			owner, err := sdk.AccAddressFromBech32(baseNFT.Owner)
-			if err != nil {
-				return err
-			}
-
-			//delete unused key
-			store.Delete(KeyNFT(denomID, baseNFT.Id))
-			store.Delete(KeyOwner(owner, denomID, baseNFT.Id))
-
-			if err := k.SaveNFT(ctx, denomID,
-				baseNFT.Id,
-				baseNFT.Name,
-				baseNFT.URI,
-				baseNFT.UriHash,
-				baseNFT.Data,
-				owner,
-			); err != nil {
-				return err
-			}
-			total++
+	iterator = sdk.KVStorePrefixIterator(store, KeyNFT(denomID, ""))
+	for ; iterator.Valid(); iterator.Next() {
+		var baseNFT types.BaseNFT
+		k.cdc.MustUnmarshal(iterator.Value(), &baseNFT)
+
+		owner, err := sdk.AccAddressFromBech32(baseNFT.Owner)
+		if err != nil {
+			return 0, err
 		}
+
+		//delete unused key
+		store.Delete(KeyNFT(denomID, baseNFT.Id))
+		store.Delete(KeyOwner(owner, denomID, baseNFT.Id))
+
+		if err := k.saveNFT(ctx, denomID,
+			baseNFT.Id,
+			baseNFT.Name,
+			baseNFT.URI,
+			baseNFT.UriHash,
+			baseNFT.Data,
+			owner,
+		); err != nil {
+			return 0, err
+		}
+		total++
 	}
-	logger.Info("migrate nft success", "nftNum", total)
-	return nil
+	logger.Info("migrate nft success", "denomID", denomID, "nftNum", total)
+	return total, nil
 }

modules/nft/migrations/v2/store_test.go

@@ -29,7 +29,7 @@ func TestMigrate(t *testing.T) {
 	cdc := app.AppCodec()
 
 	collections := prepareData(ctx, storeKey, cdc)
-	require.NoError(t, v2.Migrate(ctx, storeKey, cdc, app.NFTKeeper.Logger(ctx), app.NFTKeeper))
+	require.NoError(t, v2.Migrate(ctx, storeKey, cdc, app.NFTKeeper.Logger(ctx), app.NFTKeeper.SaveDenom))
 	check(t, ctx, app.NFTKeeper, collections)
 
 }