From 0b875aba4ade6cc50c0fd88809abc4d0855b8e1a Mon Sep 17 00:00:00 2001
From: Dan Laine <daniel.laine@avalabs.org>
Date: Tue, 5 Dec 2023 10:27:03 -0500
Subject: [PATCH] P-Chain merkledb -- include txs in merkleized state (#2398)

---
 vms/platformvm/state/state.go              | 12 +++++-------
 vms/platformvm/state/state_helpers.go      |  7 +++++++
 vms/platformvm/state/state_helpers_test.go | 10 ++++++++++
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/vms/platformvm/state/state.go b/vms/platformvm/state/state.go
index 90be03bf1727..ed2598ce974a 100644
--- a/vms/platformvm/state/state.go
+++ b/vms/platformvm/state/state.go
@@ -63,7 +63,7 @@ var (
 	merkleSingletonPrefix   = []byte{0x01}
 	merkleBlockPrefix       = []byte{0x02}
 	merkleBlockIDsPrefix    = []byte{0x03}
-	merkleTxPrefix          = []byte{0x04}
+	merkleTxsPrefix         = []byte{0x04}
 	merkleIndexUTXOsPrefix  = []byte{0x05} // to serve UTXOIDs(addr)
 	merkleUptimesPrefix     = []byte{0x06} // locally measured uptimes
 	merkleWeightDiffPrefix  = []byte{0x07} // non-merkleized validators weight diff. TODO: should we merkleize them?
@@ -277,7 +277,6 @@ type state struct {
 	// FIND a way to reduce use of these. No use in verification of addedTxs
 	// a limited windows to support APIs
 	addedTxs map[ids.ID]*txAndStatus // map of txID -> {*txs.Tx, Status}
-	txDB     database.Database
 
 	indexedUTXOsDB database.Database
 
@@ -365,7 +364,6 @@ func newState(
 		singletonDB                   = prefixdb.New(merkleSingletonPrefix, baseDB)
 		blockDB                       = prefixdb.New(merkleBlockPrefix, baseDB)
 		blockIDsDB                    = prefixdb.New(merkleBlockIDsPrefix, baseDB)
-		txDB                          = prefixdb.New(merkleTxPrefix, baseDB)
 		indexedUTXOsDB                = prefixdb.New(merkleIndexUTXOsPrefix, baseDB)
 		localUptimesDB                = prefixdb.New(merkleUptimesPrefix, baseDB)
 		flatValidatorWeightDiffsDB    = prefixdb.New(merkleWeightDiffPrefix, baseDB)
@@ -425,7 +423,6 @@ func newState(
 		blockIDDB:     blockIDsDB,
 
 		addedTxs: make(map[ids.ID]*txAndStatus),
-		txDB:     txDB,
 
 		indexedUTXOsDB: indexedUTXOsDB,
 
@@ -640,7 +637,8 @@ func (s *state) GetTx(txID ids.ID) (*txs.Tx, status.Status, error) {
 		return tx.tx, tx.status, nil
 	}
 
-	txBytes, err := s.txDB.Get(txID[:])
+	key := merkleTxKey(txID)
+	txBytes, err := s.merkleDB.Get(key)
 	if err != nil {
 		return nil, status.Unknown, err
 	}
@@ -1189,7 +1187,6 @@ func (s *state) Close() error {
 		s.flatValidatorPublicKeyDiffsDB.Close(),
 		s.localUptimesDB.Close(),
 		s.indexedUTXOsDB.Close(),
-		s.txDB.Close(),
 		s.blockDB.Close(),
 		s.blockIDDB.Close(),
 		s.merkleDB.Close(),
@@ -1635,7 +1632,8 @@ func (s *state) writeTxs() error {
 		// Note: Evict is used rather than Put here because stx may end up
 		// referencing additional data (because of shared byte slices) that
 		// would not be properly accounted for in the cache sizing.
-		if err := s.txDB.Put(txID[:], txBytes); err != nil {
+		key := merkleTxKey(txID)
+		if err := s.merkleDB.Put(key, txBytes); err != nil {
 			return fmt.Errorf("failed to add tx: %w", err)
 		}
 	}
diff --git a/vms/platformvm/state/state_helpers.go b/vms/platformvm/state/state_helpers.go
index cf8b165e2698..91532169c663 100644
--- a/vms/platformvm/state/state_helpers.go
+++ b/vms/platformvm/state/state_helpers.go
@@ -128,3 +128,10 @@ func merkleSubnetOwnersKey(subnetID ids.ID) []byte {
 	copy(key[len(delegateeRewardsPrefix):], subnetID[:])
 	return key
 }
+
+func merkleTxKey(txID ids.ID) []byte {
+	key := make([]byte, len(merkleTxsPrefix)+ids.IDLen)
+	copy(key, merkleTxsPrefix)
+	copy(key[len(merkleTxsPrefix):], txID[:])
+	return key
+}
diff --git a/vms/platformvm/state/state_helpers_test.go b/vms/platformvm/state/state_helpers_test.go
index 00547c9c8d93..1905c8f782f9 100644
--- a/vms/platformvm/state/state_helpers_test.go
+++ b/vms/platformvm/state/state_helpers_test.go
@@ -140,3 +140,13 @@ func TestDelegateeRewardsKey(t *testing.T) {
 	require.Equal(nodeID[:], key[len(prefix):len(prefix)+len(nodeID[:])])
 	require.Equal(subnetID[:], key[len(prefix)+len(nodeID[:]):])
 }
+
+func TestMerkleTxKey(t *testing.T) {
+	require := require.New(t)
+	txID := ids.GenerateTestID()
+
+	key := merkleTxKey(txID)
+	require.Len(key, len(merkleTxsPrefix)+len(txID[:]))
+	require.Equal(merkleTxsPrefix, key[:len(merkleTxsPrefix)])
+	require.Equal(txID[:], key[len(merkleTxsPrefix):])
+}