[SMST] feat: Use compact SMST proofs

pkg/client/interface.go

@@ -64,11 +64,9 @@ type SupplierClient interface {
 		ctx context.Context,
 		claimMsgs ...MsgCreateClaim,
 	) error
-	// SubmitProof sends proof messages which contain the smt.SparseMerkleClosestProof,
+	// SubmitProof sends proof messages which contain the smt.SparseCompactMerkleClosestProof,
 	// corresponding to some previously created claim for the same session.
 	// The proof is validated on-chain as part of the pocket protocol.
-	// TODO_MAINNET(#427): Use SparseCompactClosestProof here to reduce
-	// the amount of data stored on-chain.
 	SubmitProofs(
 		ctx context.Context,
 		sessionProofs ...MsgSubmitProof,

pkg/relayer/interface.go

@@ -128,7 +128,7 @@ type SessionTree interface {
 	// ProveClosest is a wrapper for the SMST's ProveClosest function. It returns the
 	// proof for the given path.
 	// This function should be called several blocks after a session has been claimed and needs to be proven.
-	ProveClosest(path []byte) (proof *smt.SparseMerkleClosestProof, err error)
+	ProveClosest(path []byte) (proof *smt.SparseCompactMerkleClosestProof, err error)
 
 	// GetClaimRoot returns the root hash of the SMST needed for creating the claim.
 	GetClaimRoot() []byte
@@ -158,4 +158,7 @@ type SessionTree interface {
 
 	// GetSupplierOperatorAddress returns the supplier operator address building this tree.
 	GetSupplierOperatorAddress() *cosmostypes.AccAddress
+
+	// GetTrieSpec returns the trie spec of the SMST.
+	GetTrieSpec() smt.TrieSpec
 }

pkg/relayer/session/sessiontree.go

@@ -50,11 +50,11 @@ type sessionTree struct {
 	// proofPath is the path for which the proof was generated.
 	proofPath []byte
 
-	// proof is the generated proof for the session given a proofPath.
-	proof *smt.SparseMerkleClosestProof
+	// compactProof is the generated compactProof for the session given a proofPath.
+	compactProof *smt.SparseCompactMerkleClosestProof
 
-	// proofBz is the marshaled proof for the session.
-	proofBz []byte
+	// compactProofBz is the marshaled proof for the session.
+	compactProofBz []byte
 
 	// treeStore is the KVStore used to store the SMST.
 	treeStore pebble.PebbleKVStore
@@ -154,9 +154,7 @@ func (st *sessionTree) Update(key, value []byte, weight uint64) error {
 // This function is intended to be called after a session has been claimed and needs to be proven.
 // If the proof has already been generated, it returns the cached proof.
 // It returns an error if the SMST has not been flushed yet (the claim has not been generated)
-// TODO_IMPROVE(#427): Compress the proof into a SparseCompactClosestMerkleProof
-// prior to submitting to chain to reduce on-chain storage requirements for proofs.
-func (st *sessionTree) ProveClosest(path []byte) (proof *smt.SparseMerkleClosestProof, err error) {
+func (st *sessionTree) ProveClosest(path []byte) (compactProof *smt.SparseCompactMerkleClosestProof, err error) {
 	st.sessionMu.Lock()
 	defer st.sessionMu.Unlock()
 
@@ -166,13 +164,13 @@ func (st *sessionTree) ProveClosest(path []byte) (proof *smt.SparseMerkleClosest
 	}
 
 	// If the proof has already been generated, return the cached proof.
-	if st.proof != nil {
+	if st.compactProof != nil {
 		// Make sure the path is the same as the one for which the proof was generated.
 		if !bytes.Equal(path, st.proofPath) {
 			return nil, ErrSessionTreeProofPathMismatch
 		}
 
-		return st.proof, nil
+		return st.compactProof, nil
 	}
 
 	// Restore the KVStore from disk since it has been closed after the claim has been generated.
@@ -184,33 +182,45 @@ func (st *sessionTree) ProveClosest(path []byte) (proof *smt.SparseMerkleClosest
 	sessionSMT := smt.ImportSparseMerkleSumTrie(st.treeStore, sha256.New(), st.claimedRoot, smt.WithValueHasher(nil))
 
 	// Generate the proof and cache it along with the path for which it was generated.
-	proof, err = sessionSMT.ProveClosest(path)
+	// There is no ProveClosest variant that generates a compact proof directly.
+	// Generate a regular SparseMerkleClosestProof then compact it.
+	proof, err := sessionSMT.ProveClosest(path)
 	if err != nil {
 		return nil, err
 	}
 
-	proofBz, err := proof.Marshal()
+	compactProof, err = smt.CompactClosestProof(proof, &sessionSMT.TrieSpec)
+	if err != nil {
+		return nil, err
+	}
+
+	compactProofBz, err := compactProof.Marshal()
 	if err != nil {
 		return nil, err
 	}
 
 	// If no error occurred, cache the proof and the path for which it was generated.
 	st.sessionSMT = sessionSMT
 	st.proofPath = path
-	st.proof = proof
-	st.proofBz = proofBz
+	st.compactProof = compactProof
+	st.compactProofBz = compactProofBz
 
-	return st.proof, nil
+	return st.compactProof, nil
 }
 
 // GetProofBz returns the marshaled proof for the session.
 func (st *sessionTree) GetProofBz() []byte {
-	return st.proofBz
+	return st.compactProofBz
+}
+
+// GetTrieSpec returns the trie spec of the SMST.
+func (st *sessionTree) GetTrieSpec() smt.TrieSpec {
+	return *st.sessionSMT.Spec()
 }
 
 // GetProof returns the proof for the SMST if it has been generated or nil otherwise.
-func (st *sessionTree) GetProof() *smt.SparseMerkleClosestProof {
-	return st.proof
+func (st *sessionTree) GetProof() *smt.SparseCompactMerkleClosestProof {
+	return st.compactProof
 }
 
 // Flush gets the root hash of the SMST needed for submitting the claim;

pkg/relayer/session/sessiontree_test.go

@@ -1,3 +1,96 @@
 package session_test
 
-// TODO: Add tests to the sessionTree logic
+import (
+	"bytes"
+	"compress/gzip"
+	"crypto/rand"
+	"testing"
+
+	"github.com/pokt-network/poktroll/pkg/crypto/protocol"
+	"github.com/pokt-network/smt"
+	"github.com/pokt-network/smt/kvstore/pebble"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSessionTree_CompactProofsAreSmallerThanNonCompactProofs(t *testing.T) {
+	// Run the test for different number of leaves.
+	for numLeaf := 10; numLeaf <= 1000000; numLeaf *= 10 {
+		// We run the test 1000 times for each number of leaves to remove the randomness bias.
+		cumulativeProofSize := 0
+		cumulativeCompactProofSize := 0
+		cumulativeGzippedProofSize := 0
+		for numLeaf := 0; numLeaf <= 1000; numLeaf++ {
+			kvStore, err := pebble.NewKVStore("")
+			require.NoError(t, err)
+
+			trie := smt.NewSparseMerkleSumTrie(kvStore, protocol.NewTrieHasher(), smt.WithValueHasher(nil))
+
+			// Insert numLeaf random leaves.
+			for i := 0; i < numLeaf; i++ {
+				key := make([]byte, 32)
+				_, err := rand.Read(key)
+				require.NoError(t, err)
+				// Insert an empty value since this does not get affected by the compaction,
+				// this is also to not favor proof compression that compresses the value too.
+				trie.Update(key, []byte{}, 1)
+			}
+
+			// Generate a random path.
+			var path = make([]byte, 32)
+			_, err = rand.Read(path)
+			require.NoError(t, err)
+
+			// Create the proof.
+			proof, err := trie.ProveClosest(path)
+			require.NoError(t, err)
+
+			proofBz, err := proof.Marshal()
+			require.NoError(t, err)
+
+			// Accumulate the proof size over 1000 runs.
+			cumulativeProofSize += len(proofBz)
+
+			// Generate the compacted proof.
+			compactProof, err := smt.CompactClosestProof(proof, &trie.TrieSpec)
+			require.NoError(t, err)
+
+			compactProofBz, err := compactProof.Marshal()
+			require.NoError(t, err)
+
+			// Accumulate the compact proof size over 1000 runs.
+			cumulativeCompactProofSize += len(compactProofBz)
+
+			// Gzip the non compacted proof.
+			var buf bytes.Buffer
+			gzipWriter := gzip.NewWriter(&buf)
+			_, err = gzipWriter.Write(proofBz)
+			require.NoError(t, err)
+			err = gzipWriter.Close()
+			require.NoError(t, err)
+
+			// Accumulate the gzipped proof size over 1000 runs.
+			cumulativeGzippedProofSize += len(buf.Bytes())
+
+			//t.Logf(
+			//	"numLeaf: %d, proofSize: %d, compactProofSize: %d, gzipProofSize: %d",
+			//	numLeaf, len(proofBz), len(compactProofBz), len(buf.Bytes()),
+			//)
+
+			// Commenting out the assertion to not fail the test since compaction is not
+			// guaranteed to always reduce the proof size.
+			//require.Less(t, len(compactProofBz), len(proofBz))
+		}
+
+		//t.Logf(
+		//	"numLeaf=%d: cumulativeProofSize: %d, cumulativeCompactProofSize: %d, cumulativeGzippedProofSize: %d",
+		//	numLeaf, cumulativeProofSize, cumulativeCompactProofSize, cumulativeGzippedProofSize,
+		//)
+
+		t.Logf(
+			"numLeaf=%d: compactionRatio: %f, compressionRatio: %f",
+			numLeaf,
+			float32(cumulativeCompactProofSize)/float32(cumulativeProofSize),
+			float32(cumulativeGzippedProofSize)/float32(cumulativeProofSize),
+		)
+	}
+}

proto/poktroll/proof/tx.proto

@@ -91,7 +91,7 @@ message MsgSubmitProof {
   string                         supplier_operator_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
   poktroll.session.SessionHeader session_header = 2;
 
-  // serialized version of *smt.SparseMerkleClosestProof
+  // serialized version of *smt.SparseCompactMerkleClosestProof
   bytes proof = 3;
 }
 

proto/poktroll/proof/types.proto

@@ -18,7 +18,7 @@ message Proof {
   string supplier_operator_address = 1 [(cosmos_proto.scalar) = "cosmos.AddressString"];
   // The session header of the session that this claim is for.
   poktroll.session.SessionHeader session_header = 2;
-  // The serialized SMST proof from the `#ClosestProof()` method.
+  // The serialized SMST compacted proof from the `#ClosestProof()` method.
   bytes closest_merkle_proof = 3;
 }
 

testutil/testtree/tree.go

@@ -121,18 +121,18 @@ func NewProof(
 	t.Helper()
 
 	// Generate a closest proof from the session tree using closestProofPath.
-	merkleProof, err := sessionTree.ProveClosest(closestProofPath)
+	merkleCompactProof, err := sessionTree.ProveClosest(closestProofPath)
 	require.NoError(t, err)
-	require.NotNil(t, merkleProof)
+	require.NotNil(t, merkleCompactProof)
 
 	// Serialize the closest merkle proof.
-	merkleProofBz, err := merkleProof.Marshal()
+	merkleCompactProofBz, err := merkleCompactProof.Marshal()
 	require.NoError(t, err)
 
 	return &prooftypes.Proof{
 		SupplierOperatorAddress: supplierOperatorAddr,
 		SessionHeader:           sessionHeader,
-		ClosestMerkleProof:      merkleProofBz,
+		ClosestMerkleProof:      merkleCompactProofBz,
 	}
 }
 

x/proof/keeper/proof.go

@@ -15,8 +15,6 @@ import (
 func (k Keeper) UpsertProof(ctx context.Context, proof types.Proof) {
 	logger := k.Logger().With("method", "UpsertProof")
 
-	// TODO_MAINNET(#427): Use the marshal method on the SparseCompactClosestProof
-	// type here instead in order to reduce space stored on chain.
 	proofBz := k.cdc.MustMarshal(&proof)
 	storeAdapter := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
 

x/proof/keeper/proof_validation.go

@@ -91,16 +91,23 @@ func (k Keeper) EnsureValidProof(
 	}
 
 	// Unmarshal the closest merkle proof from the message.
-	sparseMerkleClosestProof := &smt.SparseMerkleClosestProof{}
-	if err = sparseMerkleClosestProof.Unmarshal(proof.ClosestMerkleProof); err != nil {
+	sparseCompactMerkleClosestProof := &smt.SparseCompactMerkleClosestProof{}
+	if err = sparseCompactMerkleClosestProof.Unmarshal(proof.ClosestMerkleProof); err != nil {
 		return types.ErrProofInvalidProof.Wrapf(
 			"failed to unmarshal closest merkle proof: %s",
 			err,
 		)
 	}
 
-	// TODO_MAINNET(#427): Utilize smt.VerifyCompactClosestProof here to
-	// reduce on-chain storage requirements for proofs.
+	// SparseCompactMerkeClosestProof does not implement GetValueHash, so we need to decompact it.
+	sparseMerkleClosestProof, err := smt.DecompactClosestProof(sparseCompactMerkleClosestProof, &protocol.SmtSpec)
+	if err != nil {
+		return types.ErrProofInvalidProof.Wrapf(
+			"failed to decompact closest merkle proof: %s",
+			err,
+		)
+	}
+
 	// Get the relay request and response from the proof.GetClosestMerkleProof.
 	relayBz := sparseMerkleClosestProof.GetValueHash(&protocol.SmtSpec)
 	relay := &servicetypes.Relay{}

x/proof/keeper/proof_validation_test.go

@@ -159,7 +159,7 @@ func TestEnsureValidProof_Error(t *testing.T) {
 
 	// Store the expected error returned during deserialization of the invalid
 	// closest Merkle proof bytes.
-	sparseMerkleClosestProof := &smt.SparseMerkleClosestProof{}
+	sparseMerkleClosestProof := &smt.SparseCompactMerkleClosestProof{}
 	expectedInvalidProofUnmarshalErr := sparseMerkleClosestProof.Unmarshal(invalidClosestProofBytes)
 
 	// Construct a relay to be mangled such that it fails to deserialize in order
@@ -777,7 +777,12 @@ func getClosestRelayDifficulty(
 	closestMerkleProofPath []byte,
 ) int64 {
 	// Retrieve a merkle proof that is closest to the path provided
-	closestMerkleProof, err := sessionTree.ProveClosest(closestMerkleProofPath)
+	closestCompactMerkleProof, err := sessionTree.ProveClosest(closestMerkleProofPath)
+	require.NoError(t, err)
+
+	trieSpec := sessionTree.GetTrieSpec()
+
+	closestMerkleProof, err := smt.DecompactClosestProof(closestCompactMerkleProof, &trieSpec)
 	require.NoError(t, err)
 
 	// Extract the Relay (containing the RelayResponse & RelayRequest) from the merkle proof.