github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/ledger/complete/mtrie/forest_test.go (about)

     1  package mtrie
     2  
     3  import (
     4  	"bytes"
     5  	"math/rand"
     6  	"sort"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/onflow/flow-go/ledger"
    14  	"github.com/onflow/flow-go/ledger/common/hash"
    15  	prf "github.com/onflow/flow-go/ledger/common/proof"
    16  	"github.com/onflow/flow-go/ledger/common/testutils"
    17  	"github.com/onflow/flow-go/ledger/complete/mtrie/trie"
    18  	"github.com/onflow/flow-go/ledger/partial/ptrie"
    19  	"github.com/onflow/flow-go/module/metrics"
    20  )
    21  
    22  // TestTrieOperations tests adding removing and retrieving Trie from Forest
    23  func TestTrieOperations(t *testing.T) {
    24  
    25  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
    26  	require.NoError(t, err)
    27  
    28  	// Make new Trie (independently of MForest):
    29  	nt := trie.NewEmptyMTrie()
    30  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
    31  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
    32  
    33  	updatedTrie, _, err := trie.NewTrieWithUpdatedRegisters(nt, []ledger.Path{p1}, []ledger.Payload{*v1}, true)
    34  	require.NoError(t, err)
    35  
    36  	// Add trie
    37  	err = forest.AddTrie(updatedTrie)
    38  	require.NoError(t, err)
    39  
    40  	// Get trie
    41  	retnt, err := forest.GetTrie(updatedTrie.RootHash())
    42  	require.NoError(t, err)
    43  	require.Equal(t, retnt.RootHash(), updatedTrie.RootHash())
    44  	require.Equal(t, 2, forest.Size())
    45  }
    46  
    47  // TestTrieUpdate updates the empty trie with some values and verifies that the
    48  // written values can be retrieved from the updated trie.
    49  func TestTrieUpdate(t *testing.T) {
    50  
    51  	metricsCollector := &metrics.NoopCollector{}
    52  	forest, err := NewForest(5, metricsCollector, nil)
    53  	require.NoError(t, err)
    54  	rootHash := forest.GetEmptyRootHash()
    55  
    56  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
    57  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
    58  
    59  	paths := []ledger.Path{p1}
    60  	payloads := []*ledger.Payload{v1}
    61  	update := &ledger.TrieUpdate{RootHash: rootHash, Paths: paths, Payloads: payloads}
    62  	updatedRoot, err := forest.Update(update)
    63  	require.NoError(t, err)
    64  
    65  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
    66  	retValues, err := forest.Read(read)
    67  	require.NoError(t, err)
    68  	require.Equal(t, retValues[0], payloads[0].Value())
    69  }
    70  
    71  // TestLeftEmptyInsert tests inserting a new value into an empty sub-trie:
    72  //  1. we first construct a baseTrie holding a couple of values on the right branch [~]
    73  //  2. we update a previously non-existent register on the left branch (X)
    74  //
    75  // We verify that values for _all_ paths in the updated Trie have correct payloads
    76  func TestLeftEmptyInsert(t *testing.T) {
    77  	//////////////////////
    78  	//     insert X     //
    79  	//       ()         //
    80  	//      /  \        //
    81  	//    (X)  [~]      //
    82  	//////////////////////
    83  
    84  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
    85  	require.NoError(t, err)
    86  
    87  	// path: 1000...
    88  	p1 := pathByUint8s([]uint8{uint8(129), uint8(1)})
    89  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
    90  
    91  	// path: 1100...
    92  	p2 := pathByUint8s([]uint8{uint8(193), uint8(1)})
    93  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
    94  
    95  	paths := []ledger.Path{p1, p2}
    96  	payloads := []*ledger.Payload{v1, v2}
    97  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
    98  	baseRoot, err := forest.Update(update)
    99  	require.NoError(t, err)
   100  
   101  	baseTrie, err := forest.GetTrie(baseRoot)
   102  	require.NoError(t, err)
   103  	require.Equal(t, uint64(2), baseTrie.AllocatedRegCount())
   104  	require.Equal(t, uint64(v1.Size()+v2.Size()), baseTrie.AllocatedRegSize())
   105  
   106  	p3 := pathByUint8s([]uint8{uint8(1), uint8(1)})
   107  	v3 := payloadBySlices([]byte{'C'}, []byte{'C'})
   108  
   109  	paths = []ledger.Path{p3}
   110  	payloads = []*ledger.Payload{v3}
   111  	update = &ledger.TrieUpdate{RootHash: baseTrie.RootHash(), Paths: paths, Payloads: payloads}
   112  	updatedRoot, err := forest.Update(update)
   113  	require.NoError(t, err)
   114  
   115  	updatedTrie, err := forest.GetTrie(updatedRoot)
   116  	require.NoError(t, err)
   117  	require.Equal(t, uint64(3), updatedTrie.AllocatedRegCount())
   118  	require.Equal(t, uint64(v1.Size()+v2.Size()+v3.Size()), updatedTrie.AllocatedRegSize())
   119  	paths = []ledger.Path{p1, p2, p3}
   120  	payloads = []*ledger.Payload{v1, v2, v3}
   121  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   122  	retValues, err := forest.Read(read)
   123  	require.NoError(t, err)
   124  	for i := range paths {
   125  		require.Equal(t, retValues[i], payloads[i].Value())
   126  	}
   127  }
   128  
   129  // TestRightEmptyInsert tests inserting a new value into an empty sub-trie:
   130  //  1. we first construct a baseTrie holding a couple of values on the left branch [~]
   131  //  2. we update a previously non-existent register on the right branch (X)
   132  //
   133  // We verify that values for _all_ paths in the updated Trie have correct payloads
   134  func TestRightEmptyInsert(t *testing.T) {
   135  	///////////////////////
   136  	//     insert X      //
   137  	//       ()          //
   138  	//      /  \         //
   139  	//    [~]  (X)       //
   140  	///////////////////////
   141  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   142  	require.NoError(t, err)
   143  
   144  	// path: 0000...
   145  	p1 := pathByUint8s([]uint8{uint8(1), uint8(1)})
   146  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   147  
   148  	// path: 0100...
   149  	p2 := pathByUint8s([]uint8{uint8(64), uint8(1)})
   150  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   151  
   152  	paths := []ledger.Path{p1, p2}
   153  	payloads := []*ledger.Payload{v1, v2}
   154  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   155  	baseRoot, err := forest.Update(update)
   156  	require.NoError(t, err)
   157  
   158  	baseTrie, err := forest.GetTrie(baseRoot)
   159  	require.NoError(t, err)
   160  	require.Equal(t, uint64(2), baseTrie.AllocatedRegCount())
   161  	require.Equal(t, uint64(v1.Size()+v2.Size()), baseTrie.AllocatedRegSize())
   162  
   163  	// path: 1000...
   164  	p3 := pathByUint8s([]uint8{uint8(129), uint8(1)})
   165  	v3 := payloadBySlices([]byte{'C'}, []byte{'C'})
   166  
   167  	paths = []ledger.Path{p3}
   168  	payloads = []*ledger.Payload{v3}
   169  	update = &ledger.TrieUpdate{RootHash: baseTrie.RootHash(), Paths: paths, Payloads: payloads}
   170  	updatedRoot, err := forest.Update(update)
   171  	require.NoError(t, err)
   172  
   173  	updatedTrie, err := forest.GetTrie(updatedRoot)
   174  	require.NoError(t, err)
   175  	require.Equal(t, uint64(3), updatedTrie.AllocatedRegCount())
   176  	require.Equal(t, uint64(v1.Size()+v2.Size()+v3.Size()), updatedTrie.AllocatedRegSize())
   177  
   178  	paths = []ledger.Path{p1, p2, p3}
   179  	payloads = []*ledger.Payload{v1, v2, v3}
   180  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   181  	retValues, err := forest.Read(read)
   182  	require.NoError(t, err)
   183  	for i := range paths {
   184  		require.Equal(t, retValues[i], payloads[i].Value())
   185  	}
   186  }
   187  
   188  // TestExpansionInsert tests inserting a new value into a populated sub-trie, where a
   189  // leaf (holding a single value) would be replaced by an expanded sub-trie holding multiple value
   190  //  1. we first construct a baseTrie holding a couple of values on the right branch [~]
   191  //  2. we update a previously non-existent register on the right branch turning [~] to [~']
   192  //
   193  // We verify that values for _all_ paths in the updated Trie have correct payloads
   194  func TestExpansionInsert(t *testing.T) {
   195  	////////////////////////
   196  	// modify [~] -> [~'] //
   197  	//       ()           //
   198  	//      /  \          //
   199  	//         [~]        //
   200  	////////////////////////
   201  
   202  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   203  	require.NoError(t, err)
   204  
   205  	// path: 100000...
   206  	p1 := pathByUint8s([]uint8{uint8(129), uint8(1)})
   207  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   208  
   209  	paths := []ledger.Path{p1}
   210  	payloads := []*ledger.Payload{v1}
   211  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   212  	baseRoot, err := forest.Update(update)
   213  	require.NoError(t, err)
   214  
   215  	baseTrie, err := forest.GetTrie(baseRoot)
   216  	require.NoError(t, err)
   217  	require.Equal(t, uint64(1), baseTrie.AllocatedRegCount())
   218  	require.Equal(t, uint64(v1.Size()), baseTrie.AllocatedRegSize())
   219  
   220  	// path: 1000001...
   221  	p2 := pathByUint8s([]uint8{uint8(130), uint8(1)})
   222  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   223  
   224  	paths = []ledger.Path{p2}
   225  	payloads = []*ledger.Payload{v2}
   226  	update = &ledger.TrieUpdate{RootHash: baseTrie.RootHash(), Paths: paths, Payloads: payloads}
   227  	updatedRoot, err := forest.Update(update)
   228  	require.NoError(t, err)
   229  
   230  	updatedTrie, err := forest.GetTrie(updatedRoot)
   231  	require.NoError(t, err)
   232  	require.Equal(t, uint64(2), updatedTrie.AllocatedRegCount())
   233  	require.Equal(t, uint64(v1.Size()+v2.Size()), updatedTrie.AllocatedRegSize())
   234  
   235  	paths = []ledger.Path{p1, p2}
   236  	payloads = []*ledger.Payload{v1, v2}
   237  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   238  	retValues, err := forest.Read(read)
   239  	require.NoError(t, err)
   240  	for i := range paths {
   241  		require.Equal(t, retValues[i], payloads[i].Value())
   242  	}
   243  }
   244  
   245  // TestFullHouseInsert tests inserting a new value into a populated sub-trie, where a
   246  // leaf's value is overridden _and_ further values are added which all fall into a subtree that
   247  // replaces the leaf:
   248  //  1. we first construct a baseTrie holding a couple of values on the right branch [~]
   249  //  2. we update a previously non-existent register on the right branch turning [~] to [~']
   250  //
   251  // We verify that values for _all_ paths in the updated Trie have correct payloads
   252  func TestFullHouseInsert(t *testing.T) {
   253  	///////////////////////
   254  	//   insert ~1<X<~2  //
   255  	//       ()          //
   256  	//      /  \         //
   257  	//    [~1]  [~2]     //
   258  	///////////////////////
   259  
   260  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   261  	require.NoError(t, err)
   262  
   263  	// paths p0 forms [~1]; p1 and p2 form [~2]
   264  	// path: 0100...
   265  	p0 := pathByUint8s([]uint8{uint8(64), uint8(1)})
   266  	v0 := payloadBySlices([]byte{'0'}, []byte{'0'})
   267  	// path: 1000...
   268  	p1 := pathByUint8s([]uint8{uint8(129), uint8(1)})
   269  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   270  
   271  	// path: 1100...
   272  	p2 := pathByUint8s([]uint8{uint8(193), uint8(1)})
   273  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   274  
   275  	paths := []ledger.Path{p0, p1, p2}
   276  	payloads := []*ledger.Payload{v0, v1, v2}
   277  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   278  	baseRoot, err := forest.Update(update)
   279  	require.NoError(t, err)
   280  
   281  	baseTrie, err := forest.GetTrie(baseRoot)
   282  	require.NoError(t, err)
   283  	require.Equal(t, uint64(3), baseTrie.AllocatedRegCount())
   284  	require.Equal(t, uint64(v0.Size()+v1.Size()+v2.Size()), baseTrie.AllocatedRegSize())
   285  
   286  	// we update value for path p1 and in addition add p3 that has the same prefix `10` as p1
   287  	v1 = payloadBySlices([]byte{'X'}, []byte{'X'})
   288  
   289  	// path: 1010...
   290  	p3 := pathByUint8s([]uint8{uint8(160), uint8(1)})
   291  	v3 := payloadBySlices([]byte{'C'}, []byte{'C'})
   292  
   293  	paths = []ledger.Path{p1, p3}
   294  	payloads = []*ledger.Payload{v1, v3}
   295  	update = &ledger.TrieUpdate{RootHash: baseTrie.RootHash(), Paths: paths, Payloads: payloads}
   296  	updatedRoot, err := forest.Update(update)
   297  	require.NoError(t, err)
   298  
   299  	updatedTrie, err := forest.GetTrie(updatedRoot)
   300  	require.NoError(t, err)
   301  	require.Equal(t, uint64(4), updatedTrie.AllocatedRegCount())
   302  	require.Equal(t, uint64(v0.Size()+v1.Size()+v2.Size()+v3.Size()), updatedTrie.AllocatedRegSize())
   303  
   304  	paths = []ledger.Path{p1, p2, p3}
   305  	payloads = []*ledger.Payload{v1, v2, v3}
   306  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   307  	retValues, err := forest.Read(read)
   308  	require.NoError(t, err)
   309  	for i := range paths {
   310  		require.Equal(t, retValues[i], payloads[i].Value())
   311  	}
   312  }
   313  
   314  // TestLeafInsert inserts two keys, which only differ in their last bit.
   315  // I.e. the trie needs to be expanded to its hull depth
   316  // We verify that values for _all_ paths in the updated Trie have correct payloads
   317  func TestLeafInsert(t *testing.T) {
   318  	///////////////////////
   319  	//   insert 1, 2     //
   320  	//       ()          //
   321  	//      /  \         //
   322  	//     ()   ...      //
   323  	//          /  \     //
   324  	//         ()  ()    //
   325  	///////////////////////
   326  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   327  	require.NoError(t, err)
   328  
   329  	// path: 000...0000000100000000
   330  	p1 := testutils.PathByUint16LeftPadded(256)
   331  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   332  
   333  	// path: 000...0000000100000001
   334  	p2 := testutils.PathByUint16LeftPadded(257)
   335  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   336  
   337  	paths := []ledger.Path{p1, p2}
   338  	payloads := []*ledger.Payload{v1, v2}
   339  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   340  	updatedRoot, err := forest.Update(update)
   341  	require.NoError(t, err)
   342  
   343  	updatedTrie, err := forest.GetTrie(updatedRoot)
   344  	require.NoError(t, err)
   345  	require.Equal(t, uint64(2), updatedTrie.AllocatedRegCount())
   346  	require.Equal(t, uint64(v1.Size()+v2.Size()), updatedTrie.AllocatedRegSize())
   347  
   348  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   349  	retValues, err := forest.Read(read)
   350  	require.NoError(t, err)
   351  	for i := range paths {
   352  		require.Equal(t, retValues[i], payloads[i].Value())
   353  	}
   354  }
   355  
   356  // TestOverrideValue overrides an existing value in the trie (without any expansion)
   357  // We verify that values for _all_ paths in the updated Trie have correct payloads
   358  func TestOverrideValue(t *testing.T) {
   359  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   360  	require.NoError(t, err)
   361  
   362  	// path: 1000...
   363  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   364  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   365  
   366  	// path: 0111...
   367  	p2 := pathByUint8s([]uint8{uint8(116), uint8(129)})
   368  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   369  
   370  	paths := []ledger.Path{p1, p2}
   371  	payloads := []*ledger.Payload{v1, v2}
   372  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   373  	baseRoot, err := forest.Update(update)
   374  	require.NoError(t, err)
   375  
   376  	// path: 1000...
   377  	p3 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   378  	v3 := payloadBySlices([]byte{'C'}, []byte{'C'})
   379  
   380  	paths = []ledger.Path{p3}
   381  	payloads = []*ledger.Payload{v3}
   382  	update = &ledger.TrieUpdate{RootHash: baseRoot, Paths: paths, Payloads: payloads}
   383  	updatedRoot, err := forest.Update(update)
   384  	require.NoError(t, err)
   385  
   386  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   387  	retValues, err := forest.Read(read)
   388  	require.NoError(t, err)
   389  	require.Equal(t, retValues[0], payloads[0].Value())
   390  
   391  }
   392  
   393  // TestDuplicateOverride tests behaviour when the updates contain two different payloads for the
   394  // same path. I.e. we update with (p0, v0) and (p0, v1)
   395  // We expect that the _last_ written value is persisted in the Trie
   396  func TestDuplicateOverride(t *testing.T) {
   397  
   398  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   399  	require.NoError(t, err)
   400  
   401  	// path: 1000...
   402  	p0 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   403  	v0 := payloadBySlices([]byte{'A'}, []byte{'A'})
   404  	paths := []ledger.Path{p0}
   405  	payloads := []*ledger.Payload{v0}
   406  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   407  	baseRoot, err := forest.Update(update)
   408  	require.NoError(t, err)
   409  
   410  	v1 := payloadBySlices([]byte{'B'}, []byte{'B'})
   411  	v2 := payloadBySlices([]byte{'C'}, []byte{'C'})
   412  	paths = []ledger.Path{p0, p0}
   413  	payloads = []*ledger.Payload{v1, v2}
   414  	update = &ledger.TrieUpdate{RootHash: baseRoot, Paths: paths, Payloads: payloads}
   415  	updatedRoot, err := forest.Update(update)
   416  	require.NoError(t, err)
   417  
   418  	paths = []ledger.Path{p0}
   419  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   420  	retValues, err := forest.Read(read)
   421  	require.NoError(t, err)
   422  	require.Equal(t, retValues[0], v2.Value())
   423  
   424  }
   425  
   426  // TestReadSafety check if payload returned from a forest are safe against modification,
   427  // ie. copy of the data is returned, instead of a slice
   428  func TestReadSafety(t *testing.T) {
   429  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   430  	require.NoError(t, err)
   431  
   432  	// path: 1000...
   433  	p0 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   434  	v0 := payloadBySlices([]byte{'A'}, []byte{'A'})
   435  	paths := []ledger.Path{p0}
   436  	payloads := []*ledger.Payload{v0}
   437  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   438  	baseRoot, err := forest.Update(update)
   439  	require.NoError(t, err)
   440  
   441  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: paths}
   442  	data, err := forest.Read(read)
   443  	require.NoError(t, err)
   444  
   445  	require.Len(t, data, 1)
   446  	require.Equal(t, v0.Value(), data[0])
   447  
   448  	// modify returned slice
   449  	data[0] = []byte("new value")
   450  
   451  	// read again
   452  	data2, err := forest.Read(read)
   453  	require.NoError(t, err)
   454  	require.Len(t, data2, 1)
   455  	require.Equal(t, v0.Value(), data2[0])
   456  }
   457  
   458  // TestReadOrder tests that payloads from reading a trie are delivered in the order as specified by the paths
   459  func TestReadOrder(t *testing.T) {
   460  
   461  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   462  	require.NoError(t, err)
   463  
   464  	p1 := pathByUint8s([]uint8{uint8(116), uint8(74)})
   465  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   466  
   467  	p2 := pathByUint8s([]uint8{uint8(53), uint8(129)})
   468  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   469  
   470  	paths := []ledger.Path{p1, p2}
   471  	payloads := []*ledger.Payload{v1, v2}
   472  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   473  	testRoot, err := forest.Update(update)
   474  	require.NoError(t, err)
   475  
   476  	read := &ledger.TrieRead{RootHash: testRoot, Paths: []ledger.Path{p1, p2}}
   477  	retValues, err := forest.Read(read)
   478  	require.NoError(t, err)
   479  	require.Equal(t, len(read.Paths), len(retValues))
   480  	require.Equal(t, retValues[0], payloads[0].Value())
   481  	require.Equal(t, retValues[1], payloads[1].Value())
   482  
   483  	read = &ledger.TrieRead{RootHash: testRoot, Paths: []ledger.Path{p2, p1}}
   484  	retValues, err = forest.Read(read)
   485  	require.NoError(t, err)
   486  	require.Equal(t, len(read.Paths), len(retValues))
   487  	require.Equal(t, retValues[1], payloads[0].Value())
   488  	require.Equal(t, retValues[0], payloads[1].Value())
   489  }
   490  
   491  // TestMixRead tests reading a mixture of set and unset registers.
   492  // We expect the default payload (nil) to be returned for unset registers.
   493  func TestMixRead(t *testing.T) {
   494  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   495  	require.NoError(t, err)
   496  
   497  	// path: 01111101...
   498  	p1 := pathByUint8s([]uint8{uint8(125), uint8(23)})
   499  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   500  
   501  	// path: 10110010...
   502  	p2 := pathByUint8s([]uint8{uint8(178), uint8(152)})
   503  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   504  
   505  	paths := []ledger.Path{p1, p2}
   506  	payloads := []*ledger.Payload{v1, v2}
   507  
   508  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   509  	baseRoot, err := forest.Update(update)
   510  	require.NoError(t, err)
   511  
   512  	// path: 01101110...
   513  	p3 := pathByUint8s([]uint8{uint8(110), uint8(48)})
   514  	v3 := ledger.EmptyPayload()
   515  
   516  	// path: 00010111...
   517  	p4 := pathByUint8s([]uint8{uint8(23), uint8(82)})
   518  	v4 := ledger.EmptyPayload()
   519  
   520  	readPaths := []ledger.Path{p1, p2, p3, p4}
   521  	expectedPayloads := []*ledger.Payload{v1, v2, v3, v4}
   522  
   523  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: readPaths}
   524  	retValues, err := forest.Read(read)
   525  	require.NoError(t, err)
   526  	for i := range paths {
   527  		require.Equal(t, retValues[i], expectedPayloads[i].Value())
   528  	}
   529  }
   530  
   531  // TestReadWithDuplicatedKeys reads a the values for two keys, where both keys have the same value.
   532  // We expect that we receive the respective value twice in the return.
   533  func TestReadWithDuplicatedKeys(t *testing.T) {
   534  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   535  	require.NoError(t, err)
   536  
   537  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   538  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   539  	p2 := pathByUint8s([]uint8{uint8(116), uint8(129)})
   540  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   541  	p3 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   542  
   543  	paths := []ledger.Path{p1, p2}
   544  	payloads := []*ledger.Payload{v1, v2}
   545  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   546  	updatedRoot, err := forest.Update(update)
   547  	require.NoError(t, err)
   548  
   549  	paths = []ledger.Path{p1, p2, p3}
   550  	expectedPayloads := []*ledger.Payload{v1, v2, v1}
   551  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   552  	retValues, err := forest.Read(read)
   553  	require.NoError(t, err)
   554  	require.Equal(t, len(read.Paths), len(retValues))
   555  	for i := range paths {
   556  		require.Equal(t, retValues[i], expectedPayloads[i].Value())
   557  	}
   558  }
   559  
   560  // TestReadNonExistingPath tests reading an unset path.
   561  func TestReadNonExistingPath(t *testing.T) {
   562  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   563  	require.NoError(t, err)
   564  
   565  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   566  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   567  	paths := []ledger.Path{p1}
   568  	payloads := []*ledger.Payload{v1}
   569  
   570  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   571  	updatedRoot, err := forest.Update(update)
   572  	require.NoError(t, err)
   573  
   574  	p2 := pathByUint8s([]uint8{uint8(116), uint8(129)})
   575  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: []ledger.Path{p2}}
   576  	retValues, err := forest.Read(read)
   577  	require.NoError(t, err)
   578  	require.Equal(t, 0, len(retValues[0]))
   579  }
   580  
   581  // TestReadSinglePayload tests reading a single payload of set/unset register.
   582  func TestReadSinglePayload(t *testing.T) {
   583  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   584  	require.NoError(t, err)
   585  
   586  	// path: 01111101...
   587  	path1 := pathByUint8s([]uint8{uint8(125), uint8(23)})
   588  	payload1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   589  
   590  	// path: 10110010...
   591  	path2 := pathByUint8s([]uint8{uint8(178), uint8(152)})
   592  	payload2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   593  
   594  	paths := []ledger.Path{path1, path2}
   595  	payloads := []*ledger.Payload{payload1, payload2}
   596  
   597  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   598  	baseRoot, err := forest.Update(update)
   599  	require.NoError(t, err)
   600  
   601  	// path: 01101110...
   602  	path3 := pathByUint8s([]uint8{uint8(110), uint8(48)})
   603  	payload3 := ledger.EmptyPayload()
   604  
   605  	// path: 00010111...
   606  	path4 := pathByUint8s([]uint8{uint8(23), uint8(82)})
   607  	payload4 := ledger.EmptyPayload()
   608  
   609  	expectedPayloads := make(map[ledger.Path]*ledger.Payload)
   610  	expectedPayloads[path1] = payload1
   611  	expectedPayloads[path2] = payload2
   612  	expectedPayloads[path3] = payload3
   613  	expectedPayloads[path4] = payload4
   614  
   615  	// Batch read one payload at a time (less efficient)
   616  	for path, payload := range expectedPayloads {
   617  		read := &ledger.TrieRead{RootHash: baseRoot, Paths: []ledger.Path{path}}
   618  		retValues, err := forest.Read(read)
   619  		require.NoError(t, err)
   620  		require.Equal(t, 1, len(retValues))
   621  		if payload.IsEmpty() {
   622  			require.Equal(t, 0, len(retValues[0]))
   623  		} else {
   624  			require.Equal(t, payload.Value(), retValues[0])
   625  		}
   626  	}
   627  
   628  	// Read single value
   629  	for path, payload := range expectedPayloads {
   630  		read := &ledger.TrieReadSingleValue{RootHash: baseRoot, Path: path}
   631  		retValue, err := forest.ReadSingleValue(read)
   632  		require.NoError(t, err)
   633  		if payload.IsEmpty() {
   634  			require.Equal(t, 0, len(retValue))
   635  		} else {
   636  			require.Equal(t, payload.Value(), retValue)
   637  		}
   638  	}
   639  }
   640  
   641  // TestForkingUpdates updates a base trie in two different ways. We expect
   642  // that for each update, a new trie is added to the forest preserving the
   643  // updated values independently of the other update.
   644  func TestForkingUpdates(t *testing.T) {
   645  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   646  	require.NoError(t, err)
   647  
   648  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   649  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   650  	p2 := pathByUint8s([]uint8{uint8(116), uint8(129)})
   651  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   652  	paths := []ledger.Path{p1, p2}
   653  	payloads := []*ledger.Payload{v1, v2}
   654  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   655  	baseRoot, err := forest.Update(update)
   656  	require.NoError(t, err)
   657  
   658  	// update baseTrie -> updatedTrieA
   659  	v1a := payloadBySlices([]byte{'C'}, []byte{'C'})
   660  	p3a := pathByUint8s([]uint8{uint8(116), uint8(22)})
   661  	v3a := payloadBySlices([]byte{'D'}, []byte{'D'})
   662  	pathsA := []ledger.Path{p1, p3a}
   663  	payloadsA := []*ledger.Payload{v1a, v3a}
   664  	updateA := &ledger.TrieUpdate{RootHash: baseRoot, Paths: pathsA, Payloads: payloadsA}
   665  	updatedRootA, err := forest.Update(updateA)
   666  	require.NoError(t, err)
   667  
   668  	// update baseTrie -> updatedTrieB
   669  	v1b := payloadBySlices([]byte{'E'}, []byte{'E'})
   670  	p3b := pathByUint8s([]uint8{uint8(116), uint8(22)})
   671  	v3b := payloadBySlices([]byte{'F'}, []byte{'F'})
   672  	pathsB := []ledger.Path{p1, p3b}
   673  	payloadsB := []*ledger.Payload{v1b, v3b}
   674  	updateB := &ledger.TrieUpdate{RootHash: baseRoot, Paths: pathsB, Payloads: payloadsB}
   675  	updatedRootB, err := forest.Update(updateB)
   676  	require.NoError(t, err)
   677  
   678  	// Verify payloads are preserved
   679  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: paths}
   680  	retValues, err := forest.Read(read) // reading from original Trie
   681  	require.NoError(t, err)
   682  	for i := range paths {
   683  		require.Equal(t, retValues[i], payloads[i].Value())
   684  	}
   685  
   686  	readA := &ledger.TrieRead{RootHash: updatedRootA, Paths: pathsA}
   687  	retValues, err = forest.Read(readA) // reading from updatedTrieA
   688  	require.NoError(t, err)
   689  	for i := range paths {
   690  		require.Equal(t, retValues[i], payloadsA[i].Value())
   691  	}
   692  
   693  	readB := &ledger.TrieRead{RootHash: updatedRootB, Paths: pathsB}
   694  	retValues, err = forest.Read(readB) // reading from updatedTrieB
   695  	require.NoError(t, err)
   696  	for i := range paths {
   697  		require.Equal(t, retValues[i], payloadsB[i].Value())
   698  	}
   699  }
   700  
   701  // TestIdenticalUpdateAppliedTwice updates a base trie in the same way twice.
   702  // Hence, the forest should de-duplicate the resulting two version of the identical trie
   703  // without an error.
   704  func TestIdenticalUpdateAppliedTwice(t *testing.T) {
   705  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   706  	require.NoError(t, err)
   707  
   708  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
   709  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   710  	p2 := pathByUint8s([]uint8{uint8(116), uint8(129)})
   711  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   712  	paths := []ledger.Path{p1, p2}
   713  	payloads := []*ledger.Payload{v1, v2}
   714  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   715  	baseRoot, err := forest.Update(update)
   716  	require.NoError(t, err)
   717  
   718  	p3 := pathByUint8s([]uint8{uint8(116), uint8(22)})
   719  	v3 := payloadBySlices([]byte{'D'}, []byte{'D'})
   720  
   721  	update = &ledger.TrieUpdate{RootHash: baseRoot, Paths: []ledger.Path{p3}, Payloads: []*ledger.Payload{v3}}
   722  	updatedRootA, err := forest.Update(update)
   723  	require.NoError(t, err)
   724  	updatedRootB, err := forest.Update(update)
   725  	require.NoError(t, err)
   726  	require.Equal(t, updatedRootA, updatedRootB)
   727  
   728  	paths = []ledger.Path{p1, p2, p3}
   729  	payloads = []*ledger.Payload{v1, v2, v3}
   730  	read := &ledger.TrieRead{RootHash: updatedRootA, Paths: paths}
   731  	retValuesA, err := forest.Read(read)
   732  	require.NoError(t, err)
   733  	for i := range paths {
   734  		require.Equal(t, retValuesA[i], payloads[i].Value())
   735  	}
   736  
   737  	read = &ledger.TrieRead{RootHash: updatedRootB, Paths: paths}
   738  	retValuesB, err := forest.Read(read)
   739  	require.NoError(t, err)
   740  	for i := range paths {
   741  		require.Equal(t, retValuesB[i], payloads[i].Value())
   742  	}
   743  }
   744  
   745  func TestNonExistingInvalidProof(t *testing.T) {
   746  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   747  	require.NoError(t, err)
   748  	paths := testutils.RandomPaths(2)
   749  	payloads := testutils.RandomPayloads(len(paths), 10, 20)
   750  
   751  	existingPaths := paths[1:]
   752  	existingPayloads := payloads[1:]
   753  
   754  	nonExistingPaths := paths[:1]
   755  
   756  	activeRoot := forest.GetEmptyRootHash()
   757  
   758  	// adding a payload to a path
   759  	update := &ledger.TrieUpdate{RootHash: activeRoot, Paths: existingPaths, Payloads: existingPayloads}
   760  	activeRoot, err = forest.Update(update)
   761  	require.NoError(t, err, "error updating")
   762  
   763  	// reading proof for nonExistingPaths
   764  	read := &ledger.TrieRead{RootHash: activeRoot, Paths: nonExistingPaths}
   765  	batchProof, err := forest.Proofs(read)
   766  	require.NoError(t, err, "error generating proofs")
   767  
   768  	// now a malicious node modifies the proof to be Inclusion false
   769  	// and change the proof such that one interim node has invalid hash
   770  	batchProof.Proofs[0].Inclusion = false
   771  	batchProof.Proofs[0].Interims[0] = hash.DummyHash
   772  
   773  	// expect the VerifyTrieBatchProof should return false
   774  	require.False(t, prf.VerifyTrieBatchProof(batchProof, ledger.State(activeRoot)))
   775  }
   776  
   777  // TestRandomUpdateReadProofValueSizes repeats a sequence of actions update, read, get value sizes, and proof random paths
   778  // this simulates the common pattern of actions on flow
   779  func TestRandomUpdateReadProofValueSizes(t *testing.T) {
   780  
   781  	minPayloadByteSize := 2
   782  	maxPayloadByteSize := 10
   783  	rep := 10
   784  	maxNumPathsPerStep := 10
   785  	seed := time.Now().UnixNano()
   786  	t.Log(seed)
   787  
   788  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   789  	require.NoError(t, err)
   790  
   791  	activeRoot := forest.GetEmptyRootHash()
   792  	require.NoError(t, err)
   793  	latestPayloadByPath := make(map[ledger.Path]*ledger.Payload) // map store
   794  
   795  	for e := 0; e < rep; e++ {
   796  		paths := testutils.RandomPathsRandLen(maxNumPathsPerStep)
   797  		payloads := testutils.RandomPayloads(len(paths), minPayloadByteSize, maxPayloadByteSize)
   798  
   799  		// update map store with key values
   800  		// we use this at the end of each step to check all existing keys
   801  		for i, p := range paths {
   802  			latestPayloadByPath[p] = payloads[i]
   803  		}
   804  
   805  		// test reading for non-existing keys
   806  		nonExistingPaths := make([]ledger.Path, 0)
   807  		otherPaths := testutils.RandomPathsRandLen(maxNumPathsPerStep)
   808  		for _, p := range otherPaths {
   809  			if _, ok := latestPayloadByPath[p]; !ok {
   810  				nonExistingPaths = append(nonExistingPaths, p)
   811  			}
   812  		}
   813  		read := &ledger.TrieRead{RootHash: activeRoot, Paths: nonExistingPaths}
   814  		retValues, err := forest.Read(read)
   815  		require.NoError(t, err, "error reading - non existing paths")
   816  		for _, p := range retValues {
   817  			require.Equal(t, 0, len(p))
   818  		}
   819  
   820  		// test value sizes for non-existent keys
   821  		retValueSizes, err := forest.ValueSizes(read)
   822  		require.NoError(t, err, "error value sizes - non existent paths")
   823  		require.Equal(t, len(read.Paths), len(retValueSizes))
   824  		for _, size := range retValueSizes {
   825  			require.Equal(t, 0, size)
   826  		}
   827  
   828  		// test update
   829  		update := &ledger.TrieUpdate{RootHash: activeRoot, Paths: paths, Payloads: payloads}
   830  		activeRoot, err = forest.Update(update)
   831  		require.NoError(t, err, "error updating")
   832  
   833  		// test read
   834  		read = &ledger.TrieRead{RootHash: activeRoot, Paths: paths}
   835  		retValues, err = forest.Read(read)
   836  		require.NoError(t, err, "error reading")
   837  		for i := range payloads {
   838  			require.Equal(t, retValues[i], payloads[i].Value())
   839  		}
   840  
   841  		// test value sizes for existing keys
   842  		retValueSizes, err = forest.ValueSizes(read)
   843  		require.NoError(t, err, "error value sizes")
   844  		require.Equal(t, len(read.Paths), len(retValueSizes))
   845  		for i := range payloads {
   846  			require.Equal(t, payloads[i].Value().Size(), retValueSizes[i])
   847  		}
   848  
   849  		// test proof (mix of existing and non existing keys)
   850  		proofPaths := make([]ledger.Path, 0)
   851  		proofPaths = append(proofPaths, paths...)
   852  		proofPaths = append(proofPaths, nonExistingPaths...)
   853  
   854  		// shuffle the order of `proofPaths` to run `Proofs` on non-sorted input paths
   855  		rand.Shuffle(len(proofPaths), func(i, j int) {
   856  			proofPaths[i], proofPaths[j] = proofPaths[j], proofPaths[i]
   857  		})
   858  
   859  		// sort `proofPaths` into another slice
   860  		sortedPaths := sortedCopy(proofPaths)
   861  
   862  		read = &ledger.TrieRead{RootHash: activeRoot, Paths: proofPaths}
   863  		batchProof, err := forest.Proofs(read)
   864  		require.NoError(t, err, "error generating proofs")
   865  		assert.True(t, prf.VerifyTrieBatchProof(batchProof, ledger.State(activeRoot)))
   866  
   867  		// check `Proofs` has sorted the input paths.
   868  		// this check is needed to not weaken the SPoCK secret entropy,
   869  		// when `Proofs` is used to generate chunk data.
   870  		assert.Equal(t, sortedPaths, read.Paths)
   871  
   872  		// build a partial trie from batch proofs and check the root hash is equal
   873  		psmt, err := ptrie.NewPSMT(activeRoot, batchProof)
   874  		require.NoError(t, err, "error building partial trie")
   875  		assert.Equal(t, psmt.RootHash(), activeRoot)
   876  
   877  		// check payloads for all existing paths
   878  		allPaths := make([]ledger.Path, 0, len(latestPayloadByPath))
   879  		allPayloads := make([]*ledger.Payload, 0, len(latestPayloadByPath))
   880  		for p, v := range latestPayloadByPath {
   881  			allPaths = append(allPaths, p)
   882  			allPayloads = append(allPayloads, v)
   883  		}
   884  
   885  		read = &ledger.TrieRead{RootHash: activeRoot, Paths: allPaths}
   886  		retValues, err = forest.Read(read)
   887  		require.NoError(t, err)
   888  		for i, v := range allPayloads {
   889  			assert.Equal(t, retValues[i], v.Value())
   890  		}
   891  
   892  		// check value sizes for all existing paths
   893  		retValueSizes, err = forest.ValueSizes(read)
   894  		require.NoError(t, err)
   895  		require.Equal(t, len(read.Paths), len(retValueSizes))
   896  		for i, v := range allPayloads {
   897  			assert.Equal(t, v.Value().Size(), retValueSizes[i])
   898  		}
   899  	}
   900  }
   901  
   902  func sortedCopy(paths []ledger.Path) []ledger.Path {
   903  	sortedPaths := make([]ledger.Path, len(paths))
   904  	copy(sortedPaths, paths)
   905  	sort.Slice(sortedPaths, func(i, j int) bool {
   906  		return bytes.Compare(sortedPaths[i][:], sortedPaths[j][:]) < 0
   907  	})
   908  	return sortedPaths
   909  }
   910  
   911  // TestProofGenerationInclusion tests that inclusion proofs generated by a Trie pass verification
   912  func TestProofGenerationInclusion(t *testing.T) {
   913  
   914  	metricsCollector := &metrics.NoopCollector{}
   915  	forest, err := NewForest(5, metricsCollector, nil)
   916  	require.NoError(t, err)
   917  	emptyRoot := forest.GetEmptyRootHash()
   918  
   919  	p1 := pathByUint8s([]uint8{uint8(1), uint8(74)})
   920  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
   921  	p2 := pathByUint8s([]uint8{uint8(2), uint8(74)})
   922  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
   923  	p3 := pathByUint8s([]uint8{uint8(130), uint8(74)})
   924  	v3 := payloadBySlices([]byte{'C'}, []byte{'C'})
   925  	p4 := pathByUint8s([]uint8{uint8(131), uint8(74)})
   926  	v4 := payloadBySlices([]byte{'D'}, []byte{'D'})
   927  	paths := []ledger.Path{p1, p2, p3, p4}
   928  	payloads := []*ledger.Payload{v1, v2, v3, v4}
   929  
   930  	// shuffle the order of `proofPaths` to run `Proofs` on non-sorted input paths
   931  	rand.Shuffle(len(paths), func(i, j int) {
   932  		paths[i], paths[j] = paths[j], paths[i]
   933  	})
   934  
   935  	// sort `proofPaths` into another slice
   936  	sortedPaths := sortedCopy(paths)
   937  
   938  	update := &ledger.TrieUpdate{RootHash: emptyRoot, Paths: paths, Payloads: payloads}
   939  	updatedRoot, err := forest.Update(update)
   940  	require.NoError(t, err)
   941  	read := &ledger.TrieRead{RootHash: updatedRoot, Paths: paths}
   942  	proof, err := forest.Proofs(read)
   943  	require.NoError(t, err)
   944  
   945  	// verify batch proofs.
   946  	assert.True(t, prf.VerifyTrieBatchProof(proof, ledger.State(updatedRoot)))
   947  
   948  	// check `Proofs` has sorted the input paths.
   949  	// this check is needed to not weaken the SPoCK secret entropy,
   950  	// when `Proofs` is used to generate chunk data.
   951  	assert.Equal(t, sortedPaths, read.Paths)
   952  }
   953  
   954  func payloadBySlices(keydata []byte, valuedata []byte) *ledger.Payload {
   955  	key := ledger.Key{KeyParts: []ledger.KeyPart{{Type: 0, Value: keydata}}}
   956  	value := ledger.Value(valuedata)
   957  	return ledger.NewPayload(key, value)
   958  }
   959  
   960  func pathByUint8s(inputs []uint8) ledger.Path {
   961  	var b ledger.Path
   962  	copy(b[:], inputs)
   963  	return b
   964  }
   965  
   966  // TestValueSizesOrder tests returned value sizes are in the order as specified by the paths
   967  func TestValueSizesOrder(t *testing.T) {
   968  
   969  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
   970  	require.NoError(t, err)
   971  
   972  	// path: 01111101...
   973  	p1 := pathByUint8s([]uint8{uint8(125), uint8(23)})
   974  	v1 := testutils.RandomPayload(1, 100)
   975  
   976  	// path: 10110010...
   977  	p2 := pathByUint8s([]uint8{uint8(178), uint8(152)})
   978  	v2 := testutils.RandomPayload(1, 100)
   979  
   980  	paths := []ledger.Path{p1, p2}
   981  	payloads := []*ledger.Payload{v1, v2}
   982  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
   983  	baseRoot, err := forest.Update(update)
   984  	require.NoError(t, err)
   985  
   986  	// Get value sizes for paths {p1, p2}
   987  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: []ledger.Path{p1, p2}}
   988  	retValueSizes, err := forest.ValueSizes(read)
   989  	require.NoError(t, err)
   990  	require.Equal(t, len(read.Paths), len(retValueSizes))
   991  	require.Equal(t, v1.Value().Size(), retValueSizes[0])
   992  	require.Equal(t, v2.Value().Size(), retValueSizes[1])
   993  
   994  	// Get value sizes for paths {p2, p1}
   995  	read = &ledger.TrieRead{RootHash: baseRoot, Paths: []ledger.Path{p2, p1}}
   996  	retValueSizes, err = forest.ValueSizes(read)
   997  	require.NoError(t, err)
   998  	require.Equal(t, len(read.Paths), len(retValueSizes))
   999  	require.Equal(t, v2.Value().Size(), retValueSizes[0])
  1000  	require.Equal(t, v1.Value().Size(), retValueSizes[1])
  1001  }
  1002  
  1003  // TestMixGetValueSizes tests value sizes of a mix of set and unset registers.
  1004  // We expect value size 0 to be returned for unset registers.
  1005  func TestMixGetValueSizes(t *testing.T) {
  1006  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
  1007  	require.NoError(t, err)
  1008  
  1009  	// path: 01111101...
  1010  	p1 := pathByUint8s([]uint8{uint8(125), uint8(23)})
  1011  	v1 := testutils.RandomPayload(1, 100)
  1012  
  1013  	// path: 10110010...
  1014  	p2 := pathByUint8s([]uint8{uint8(178), uint8(152)})
  1015  	v2 := testutils.RandomPayload(1, 100)
  1016  
  1017  	paths := []ledger.Path{p1, p2}
  1018  	payloads := []*ledger.Payload{v1, v2}
  1019  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
  1020  	baseRoot, err := forest.Update(update)
  1021  	require.NoError(t, err)
  1022  
  1023  	// path: 01101110...
  1024  	p3 := pathByUint8s([]uint8{uint8(110), uint8(48)})
  1025  
  1026  	// path: 00010111...
  1027  	p4 := pathByUint8s([]uint8{uint8(23), uint8(82)})
  1028  
  1029  	readPaths := []ledger.Path{p1, p2, p3, p4}
  1030  	expectedValueSizes := []int{v1.Value().Size(), v2.Value().Size(), 0, 0}
  1031  
  1032  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: readPaths}
  1033  	retValueSizes, err := forest.ValueSizes(read)
  1034  	require.NoError(t, err)
  1035  	require.Equal(t, len(read.Paths), len(retValueSizes))
  1036  	for i := range read.Paths {
  1037  		require.Equal(t, expectedValueSizes[i], retValueSizes[i])
  1038  	}
  1039  }
  1040  
  1041  // TestValueSizesWithDuplicatedKeys gets value sizes for two keys,
  1042  // where both keys have the same value.
  1043  // We expect to receive same value sizes twice.
  1044  func TestValueSizesWithDuplicatedKeys(t *testing.T) {
  1045  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
  1046  	require.NoError(t, err)
  1047  
  1048  	// path: 01111101...
  1049  	p1 := pathByUint8s([]uint8{uint8(125), uint8(23)})
  1050  	v1 := testutils.RandomPayload(1, 100)
  1051  
  1052  	// path: 10110010...
  1053  	p2 := pathByUint8s([]uint8{uint8(178), uint8(152)})
  1054  	v2 := testutils.RandomPayload(1, 100)
  1055  
  1056  	// same path as p1
  1057  	p3 := pathByUint8s([]uint8{uint8(125), uint8(23)})
  1058  
  1059  	paths := []ledger.Path{p1, p2}
  1060  	payloads := []*ledger.Payload{v1, v2}
  1061  	update := &ledger.TrieUpdate{RootHash: forest.GetEmptyRootHash(), Paths: paths, Payloads: payloads}
  1062  	baseRoot, err := forest.Update(update)
  1063  	require.NoError(t, err)
  1064  
  1065  	readPaths := []ledger.Path{p1, p2, p3}
  1066  	expectedValueSizes := []int{v1.Value().Size(), v2.Value().Size(), v1.Value().Size()}
  1067  
  1068  	read := &ledger.TrieRead{RootHash: baseRoot, Paths: readPaths}
  1069  	retValueSizes, err := forest.ValueSizes(read)
  1070  	require.NoError(t, err)
  1071  	require.Equal(t, len(read.Paths), len(retValueSizes))
  1072  	for i := range read.Paths {
  1073  		require.Equal(t, expectedValueSizes[i], retValueSizes[i])
  1074  	}
  1075  }
  1076  
  1077  func TestPurgeCacheExcept(t *testing.T) {
  1078  	forest, err := NewForest(5, &metrics.NoopCollector{}, nil)
  1079  	require.NoError(t, err)
  1080  
  1081  	nt := trie.NewEmptyMTrie()
  1082  	p1 := pathByUint8s([]uint8{uint8(53), uint8(74)})
  1083  	v1 := payloadBySlices([]byte{'A'}, []byte{'A'})
  1084  
  1085  	updatedTrie1, _, err := trie.NewTrieWithUpdatedRegisters(nt, []ledger.Path{p1}, []ledger.Payload{*v1}, true)
  1086  	require.NoError(t, err)
  1087  
  1088  	err = forest.AddTrie(updatedTrie1)
  1089  	require.NoError(t, err)
  1090  
  1091  	p2 := pathByUint8s([]uint8{uint8(12), uint8(34)})
  1092  	v2 := payloadBySlices([]byte{'B'}, []byte{'B'})
  1093  
  1094  	updatedTrie2, _, err := trie.NewTrieWithUpdatedRegisters(nt, []ledger.Path{p2}, []ledger.Payload{*v2}, true)
  1095  	require.NoError(t, err)
  1096  
  1097  	err = forest.AddTrie(updatedTrie2)
  1098  	require.NoError(t, err)
  1099  	require.Equal(t, 3, forest.tries.Count())
  1100  
  1101  	err = forest.PurgeCacheExcept(updatedTrie2.RootHash())
  1102  	require.NoError(t, err)
  1103  	require.Equal(t, 1, forest.tries.Count())
  1104  
  1105  	ret, err := forest.GetTrie(updatedTrie2.RootHash())
  1106  	require.NoError(t, err)
  1107  	require.Equal(t, ret, updatedTrie2)
  1108  
  1109  	_, err = forest.GetTrie(updatedTrie1.RootHash())
  1110  	require.Error(t, err)
  1111  
  1112  	// test purge with non existing trie
  1113  	err = forest.PurgeCacheExcept(updatedTrie1.RootHash())
  1114  	require.Error(t, err)
  1115  
  1116  	ret, err = forest.GetTrie(updatedTrie2.RootHash())
  1117  	require.NoError(t, err)
  1118  	require.Equal(t, ret, updatedTrie2)
  1119  
  1120  	_, err = forest.GetTrie(updatedTrie1.RootHash())
  1121  	require.Error(t, err)
  1122  
  1123  	// test purge when only a single target trie exist there
  1124  	err = forest.PurgeCacheExcept(updatedTrie2.RootHash())
  1125  	require.NoError(t, err)
  1126  	require.Equal(t, 1, forest.tries.Count())
  1127  }