Add balance support to flatKV

sei-db/common/keys/evm.go

@@ -25,6 +25,7 @@ var (
 	codeKeyPrefix     = []byte{0x07}
 	codeHashKeyPrefix = []byte{0x08}
 	nonceKeyPrefix    = []byte{0x0a}
+	balanceKeyPrefix  = []byte{0x21}
 )
 
 // StateKeyPrefix returns the storage state key prefix (0x03).
@@ -34,18 +35,21 @@ func StateKeyPrefix() []byte { return stateKeyPrefix }
 // EVMKeyKind identifies an EVM key family.
 type EVMKeyKind uint8
 
+// These values are in-memory routing tags, renumbered whenever a kind is added. Writing one into a
+// key, a value, or any other stored or wire format is forbidden.
 const (
 	EVMKeyEmpty    EVMKeyKind = iota // Returned only for zero-length keys
 	EVMKeyNonce                      // Stripped key: 20-byte address
 	EVMKeyCodeHash                   // Stripped key: 20-byte address
+	EVMKeyBalance                    // Stripped key: 20-byte address
 	EVMKeyCode                       // Stripped key: 20-byte address
 	EVMKeyStorage                    // Stripped key: addr||slot (20+32 bytes)
 	EVMKeyMisc                       // Full original key preserved (address mappings, codesize, etc.)
 )
 
 // ParseEVMKey parses an EVM key from the x/evm store keyspace.
 //
-// For optimized keys (nonce, code, codehash, storage), keyBytes is the stripped key.
+// For optimized keys (nonce, code, codehash, storage, balance), keyBytes is the stripped key.
 // For misc keys (all other EVM data including codesize), keyBytes is the full original key.
 // Only returns EVMKeyEmpty for zero-length keys.
 func ParseEVMKey(key []byte) (kind EVMKeyKind, keyBytes []byte) {
@@ -77,6 +81,12 @@ func ParseEVMKey(key []byte) (kind EVMKeyKind, keyBytes []byte) {
 			return EVMKeyMisc, key
 		}
 		return EVMKeyStorage, key[len(stateKeyPrefix):]
+
+	case bytes.HasPrefix(key, balanceKeyPrefix):
+		if len(key) != len(balanceKeyPrefix)+AddressLen {
+			return EVMKeyMisc, key
+		}
+		return EVMKeyBalance, key[len(balanceKeyPrefix):]
 	}
 
 	// All other EVM keys go to the misc store (address mappings, codesize, etc.)
@@ -95,6 +105,8 @@ func EVMKeyPrefixByte(kind EVMKeyKind) (byte, bool) {
 		return codeHashKeyPrefix[0], true
 	case EVMKeyCode:
 		return codeKeyPrefix[0], true
+	case EVMKeyBalance:
+		return balanceKeyPrefix[0], true
 	default:
 		return 0, false
 	}
