Skip to content

Commit

Permalink
[SMST] feat: Use compact SMST proofs (#823)
Browse files Browse the repository at this point in the history
## Summary

This PR changes the proofs being sent in `MsgSubmitPRoof` and stored by
the `SubmitProof` keeper from `SparseMerkleClosestProof` to
`SparseCompactMerkleClosestProof`.

## Issue

`Proofs` one of the most block-space consuming primitives. The SMST
support compact proofs that could help reduce their sizes.

- #758 

## Type of change

Select one or more from the following:

- [x] New feature, functionality or library
- [x] Consensus breaking; add the `consensus-breaking` label if so. See
#791 for details
- [ ] Bug fix
- [ ] Code health or cleanup
- [ ] Documentation
- [ ] Other (specify)

## Testing

- [x] **Unit Tests**: `make go_develop_and_test`
- [x] **LocalNet E2E Tests**: `make test_e2e`
- [ ] **DevNet E2E Tests**: Add the `devnet-test-e2e` label to the PR.

## Sanity Checklist

- [x] I have tested my changes using the available tooling
- [x] I have commented my code
- [x] I have performed a self-review of my own code; both comments &
source code
  • Loading branch information
red-0ne authored Sep 24, 2024
1 parent b39f84b commit 6ed08b4
Show file tree
Hide file tree
Showing 14 changed files with 163 additions and 42 deletions.
2 changes: 1 addition & 1 deletion api/poktroll/proof/tx.pulsar.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion api/poktroll/proof/types.pulsar.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 1 addition & 3 deletions pkg/client/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,11 +65,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,
Expand Down
5 changes: 4 additions & 1 deletion pkg/relayer/interface.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
44 changes: 27 additions & 17 deletions pkg/relayer/session/sessiontree.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand All @@ -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.
Expand All @@ -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;
Expand Down
104 changes: 103 additions & 1 deletion pkg/relayer/session/sessiontree_test.go
Original file line number Diff line number Diff line change
@@ -1,3 +1,105 @@
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"
)

const (
// Test multiple SMST sizes to see how the compaction ratio changes when the number
// of leaves increases.
// maxLeafs is the maximum number of leaves to test, after which the test stops.
maxLeafs = 10000
// Since the inserted leaves are random, we run the test for a given leaf count
// multiple times to remove the randomness bias.
numIterations = 100
)

// No significant performance gains were observed when using compact proofs compared
// to non-compact proofs.
// In fact, compact proofs appear to be less efficient than gzipped proofs, even
// without considering the "proof closest value" compression.
// For a sample comparison between compression and compaction ratios, see:
// https://github.com/pokt-network/poktroll/pull/823#issuecomment-2363987920
func TestSessionTree_CompactProofsAreSmallerThanNonCompactProofs(t *testing.T) {
// Run the test for different number of leaves.
for numLeafs := 10; numLeafs <= maxLeafs; numLeafs *= 10 {
cumulativeProofSize := 0
cumulativeCompactProofSize := 0
cumulativeGzippedProofSize := 0
// We run the test numIterations times for each number of leaves to remove the randomness bias.
for iteration := 0; iteration <= numIterations; iteration++ {
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 < numLeafs; 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 numIterations 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 numIterations 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 numIterations runs.
cumulativeGzippedProofSize += len(buf.Bytes())
}

// Calculate how much more efficient compact SMT proofs are compared to non-compact proofs.
compactionRatio := float32(cumulativeProofSize) / float32(cumulativeCompactProofSize)

// Claculate how much more efficient gzipped proofs are compared to non-compact proofs.
compressionRatio := float32(cumulativeProofSize) / float32(cumulativeGzippedProofSize)

// Gzip compression is more efficient than SMT compaction.
require.Greater(t, compressionRatio, compactionRatio)

t.Logf(
"numLeaf=%d: compactionRatio: %f, compressionRatio: %f",
numLeafs, compactionRatio, compressionRatio,
)
}
}
2 changes: 1 addition & 1 deletion proto/poktroll/proof/tx.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
2 changes: 1 addition & 1 deletion proto/poktroll/proof/types.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
8 changes: 4 additions & 4 deletions testutil/testtree/tree.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
}

Expand Down
2 changes: 0 additions & 2 deletions x/proof/keeper/proof.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))

Expand Down
15 changes: 11 additions & 4 deletions x/proof/keeper/proof_validation.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down
11 changes: 7 additions & 4 deletions x/proof/keeper/proof_validation_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -156,8 +156,8 @@ func TestEnsureValidProof_Error(t *testing.T) {

// Store the expected error returned during deserialization of the invalid
// closest Merkle proof bytes.
sparseMerkleClosestProof := &smt.SparseMerkleClosestProof{}
expectedInvalidProofUnmarshalErr := sparseMerkleClosestProof.Unmarshal(invalidClosestProofBytes)
sparseCompactMerkleClosestProof := &smt.SparseCompactMerkleClosestProof{}
expectedInvalidProofUnmarshalErr := sparseCompactMerkleClosestProof.Unmarshal(invalidClosestProofBytes)

// Construct a relay to be mangled such that it fails to deserialize in order
// to set the error expectation for the relevant test case.
Expand Down Expand Up @@ -611,9 +611,12 @@ func TestEnsureValidProof_Error(t *testing.T) {
)

// Extract relayHash to check below that it's difficulty is insufficient
sparseMerkleClosestProof := &smt.SparseMerkleClosestProof{}
err = sparseMerkleClosestProof.Unmarshal(proof.ClosestMerkleProof)
err = sparseCompactMerkleClosestProof.Unmarshal(proof.ClosestMerkleProof)
require.NoError(t, err)
var sparseMerkleClosestProof *smt.SparseMerkleClosestProof
sparseMerkleClosestProof, err = smt.DecompactClosestProof(sparseCompactMerkleClosestProof, &protocol.SmtSpec)
require.NoError(t, err)

relayBz := sparseMerkleClosestProof.GetValueHash(&protocol.SmtSpec)
relayHashArr := protocol.GetRelayHashFromBytes(relayBz)
relayHash := relayHashArr[:]
Expand Down
2 changes: 1 addition & 1 deletion x/proof/types/tx.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion x/proof/types/types.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 6ed08b4

Please sign in to comment.