Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[SMST] feat: Use compact SMST proofs #823

Merged
merged 11 commits into from
Sep 24, 2024
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) {
red-0ne marked this conversation as resolved.
Show resolved Hide resolved
// 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.

Loading