fix ss import for raw flatkv modules

sei-db/common/evm/keys.go

@@ -18,6 +18,13 @@ const EVMStoreKey = "evm"
 // layers need to recognise this name and treat it as EVM data.
 const EVMFlatKVStoreKey = "evm_flatkv"
 
+const (
+	EVMFlatKVAccountStoreKey = "evm_flatkv_account"
+	EVMFlatKVCodeStoreKey    = "evm_flatkv_code"
+	EVMFlatKVStorageStoreKey = "evm_flatkv_storage"
+	EVMFlatKVLegacyStoreKey  = "evm_flatkv_legacy"
+)
+
 // EVM key prefixes — mirrored from x/evm/types/keys.go.
 // These are immutable on-disk format markers; changing them would break
 // all existing state, so duplicating here is safe and avoids pulling in the
@@ -135,3 +142,23 @@ func InternalKeyLen(kind EVMKeyKind) int {
 		return 0
 	}
 }
+
+func FlatKVStoreKeyKind(storeKey string) (EVMKeyKind, bool) {
+	switch storeKey {
+	case EVMFlatKVAccountStoreKey:
+		return EVMKeyNonce, true
+	case EVMFlatKVCodeStoreKey:
+		return EVMKeyCode, true
+	case EVMFlatKVStorageStoreKey:
+		return EVMKeyStorage, true
+	case EVMFlatKVLegacyStoreKey:
+		return EVMKeyLegacy, true
+	default:
+		return EVMKeyEmpty, false
+	}
+}
+
+func IsFlatKVStoreKey(storeKey string) bool {
+	_, ok := FlatKVStoreKeyKind(storeKey)
+	return ok
+}

sei-db/state_db/ss/composite/store.go