@@ -123,7 +135,7 @@ func InternalKeyLen(kind EVMKeyKind) int {
 	switch kind {
 	case EVMKeyStorage:
 		return AddressLen + slotLen // 52 bytes
-	case EVMKeyNonce, EVMKeyCodeHash, EVMKeyCode:
+	case EVMKeyNonce, EVMKeyCodeHash, EVMKeyCode, EVMKeyBalance:
 		return AddressLen // 20 bytes
 	default:
 		return 0

sei-db/common/keys/evm_test.go

@@ -64,6 +64,12 @@ func TestParseEVMKey(t *testing.T) {
 			wantKind:  EVMKeyStorage,
 			wantBytes: concat(addr, slot),
 		},
+		{
+			name:      "Balance",
+			key:       concat(balanceKeyPrefix, addr),
+			wantKind:  EVMKeyBalance,
+			wantBytes: addr,
+		},
 		// Legacy keys - keep full key (address mappings, unknown prefix, malformed, etc.)
 		{
 			name:      "EVMAddressToSeiAddress goes to Legacy",
@@ -119,6 +125,18 @@ func TestParseEVMKey(t *testing.T) {
 			wantKind:  EVMKeyMisc,
 			wantBytes: concat(concat(concat(stateKeyPrefix, addr), slot), []byte{0x00}),
 		},
+		{
+			name:      "BalanceTooShort goes to Legacy",
+			key:       balanceKeyPrefix,
+			wantKind:  EVMKeyMisc,
+			wantBytes: balanceKeyPrefix,
+		},
+		{
+			name:      "BalanceWrongLenLong goes to Legacy",
+			key:       concat(balanceKeyPrefix, concat(addr, []byte{0x00})),
+			wantKind:  EVMKeyMisc,
+			wantBytes: concat(balanceKeyPrefix, concat(addr, []byte{0x00})),
+		},
 	}
 
 	for _, tc := range tests {
@@ -177,6 +195,12 @@ func TestBuildMemIAVLEVMKey(t *testing.T) {
 			keyBytes: concat(addr, slot),
 			want:     concat(stateKeyPrefix, concat(addr, slot)),
 		},
+		{
+			name:     "Balance",
+			kind:     EVMKeyBalance,
+			keyBytes: addr,
+			want:     concat(balanceKeyPrefix, addr),
+		},
 	}
 
 	for _, tc := range tests {
@@ -192,4 +216,18 @@ func TestInternalKeyLen(t *testing.T) {
 	require.Equal(t, AddressLen, InternalKeyLen(EVMKeyNonce))
 	require.Equal(t, AddressLen, InternalKeyLen(EVMKeyCodeHash))
 	require.Equal(t, AddressLen, InternalKeyLen(EVMKeyCode))
+	require.Equal(t, AddressLen, InternalKeyLen(EVMKeyBalance))
+}
+
+// The prefix bytes this package mirrors from x/evm/types must stay distinct: ParseEVMKey classifies on
+// the first byte alone, so two families sharing one would silently reparse each other's rows.
+func TestEVMKeyPrefixesAreDistinct(t *testing.T) {
+	seen := map[byte]EVMKeyKind{}
+	for _, kind := range []EVMKeyKind{EVMKeyNonce, EVMKeyCodeHash, EVMKeyCode, EVMKeyStorage, EVMKeyBalance} {
+		prefix, ok := EVMKeyPrefixByte(kind)
+		require.True(t, ok, "kind %v has no prefix byte", kind)
+		previous, duplicate := seen[prefix]
+		require.False(t, duplicate, "kinds %v and %v both use prefix 0x%02x", previous, kind, prefix)
+		seen[prefix] = kind
+	}
 }

sei-db/state_db/sc/composite/random_test_framework_test.go

@@ -244,6 +244,10 @@ func randomCodeHashValue(rng *testutil.TestRandom) []byte {
 	return ensureNonZero(randomTestBytes(rng, vtype.CodeHashLen))
 }
 
+func randomBalanceValue(rng *testutil.TestRandom) []byte {
+	return ensureNonZero(randomTestBytes(rng, vtype.BalanceLen))
+}
+
 func randomStorageValue(rng *testutil.TestRandom) []byte {
 	return ensureNonZero(randomTestBytes(rng, vtype.SlotLen))
 }
@@ -284,16 +288,17 @@ func randomLegacyEVMKey(rng *testutil.TestRandom) []byte {
 //   - storage: one storageDB row  (0x03 || addr || slot)
 //   - code:    one codeDB row      (0x07 || addr)
 //   - account: one accountDB row   (0x0a || addr) — a nonce is always written
-//     (a zero nonce reads back as absent), and with ~50% probability a code
-//     hash is written for the SAME address. That second case is the account-map
-//     "collision": the nonce and code hash merge into one physical account row,
-//     exercising the merged-account read / iterate / migrate paths.
+//     (a zero nonce reads back as absent), and with ~50% probability each, a code
+//     hash and a balance are written for the SAME address. Those cases are the
+//     account-map "collision": the nonce, code hash and balance merge into one
+//     physical account row, exercising the merged-account read / iterate /
+//     migrate paths.
 //   - legacy:  one legacyDB row     (0x01 || suffix) — a non-optimized EVM key
 //     (address mappings, codesize, etc.) with a variable-length, occasionally
 //     empty value, populating flatkv's EVM legacy lane.
 //
 // Returning a slice (rather than one pair) is what lets a single logical
-// account own both a nonce and a code hash within one block.
+// account own a nonce, a code hash and a balance within one block.
 func newRandomEVMEntry(rng *testutil.TestRandom) []*proto.KVPair {
 	switch rng.Intn(4) {
 	case 0:
@@ -307,6 +312,12 @@ func newRandomEVMEntry(rng *testutil.TestRandom) []*proto.KVPair {
 				Value: randomCodeHashValue(rng),
 			})
 		}
+		if rng.Intn(2) == 0 {
+			pairs = append(pairs, &proto.KVPair{
+				Key:   keys.BuildEVMKey(keys.EVMKeyBalance, addr),
+				Value: randomBalanceValue(rng),
+			})
+		}
 		return pairs
 	case 1:
 		addr := randomTestBytes(rng, keys.AddressLen)
@@ -329,6 +340,8 @@ func freshEVMValue(rng *testutil.TestRandom, key []byte) []byte {
 		return randomNonceValue(rng)
 	case keys.EVMKeyCodeHash:
 		return randomCodeHashValue(rng)
+	case keys.EVMKeyBalance:
+		return randomBalanceValue(rng)
 	case keys.EVMKeyCode:
 		return randomCodeValue(rng)
 	case keys.EVMKeyStorage:
@@ -494,12 +507,11 @@ func simulateBlocks(
 				&proto.KVPair{Key: []byte(kp.key), Value: value})
 		}
 
-		// Delete existing keys. Deleting an account's nonce must also delete
-		// that address's code hash in the same block: flatkv merges both into a
-		// single physical account row, so dropping only the nonce would leave a
-		// live (now nonce-zero) row whose nonce reads back via the phantom-nonce
-		// path — present in flatkv but absent from the oracle. Expanding the
-		// delete set deterministically (in sample order, deduped) keeps the
+		// Delete existing keys. Deleting an account's nonce must also delete that address's other
+		// merged fields — code hash and balance — in the same block: flatkv merges all three into a
+		// single physical account row, so dropping only the nonce would leave a live (now nonce-zero)
+		// row whose nonce reads back via the phantom-nonce path — present in flatkv but absent from
+		// the oracle. Expanding the delete set deterministically (in sample order, deduped) keeps the
 		// generated changeset byte-identical for a given seed.
 		toDelete := make([]keyPair, 0, p.deletesPerBlock)
 		inDelete := make(map[keyPair]struct{}, p.deletesPerBlock)
@@ -515,8 +527,13 @@ func simulateBlocks(
 			if kp.store != keys.EVMStoreKey {
 				continue
 			}
-			if kind, stripped := keys.ParseEVMKey([]byte(kp.key)); kind == keys.EVMKeyNonce {
-				sibling := keyPair{store: kp.store, key: string(keys.BuildEVMKey(keys.EVMKeyCodeHash, stripped))}
+			kind, stripped := keys.ParseEVMKey([]byte(kp.key))
+			if kind != keys.EVMKeyNonce {
+				continue
+			}
+			for _, siblingKind := range []keys.EVMKeyKind{keys.EVMKeyCodeHash, keys.EVMKeyBalance} {
+				siblingKey := keys.BuildEVMKey(siblingKind, stripped)
+				sibling := keyPair{store: kp.store, key: string(siblingKey)}
 				if keysInUse.Contains(sibling) {
 					addDelete(sibling)
 				}
@@ -1002,21 +1019,23 @@ type flatKVExpectedRow struct {
 	storageValue [32]byte // rowStorage
 	nonce        uint64   // rowAccount
 	codeHash     [32]byte // rowAccount
+	balance      [32]byte // rowAccount
 	code         []byte   // rowCode
 	legacyValue  []byte   // rowLegacy
 }
 
 // oracleToFlatKVRows projects the oracle into the physical row layout flatkv
 // uses internally, keyed by the physical (module-prefixed) key. Only stores the
-// placement model routes to flatkv are included. The EVM nonce and code hash
-// for a single address are merged into one account row, exactly as flatkv's
-// accountDB stores them — this is what makes the row-by-row check sensitive to
-// the account-merge logic. Valid only for steady-state placement.
+// placement model routes to flatkv are included. The EVM nonce, code hash and
+// balance for a single address are merged into one account row, exactly as
+// flatkv's accountDB stores them — this is what makes the row-by-row check
+// sensitive to the account-merge logic. Valid only for steady-state placement.
 func oracleToFlatKVRows(
 	oracle *storeOracle, placement func(store string) backendPlacement) map[string]flatKVExpectedRow {
 	type acct struct {
 		nonce    uint64
 		codeHash [32]byte
+		balance  [32]byte
 	}
 	accounts := map[string]*acct{}
 	getAcct := func(addr string) *acct {
@@ -1057,6 +1076,8 @@ func oracleToFlatKVRows(
 				getAcct(string(stripped)).nonce = binary.BigEndian.Uint64(v)
 			case keys.EVMKeyCodeHash:
 				copy(getAcct(string(stripped)).codeHash[:], v)
+			case keys.EVMKeyBalance:
+				copy(getAcct(string(stripped)).balance[:], v)
 			default: // EVMKeyMisc: identity-mapped under the "evm/" prefix
 				rows[string(ktype.ModulePhysicalKey(keys.EVMStoreKey, []byte(k)))] =
 					flatKVExpectedRow{kind: rowLegacy, legacyValue: append([]byte(nil), v...)}
@@ -1066,7 +1087,7 @@ func oracleToFlatKVRows(
 
 	for addr, a := range accounts {
 		rows[string(ktype.EVMPhysicalKey(ktype.EVMKeyAccount, []byte(addr)))] =
-			flatKVExpectedRow{kind: rowAccount, nonce: a.nonce, codeHash: a.codeHash}
+			flatKVExpectedRow{kind: rowAccount, nonce: a.nonce, codeHash: a.codeHash, balance: a.balance}
 	}
 	return rows
 }
@@ -1134,9 +1155,7 @@ func assertFlatKVRowMatches(t *testing.T, physKey, rawVal []byte, exp flatKVExpe
 		require.NoError(t, err, "decode account row %x", physKey)
 		require.Equal(t, exp.nonce, ad.GetNonce(), "account nonce mismatch for %x", physKey)
 		require.Equal(t, exp.codeHash[:], ad.GetCodeHash()[:], "account code hash mismatch for %x", physKey)
-		var zeroBalance vtype.Balance
-		require.Equal(t, zeroBalance[:], ad.GetBalance()[:],
-			"account balance must be zero (balances are not stored in flatkv yet) for %x", physKey)
+		require.Equal(t, exp.balance[:], ad.GetBalance()[:], "account balance mismatch for %x", physKey)
 	case rowLegacy:
 		ld, err := vtype.DeserializeMiscData(rawVal)
 		require.NoError(t, err, "decode legacy row %x", physKey)
@@ -1153,7 +1172,11 @@ func assertFlatKVRowMatches(t *testing.T, physKey, rawVal []byte, exp flatKVExpe
 func assertFlatKVMapsExercised(t *testing.T, oracle *storeOracle, placement func(store string) backendPlacement) {
 	t.Helper()
 	var storageRows, codeRows, legacyRows int
-	type acctFlags struct{ nonce, codeHash bool }
+	type acctFlags struct {
+		nonce    bool
+		codeHash bool
+		balance  bool
+	}
 	accounts := map[string]*acctFlags{}
 	flag := func(addr string) *acctFlags {
 		a, ok := accounts[addr]
@@ -1185,18 +1208,23 @@ func assertFlatKVMapsExercised(t *testing.T, oracle *storeOracle, placement func
 				flag(string(stripped)).nonce = true
 			case keys.EVMKeyCodeHash:
 				flag(string(stripped)).codeHash = true
+			case keys.EVMKeyBalance:
+				flag(string(stripped)).balance = true
 			default:
 				legacyRows++
 			}
 		}
 	}
 
-	var accountRows, collisions int
+	var accountRows, collisions, balanceAccounts int
 	for _, af := range accounts {
 		accountRows++
 		if af.nonce && af.codeHash {
 			collisions++
 		}
+		if af.balance {
+			balanceAccounts++
+		}
 	}
 
 	evmPlacement := placement(keys.EVMStoreKey)
@@ -1206,6 +1234,7 @@ func assertFlatKVMapsExercised(t *testing.T, oracle *storeOracle, placement func
 		require.Positive(t, accountRows, "expected account-map rows in flatkv")
 		require.Positive(t, collisions,
 			"expected at least one account with both a nonce and a code hash (account-map collision)")
+		require.Positive(t, balanceAccounts, "expected at least one account holding a balance")
 	}
 
 	legacyExpected := false