github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/mmr/util_test.go (about)

     1  package mmr
     2  
     3  import (
     4  	"crypto/sha256"
     5  	"math"
     6  	"math/bits"
     7  	"testing"
     8  
     9  	"gotest.tools/assert"
    10  )
    11  
    12  func TestInclusionProof(t *testing.T) {
    13  	size := uint64(0b00111110)
    14  
    15  	completeSize := uint64(0b10000)
    16  	assert.Assert(t, minPeakHeight(completeSize) == 4 && maxPeakHeight(completeSize) == 4)
    17  	assert.Assert(t, minPeakHeight(size) == 1)
    18  	assert.Assert(t, maxPeakHeight(size) == 5)
    19  	assert.Assert(t, peakCount(size) == 5)
    20  	assert.Assert(t, leftPeak(size) == 2*0b100000-1 && leftPeakLeaf(size) == 0b100000)
    21  
    22  	sizes := getMoutainSizes(size)
    23  	assert.Assert(t,
    24  		sizes[0] == 2*0b100000-1 &&
    25  			sizes[1] == 2*0b10000-1 &&
    26  			sizes[2] == 2*0b1000-1 &&
    27  			sizes[3] == 2*0b100-1 &&
    28  			sizes[4] == 2*0b10-1)
    29  
    30  	indices := getMoutainPeaks(size)
    31  	assert.Assert(t,
    32  		indices[0] == 2*0b100000-1 &&
    33  			indices[1] == 2*0b100000-1+2*0b10000-1 &&
    34  			indices[2] == 2*0b100000-1+2*0b10000-1+2*0b1000-1 &&
    35  			indices[3] == 2*0b100000-1+2*0b10000-1+2*0b1000-1+2*0b100-1 &&
    36  			indices[4] == 2*0b100000-1+2*0b10000-1+2*0b1000-1+2*0b100-1+2*0b10-1)
    37  
    38  	store, err := NewFileHashStore("/tmp/hs.log", 0)
    39  	assert.Assert(t, err == nil)
    40  	defer store.Close()
    41  
    42  	mmr := NewMMR(0, nil, NewHasher([]byte{0}, []byte{1}), store)
    43  	h1 := sha256.Sum256([]byte{1})
    44  	mmr.PushHash(h1, false)
    45  	h1Idx := mmr.Size() - 1
    46  
    47  	h2 := sha256.Sum256([]byte{2})
    48  	ap2 := mmr.PushHash(h2, true)
    49  	h2Idx := mmr.Size() - 1
    50  	assert.Assert(t, len(ap2) == proofLength(h2Idx, mmr.Size()) && ap2[0] == h1)
    51  
    52  	rootHash2 := mmr.Root()
    53  
    54  	h3 := sha256.Sum256([]byte{3})
    55  	ap3 := mmr.PushHash(h3, true)
    56  	h3Idx := mmr.Size() - 1
    57  	assert.Assert(t, len(ap3) == proofLength(h3Idx, mmr.Size()))
    58  
    59  	assert.Assert(t, ComputeRoot([]HashType{h1, h2, h3}) == mmr.Root())
    60  
    61  	// h2's proof is returned by Push
    62  	err = mmr.VerifyInclusion(h2, rootHash2, h2Idx, h2Idx+1, ap2)
    63  	assert.Assert(t, err == nil)
    64  
    65  	// h3's proof is returned by Push
    66  	err = mmr.VerifyInclusion(h3, mmr.Root(), h3Idx, h3Idx+1, ap3)
    67  	assert.Assert(t, err == nil)
    68  
    69  	// generate proof for h1 wrt current root
    70  	proof, err := mmr.InclusionProof(h1Idx, mmr.Size())
    71  	assert.Assert(t, err == nil)
    72  	err = mmr.VerifyInclusion(h1, mmr.Root(), h1Idx, mmr.Size(), proof)
    73  	assert.Assert(t, err == nil)
    74  
    75  	// test getMoutainPeaks
    76  	peaks := getMoutainPeaks(8)
    77  	assert.Assert(t, len(peaks) == 1 && peaks[0] == 15)
    78  
    79  }
    80  
    81  func TestConsistencyProof(t *testing.T) {
    82  	store, err := NewFileHashStore("/tmp/hs.log", 0)
    83  	assert.Assert(t, err == nil)
    84  	defer store.Close()
    85  
    86  	m := NewMMR(0, nil, NewHasher([]byte{0}, []byte{1}), store)
    87  
    88  	n := uint64(7)
    89  	for i := uint64(0); i < n; i++ {
    90  		h := sha256.Sum256([]byte{byte(i + 1)})
    91  		m.PushHash(h, false)
    92  	}
    93  
    94  	cmp := []int{3, 2, 4, 1, 4, 3, 0}
    95  	for i := uint64(0); i < n; i++ {
    96  		proof, err := m.ConsistencyProof(uint64(i+1), n)
    97  		assert.Assert(t, err == nil && len(proof) == cmp[i])
    98  	}
    99  }
   100  
   101  func TestBits(t *testing.T) {
   102  	assert.Assert(t, bits.Len32(256) == 9)
   103  	assert.Assert(t, bits.Len32(0) == 0)
   104  	assert.Assert(t, bits.Len32(1) == 1)
   105  	assert.Assert(t, math.Ceil(math.Log2(float64(256))) == 8)
   106  	assert.Assert(t, 1<<8 == 256)
   107  }