@@ -1,6 +1,7 @@
 package composite
 
 import (
+	"encoding/binary"
 	"fmt"
 	"path/filepath"
 	"sync"
@@ -10,6 +11,7 @@ import (
 	"github.com/sei-protocol/sei-chain/sei-db/config"
 	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
 	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/flatkv/vtype"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/ss/backend"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/ss/cosmos"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/ss/evm"
@@ -258,32 +260,101 @@ func stripEVMFromChangesets(changesets []*proto.NamedChangeSet) []*proto.NamedCh
 	return stripped
 }
 
-func normalizeSnapshotNode(node types.SnapshotNode) types.SnapshotNode {
+func normalizeSnapshotNodes(node types.SnapshotNode) ([]types.SnapshotNode, error) {
 	if node.StoreKey == commonevm.EVMFlatKVStoreKey {
 		node.StoreKey = evm.EVMStoreKey
+		return []types.SnapshotNode{node}, nil
+	}
+	if !commonevm.IsFlatKVStoreKey(node.StoreKey) {
+		return []types.SnapshotNode{node}, nil
+	}
+
+	switch node.StoreKey {
+	case commonevm.EVMFlatKVAccountStoreKey:
+		return normalizeFlatKVAccountNode(node)
+	case commonevm.EVMFlatKVCodeStoreKey:
+		return normalizeFlatKVCodeNode(node)
+	case commonevm.EVMFlatKVStorageStoreKey:
+		return normalizeFlatKVStorageNode(node)
+	case commonevm.EVMFlatKVLegacyStoreKey:
+		return normalizeFlatKVLegacyNode(node)
+	default:
+		return nil, fmt.Errorf("unsupported flatkv snapshot module: %s", node.StoreKey)
 	}
-	return node
 }
 
-func (s *CompositeStateStore) Import(version int64, ch <-chan types.SnapshotNode) error {
-	if s.evmStore == nil || s.config.WriteMode == config.CosmosOnlyWrite {
-		// Normalize evm_flatkv → evm so downstream routing and storage work
-		// correctly regardless of whether the snapshot was exported with the
-		// FlatKV module or only the legacy evm module.
-		normalized := make(chan types.SnapshotNode, cap(ch))
-		go func() {
-			defer close(normalized)
-			for node := range ch {
-				normalized <- normalizeSnapshotNode(node)
-			}
-		}()
-		return s.cosmosStore.Import(version, normalized)
+func normalizeFlatKVAccountNode(node types.SnapshotNode) ([]types.SnapshotNode, error) {
+	accountData, err := vtype.DeserializeAccountData(node.Value)
+	if err != nil {
+		return nil, fmt.Errorf("deserialize flatkv account value: %w", err)
+	}
+
+	nonceValue := make([]byte, vtype.NonceLen)
+	binary.BigEndian.PutUint64(nonceValue, accountData.GetNonce())
+	nodes := []types.SnapshotNode{{
+		StoreKey: evm.EVMStoreKey,
+		Key:      commonevm.BuildMemIAVLEVMKey(commonevm.EVMKeyNonce, node.Key),
+		Value:    nonceValue,
+	}}
+
+	codeHash := accountData.GetCodeHash()
+	var zeroHash vtype.CodeHash
+	if codeHash != nil && *codeHash != zeroHash {
+		nodes = append(nodes, types.SnapshotNode{
+			StoreKey: evm.EVMStoreKey,
+			Key:      commonevm.BuildMemIAVLEVMKey(commonevm.EVMKeyCodeHash, node.Key),
+			Value:    append([]byte(nil), codeHash[:]...),
+		})
 	}
 
+	return nodes, nil
+}
+
+func normalizeFlatKVCodeNode(node types.SnapshotNode) ([]types.SnapshotNode, error) {
+	codeData, err := vtype.DeserializeCodeData(node.Value)
+	if err != nil {
+		return nil, fmt.Errorf("deserialize flatkv code value: %w", err)
+	}
+	return []types.SnapshotNode{{
+		StoreKey: evm.EVMStoreKey,
+		Key:      commonevm.BuildMemIAVLEVMKey(commonevm.EVMKeyCode, node.Key),
+		Value:    append([]byte(nil), codeData.GetBytecode()...),
+	}}, nil
+}
+
+func normalizeFlatKVStorageNode(node types.SnapshotNode) ([]types.SnapshotNode, error) {
+	storageData, err := vtype.DeserializeStorageData(node.Value)
+	if err != nil {
+		return nil, fmt.Errorf("deserialize flatkv storage value: %w", err)
+	}
+	return []types.SnapshotNode{{
+		StoreKey: evm.EVMStoreKey,
+		Key:      commonevm.BuildMemIAVLEVMKey(commonevm.EVMKeyStorage, node.Key),
+		Value:    append([]byte(nil), storageData.GetValue()[:]...),
+	}}, nil
+}
+
+func normalizeFlatKVLegacyNode(node types.SnapshotNode) ([]types.SnapshotNode, error) {
+	legacyData, err := vtype.DeserializeLegacyData(node.Value)
+	if err != nil {
+		return nil, fmt.Errorf("deserialize flatkv legacy value: %w", err)
+	}
+	return []types.SnapshotNode{{
+		StoreKey: evm.EVMStoreKey,
+		Key:      append([]byte(nil), node.Key...),
+		Value:    legacyData.GetValue(),
+	}}, nil
+}
+
+func (s *CompositeStateStore) Import(version int64, ch <-chan types.SnapshotNode) error {
 	splitWrite := s.config.WriteMode == config.SplitWrite
+	importToEVM := s.evmStore != nil && s.config.WriteMode != config.CosmosOnlyWrite
 
 	cosmosCh := make(chan types.SnapshotNode, 100)
-	evmCh := make(chan types.SnapshotNode, 100)
+	var evmCh chan types.SnapshotNode
+	if importToEVM {
+		evmCh = make(chan types.SnapshotNode, 100)
+	}
 	importErrCh := make(chan error, 2)
 
 	var wg sync.WaitGroup
@@ -292,7 +363,9 @@ func (s *CompositeStateStore) Import(version int64, ch <-chan types.SnapshotNode
 	closeImportChans := func() {
 		closeOnce.Do(func() {
 			close(cosmosCh)
-			close(evmCh)
+			if evmCh != nil {
+				close(evmCh)
+			}
 		})
 	}
 
@@ -304,13 +377,15 @@ func (s *CompositeStateStore) Import(version int64, ch <-chan types.SnapshotNode
 		}
 	}()
 
-	wg.Add(1)
-	go func() {
-		defer wg.Done()
-		if err := s.evmStore.Import(version, evmCh); err != nil {
-			importErrCh <- err
-		}
-	}()
+	if importToEVM {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			if err := s.evmStore.Import(version, evmCh); err != nil {
+				importErrCh <- err
+			}
+		}()
+	}
 
 	var importErr error
 	drainImportErr := func() {
@@ -345,21 +420,27 @@ func (s *CompositeStateStore) Import(version int64, ch <-chan types.SnapshotNode
 	}
 
 	for node := range ch {
-		node = normalizeSnapshotNode(node)
+		normalizedNodes, err := normalizeSnapshotNodes(node)
+		if err != nil && importErr == nil {
+			importErr = err
+			closeImportChans()
+		}
 		drainImportErr()
 		if importErr != nil {
 			continue
 		}
 
-		isEVM := node.StoreKey == evm.EVMStoreKey
-		if !isEVM || !splitWrite {
-			if err := sendNode(cosmosCh, node); err != nil {
-				continue
+		for _, normalizedNode := range normalizedNodes {
+			isEVM := normalizedNode.StoreKey == evm.EVMStoreKey
+			if !isEVM || !splitWrite || !importToEVM {
+				if err := sendNode(cosmosCh, normalizedNode); err != nil {
+					continue
+				}
 			}
-		}
-		if isEVM {
-			if err := sendNode(evmCh, node); err != nil {
-				continue
+			if isEVM && importToEVM {
+				if err := sendNode(evmCh, normalizedNode); err != nil {
+					continue
+				}
 			}
 		}
 	}

sei-db/state_db/ss/composite/store_test.go

@@ -12,6 +12,7 @@ import (
 	"github.com/sei-protocol/sei-chain/sei-db/config"
 	"github.com/sei-protocol/sei-chain/sei-db/db_engine/types"
 	"github.com/sei-protocol/sei-chain/sei-db/proto"
+	"github.com/sei-protocol/sei-chain/sei-db/state_db/sc/flatkv/vtype"
 	"github.com/sei-protocol/sei-chain/sei-db/state_db/ss/evm"
 	"github.com/stretchr/testify/require"
 )
@@ -1493,6 +1494,20 @@ func feedNodes(ch chan<- types.SnapshotNode, nodes []types.SnapshotNode) {
 	close(ch)
 }
 
+func mustSSCodeHash(t *testing.T, value []byte) *vtype.CodeHash {
+	t.Helper()
+	parsed, err := vtype.ParseCodeHash(value)
+	require.NoError(t, err)
+	return parsed
+}
+
+func mustSSStorageValue(t *testing.T, value []byte) *[32]byte {
+	t.Helper()
+	parsed, err := vtype.ParseStorageValue(value)
+	require.NoError(t, err)
+	return parsed
+}
+
 func TestImport_OnlyEvmModule(t *testing.T) {
 	for _, mode := range []config.WriteMode{config.DualWrite, config.SplitWrite, config.CosmosOnlyWrite} {
 		t.Run("WriteMode="+string(mode), func(t *testing.T) {
@@ -1579,6 +1594,112 @@ func TestImport_OnlyEvmFlatkvModule(t *testing.T) {
 	}
 }
 
+func TestImport_RawFlatKVModules(t *testing.T) {
+	for _, mode := range []config.WriteMode{config.DualWrite, config.SplitWrite, config.CosmosOnlyWrite} {
+		t.Run("WriteMode="+string(mode), func(t *testing.T) {
+			store, cleanup := setupImportTestStore(t, mode)
+			defer cleanup()
+
+			addr := make([]byte, 20)
+			addr[0] = 0xAA
+			slot := make([]byte, 32)
+			slot[0] = 0xBB
+
+			nonceValue := []byte{0, 0, 0, 0, 0, 0, 0, 7}
+			codeHashValue := make([]byte, vtype.CodeHashLen)
+			codeHashValue[31] = 0xCC
+			codeValue := []byte{0x60, 0x80, 0x52}
+			storageValue := make([]byte, 32)
+			storageValue[31] = 0xDD
+			legacyKey := append([]byte{0x09}, addr...)
+			legacyValue := []byte{0x00, 0x20}
+
+			accountData := vtype.NewAccountData().
+				SetBlockHeight(11).
+				SetNonce(7).
+				SetCodeHash(mustSSCodeHash(t, codeHashValue))
+			codeData := vtype.NewCodeData().
+				SetBlockHeight(11).
+				SetBytecode(codeValue)
+			storageData := vtype.NewStorageData().
+				SetBlockHeight(11).
+				SetValue(mustSSStorageValue(t, storageValue))
+			legacyData := vtype.NewLegacyData().
+				SetValue(legacyValue)
+
+			ch := make(chan types.SnapshotNode, 10)
+			nodes := []types.SnapshotNode{
+				{StoreKey: "bank", Key: []byte("supply"), Value: []byte("7000")},
+				{StoreKey: commonevm.EVMFlatKVAccountStoreKey, Key: addr, Value: accountData.Serialize()},
+				{StoreKey: commonevm.EVMFlatKVCodeStoreKey, Key: addr, Value: codeData.Serialize()},
+				{StoreKey: commonevm.EVMFlatKVStorageStoreKey, Key: append(append([]byte(nil), addr...), slot...), Value: storageData.Serialize()},
+				{StoreKey: commonevm.EVMFlatKVLegacyStoreKey, Key: legacyKey, Value: legacyData.Serialize()},
+			}
+			go feedNodes(ch, nodes)
+
+			err := store.Import(1, ch)
+			require.NoError(t, err)
+
+			bankVal, err := store.cosmosStore.Get("bank", 1, []byte("supply"))
+			require.NoError(t, err)
+			require.Equal(t, []byte("7000"), bankVal)
+
+			nonceKey := append([]byte{0x0a}, addr...)
+			codeHashKey := append([]byte{0x08}, addr...)
+			codeKey := append([]byte{0x07}, addr...)
+			storageKey := append([]byte{0x03}, append(append([]byte(nil), addr...), slot...)...)
+
+			if mode == config.SplitWrite {
+				cosmosNonce, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, nonceKey)
+				require.NoError(t, err)
+				require.Nil(t, cosmosNonce)
+			} else {
+				cosmosNonce, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, nonceKey)
+				require.NoError(t, err)
+				require.Equal(t, nonceValue, cosmosNonce)
+
+				cosmosCodeHash, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, codeHashKey)
+				require.NoError(t, err)
+				require.Equal(t, codeHashValue, cosmosCodeHash)
+
+				cosmosCode, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, codeKey)
+				require.NoError(t, err)
+				require.Equal(t, codeValue, cosmosCode)
+
+				cosmosStorage, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, storageKey)
+				require.NoError(t, err)
+				require.Equal(t, storageValue, cosmosStorage)
+
+				cosmosLegacy, err := store.cosmosStore.Get(evm.EVMStoreKey, 1, legacyKey)
+				require.NoError(t, err)
+				require.Equal(t, legacyValue, cosmosLegacy)
+			}
+
+			if store.evmStore != nil && mode != config.CosmosOnlyWrite {
+				evmNonce, err := store.evmStore.Get(evm.EVMStoreKey, 1, nonceKey)
+				require.NoError(t, err)
+				require.Equal(t, nonceValue, evmNonce)
+
+				evmCodeHash, err := store.evmStore.Get(evm.EVMStoreKey, 1, codeHashKey)
+				require.NoError(t, err)
+				require.Equal(t, codeHashValue, evmCodeHash)
+
+				evmCode, err := store.evmStore.Get(evm.EVMStoreKey, 1, codeKey)
+				require.NoError(t, err)
+				require.Equal(t, codeValue, evmCode)
+
+				evmStorage, err := store.evmStore.Get(evm.EVMStoreKey, 1, storageKey)
+				require.NoError(t, err)
+				require.Equal(t, storageValue, evmStorage)
+
+				evmLegacy, err := store.evmStore.Get(evm.EVMStoreKey, 1, legacyKey)
+				require.NoError(t, err)
+				require.Equal(t, legacyValue, evmLegacy)
+			}
+		})
+	}
+}
+
 func TestImport_BothEvmAndEvmFlatkv(t *testing.T) {
 	for _, mode := range []config.WriteMode{config.DualWrite, config.SplitWrite} {
 		t.Run("WriteMode="+string(mode), func(t *testing.T) {