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

     1  package flattener_test
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/require"
     8  
     9  	"github.com/onflow/flow-go/ledger"
    10  	"github.com/onflow/flow-go/ledger/common/testutils"
    11  	"github.com/onflow/flow-go/ledger/complete/mtrie/flattener"
    12  	"github.com/onflow/flow-go/ledger/complete/mtrie/node"
    13  	"github.com/onflow/flow-go/ledger/complete/mtrie/trie"
    14  )
    15  
    16  func TestEmptyTrie(t *testing.T) {
    17  	emptyTrie := trie.NewEmptyMTrie()
    18  
    19  	itr := flattener.NewNodeIterator(emptyTrie.RootNode())
    20  	require.True(t, nil == itr.Value()) // initial iterator should return nil
    21  
    22  	require.False(t, itr.Next())
    23  	require.Equal(t, emptyTrie.RootNode(), itr.Value())
    24  	require.Equal(t, emptyTrie.RootNode(), itr.Value()) // test that recalling twice has no problem
    25  	require.False(t, itr.Next())
    26  	require.True(t, nil == itr.Value())
    27  }
    28  
    29  func TestTrieWithOneNode(t *testing.T) {
    30  	emptyTrie := trie.NewEmptyMTrie()
    31  
    32  	// key: 0000...
    33  	p1 := testutils.PathByUint8(1)
    34  	v1 := testutils.LightPayload8('A', 'a')
    35  
    36  	paths := []ledger.Path{p1}
    37  	payloads := []ledger.Payload{*v1}
    38  
    39  	testTrie, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
    40  	require.NoError(t, err)
    41  
    42  	itr := flattener.NewNodeIterator(testTrie.RootNode())
    43  	// by default Value returns nil, even if there is node
    44  	require.Nil(t, itr.Value())
    45  
    46  	// has one node
    47  	require.Equal(t, true, itr.Next())
    48  	require.NotNil(t, itr.Value())
    49  
    50  	// no more node
    51  	require.Equal(t, false, itr.Next())
    52  	require.Nil(t, itr.Value())
    53  }
    54  
    55  func TestPopulatedTrie(t *testing.T) {
    56  	emptyTrie := trie.NewEmptyMTrie()
    57  
    58  	// key: 0000...
    59  	p1 := testutils.PathByUint8(0)
    60  	v1 := testutils.LightPayload8('A', 'a')
    61  
    62  	// key: 0100....
    63  	p2 := testutils.PathByUint8(64)
    64  	v2 := testutils.LightPayload8('B', 'b')
    65  
    66  	paths := []ledger.Path{p1, p2}
    67  	payloads := []ledger.Payload{*v1, *v2}
    68  
    69  	testTrie, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
    70  	require.NoError(t, err)
    71  
    72  	for itr := flattener.NewNodeIterator(testTrie.RootNode()); itr.Next(); {
    73  		fmt.Println(itr.Value().FmtStr("", ""))
    74  		fmt.Println()
    75  	}
    76  
    77  	itr := flattener.NewNodeIterator(testTrie.RootNode())
    78  
    79  	require.True(t, itr.Next())
    80  	p1_leaf := itr.Value()
    81  	require.Equal(t, p1, *p1_leaf.Path())
    82  	require.Equal(t, v1, p1_leaf.Payload())
    83  
    84  	require.True(t, itr.Next())
    85  	p2_leaf := itr.Value()
    86  	require.Equal(t, p2, *p2_leaf.Path())
    87  	require.Equal(t, v2, p2_leaf.Payload())
    88  
    89  	require.True(t, itr.Next())
    90  	p_parent := itr.Value()
    91  	require.Equal(t, p1_leaf, p_parent.LeftChild())
    92  	require.Equal(t, p2_leaf, p_parent.RightChild())
    93  
    94  	require.True(t, itr.Next())
    95  	root := itr.Value()
    96  	require.Equal(t, testTrie.RootNode(), root)
    97  	require.Equal(t, p_parent, root.LeftChild())
    98  	require.True(t, nil == root.RightChild())
    99  
   100  	require.False(t, itr.Next())
   101  	require.True(t, nil == itr.Value())
   102  }
   103  
   104  func TestUniqueNodeIterator(t *testing.T) {
   105  	t.Run("empty trie", func(t *testing.T) {
   106  		emptyTrie := trie.NewEmptyMTrie()
   107  
   108  		// visitedNodes is nil
   109  		itr := flattener.NewUniqueNodeIterator(emptyTrie.RootNode(), nil)
   110  		require.False(t, itr.Next())
   111  		require.True(t, nil == itr.Value()) // initial iterator should return nil
   112  
   113  		// visitedNodes is empty map
   114  		visitedNodes := make(map[*node.Node]uint64)
   115  		itr = flattener.NewUniqueNodeIterator(emptyTrie.RootNode(), visitedNodes)
   116  		require.False(t, itr.Next())
   117  		require.True(t, nil == itr.Value()) // initial iterator should return nil
   118  	})
   119  
   120  	t.Run("trie", func(t *testing.T) {
   121  		emptyTrie := trie.NewEmptyMTrie()
   122  
   123  		// key: 0000...
   124  		p1 := testutils.PathByUint8(0)
   125  		v1 := testutils.LightPayload8('A', 'a')
   126  
   127  		// key: 0100....
   128  		p2 := testutils.PathByUint8(64)
   129  		v2 := testutils.LightPayload8('B', 'b')
   130  
   131  		paths := []ledger.Path{p1, p2}
   132  		payloads := []ledger.Payload{*v1, *v2}
   133  
   134  		updatedTrie, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
   135  		require.NoError(t, err)
   136  
   137  		//              n4
   138  		//             /
   139  		//            /
   140  		//          n3
   141  		//        /     \
   142  		//      /         \
   143  		//   n1 (p1/v1)     n2 (p2/v2)
   144  		//
   145  
   146  		expectedNodes := []*node.Node{
   147  			updatedTrie.RootNode().LeftChild().LeftChild(),  // n1
   148  			updatedTrie.RootNode().LeftChild().RightChild(), // n2
   149  			updatedTrie.RootNode().LeftChild(),              // n3
   150  			updatedTrie.RootNode(),                          // n4
   151  		}
   152  
   153  		// visitedNodes is nil
   154  		i := 0
   155  		for itr := flattener.NewUniqueNodeIterator(updatedTrie.RootNode(), nil); itr.Next(); {
   156  			n := itr.Value()
   157  			require.True(t, i < len(expectedNodes))
   158  			require.Equal(t, expectedNodes[i], n)
   159  			i++
   160  		}
   161  		require.Equal(t, i, len(expectedNodes))
   162  
   163  		// visitedNodes is not nil, but it's pointless for iterating a single trie because
   164  		// there isn't any shared sub-trie.
   165  		visitedNodes := make(map[*node.Node]uint64)
   166  		i = 0
   167  		for itr := flattener.NewUniqueNodeIterator(updatedTrie.RootNode(), visitedNodes); itr.Next(); {
   168  			n := itr.Value()
   169  			visitedNodes[n] = uint64(i)
   170  
   171  			require.True(t, i < len(expectedNodes))
   172  			require.Equal(t, expectedNodes[i], n)
   173  			i++
   174  		}
   175  		require.Equal(t, i, len(expectedNodes))
   176  	})
   177  
   178  	t.Run("forest", func(t *testing.T) {
   179  
   180  		// tries is a slice of mtries to guarantee order.
   181  		var tries []*trie.MTrie
   182  
   183  		emptyTrie := trie.NewEmptyMTrie()
   184  
   185  		// key: 0000...
   186  		p1 := testutils.PathByUint8(0)
   187  		v1 := testutils.LightPayload8('A', 'a')
   188  
   189  		// key: 0100....
   190  		p2 := testutils.PathByUint8(64)
   191  		v2 := testutils.LightPayload8('B', 'b')
   192  
   193  		paths := []ledger.Path{p1, p2}
   194  		payloads := []ledger.Payload{*v1, *v2}
   195  
   196  		trie1, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
   197  		require.NoError(t, err)
   198  
   199  		// trie1
   200  		//              n4
   201  		//             /
   202  		//            /
   203  		//          n3
   204  		//        /     \
   205  		//      /         \
   206  		//   n1 (p1/v1)     n2 (p2/v2)
   207  		//
   208  
   209  		tries = append(tries, trie1)
   210  
   211  		// New trie reuses its parent's left sub-trie.
   212  
   213  		// key: 1000...
   214  		p3 := testutils.PathByUint8(128)
   215  		v3 := testutils.LightPayload8('C', 'c')
   216  
   217  		// key: 1100....
   218  		p4 := testutils.PathByUint8(192)
   219  		v4 := testutils.LightPayload8('D', 'd')
   220  
   221  		paths = []ledger.Path{p3, p4}
   222  		payloads = []ledger.Payload{*v3, *v4}
   223  
   224  		trie2, _, err := trie.NewTrieWithUpdatedRegisters(trie1, paths, payloads, true)
   225  		require.NoError(t, err)
   226  
   227  		// trie2
   228  		//              n8
   229  		//             /   \
   230  		//            /      \
   231  		//          n3       n7
   232  		//       (shared)   /   \
   233  		//                /       \
   234  		//              n5         n6
   235  		//            (p3/v3)    (p4/v4)
   236  
   237  		tries = append(tries, trie2)
   238  
   239  		// New trie reuses its parent's right sub-trie, and left sub-trie's leaf node.
   240  
   241  		// key: 0000...
   242  		v5 := testutils.LightPayload8('E', 'e')
   243  
   244  		paths = []ledger.Path{p1}
   245  		payloads = []ledger.Payload{*v5}
   246  
   247  		trie3, _, err := trie.NewTrieWithUpdatedRegisters(trie2, paths, payloads, true)
   248  		require.NoError(t, err)
   249  
   250  		// trie3
   251  		//              n11
   252  		//             /   \
   253  		//            /      \
   254  		//          n10       n7
   255  		//         /   \    (shared)
   256  		//       /       \
   257  		//     n9         n2
   258  		//  (p1/v5)    (shared)
   259  
   260  		tries = append(tries, trie3)
   261  
   262  		expectedNodes := []*node.Node{
   263  			// unique nodes from trie1
   264  			trie1.RootNode().LeftChild().LeftChild(),  // n1
   265  			trie1.RootNode().LeftChild().RightChild(), // n2
   266  			trie1.RootNode().LeftChild(),              // n3
   267  			trie1.RootNode(),                          // n4
   268  			// unique nodes from trie2
   269  			trie2.RootNode().RightChild().LeftChild(),  // n5
   270  			trie2.RootNode().RightChild().RightChild(), // n6
   271  			trie2.RootNode().RightChild(),              // n7
   272  			trie2.RootNode(),                           // n8
   273  			// unique nodes from trie3
   274  			trie3.RootNode().LeftChild().LeftChild(), // n9
   275  			trie3.RootNode().LeftChild(),             // n10
   276  			trie3.RootNode(),                         // n11
   277  		}
   278  
   279  		// Use visitedNodes to prevent revisiting shared sub-tries.
   280  		visitedNodes := make(map[*node.Node]uint64)
   281  		i := 0
   282  		for _, trie := range tries {
   283  			for itr := flattener.NewUniqueNodeIterator(trie.RootNode(), visitedNodes); itr.Next(); {
   284  				n := itr.Value()
   285  				visitedNodes[n] = uint64(i)
   286  
   287  				require.True(t, i < len(expectedNodes))
   288  				require.Equal(t, expectedNodes[i], n)
   289  				i++
   290  			}
   291  		}
   292  		require.Equal(t, i, len(expectedNodes))
   293  	})
   294  
   295  	t.Run("subtries", func(t *testing.T) {
   296  
   297  		emptyTrie := trie.NewEmptyMTrie()
   298  
   299  		// key: 0000...
   300  		p1 := testutils.PathByUint8(0)
   301  		v1 := testutils.LightPayload8('A', 'a')
   302  
   303  		// key: 0100....
   304  		p2 := testutils.PathByUint8(64)
   305  		v2 := testutils.LightPayload8('B', 'b')
   306  
   307  		paths := []ledger.Path{p1, p2}
   308  		payloads := []ledger.Payload{*v1, *v2}
   309  
   310  		trie1, _, err := trie.NewTrieWithUpdatedRegisters(emptyTrie, paths, payloads, true)
   311  		require.NoError(t, err)
   312  
   313  		// trie1
   314  		//              n4
   315  		//             /
   316  		//            /
   317  		//          n3
   318  		//        /     \
   319  		//      /         \
   320  		//   n1 (p1/v1)     n2 (p2/v2)
   321  		//
   322  
   323  		// New trie reuses its parent's left sub-trie.
   324  
   325  		// key: 1000...
   326  		p3 := testutils.PathByUint8(128)
   327  		v3 := testutils.LightPayload8('C', 'c')
   328  
   329  		// key: 1100....
   330  		p4 := testutils.PathByUint8(192)
   331  		v4 := testutils.LightPayload8('D', 'd')
   332  
   333  		paths = []ledger.Path{p3, p4}
   334  		payloads = []ledger.Payload{*v3, *v4}
   335  
   336  		trie2, _, err := trie.NewTrieWithUpdatedRegisters(trie1, paths, payloads, true)
   337  		require.NoError(t, err)
   338  
   339  		// trie2
   340  		//              n8
   341  		//             /   \
   342  		//            /      \
   343  		//          n3       n7
   344  		//       (shared)   /   \
   345  		//                /       \
   346  		//              n5         n6
   347  		//            (p3/v3)    (p4/v4)
   348  
   349  		// New trie reuses its parent's right sub-trie, and left sub-trie's leaf node.
   350  
   351  		// key: 0000...
   352  		v5 := testutils.LightPayload8('E', 'e')
   353  
   354  		paths = []ledger.Path{p1}
   355  		payloads = []ledger.Payload{*v5}
   356  
   357  		trie3, _, err := trie.NewTrieWithUpdatedRegisters(trie2, paths, payloads, true)
   358  		require.NoError(t, err)
   359  
   360  		// trie3
   361  		//              n11
   362  		//             /   \
   363  		//            /      \
   364  		//          n10       n7
   365  		//         /   \    (shared)
   366  		//       /       \
   367  		//     n9         n2
   368  		//  (p1/v5)    (shared)
   369  
   370  		leftSubtries := []*node.Node{
   371  			trie1.RootNode().LeftChild(),
   372  			trie2.RootNode().LeftChild(),
   373  			trie3.RootNode().LeftChild(),
   374  		}
   375  
   376  		expectedNodesInLeftSubtries := []*node.Node{
   377  			// unique nodes from trie1
   378  			trie1.RootNode().LeftChild().LeftChild(),  // n1
   379  			trie1.RootNode().LeftChild().RightChild(), // n2
   380  			trie1.RootNode().LeftChild(),              // n3
   381  			// unique nodes from trie3
   382  			trie3.RootNode().LeftChild().LeftChild(), // n9
   383  			trie3.RootNode().LeftChild(),             // n10
   384  		}
   385  
   386  		rightSubtries := []*node.Node{
   387  			trie1.RootNode().RightChild(),
   388  			trie2.RootNode().RightChild(),
   389  			trie3.RootNode().RightChild(),
   390  		}
   391  
   392  		expectedNodesInRightSubtries := []*node.Node{
   393  			// unique nodes from trie2
   394  			trie2.RootNode().RightChild().LeftChild(),  // n5
   395  			trie2.RootNode().RightChild().RightChild(), // n6
   396  			trie2.RootNode().RightChild(),              // n7
   397  		}
   398  
   399  		testcases := []struct {
   400  			roots         []*node.Node
   401  			expectedNodes []*node.Node
   402  		}{
   403  			{leftSubtries, expectedNodesInLeftSubtries},
   404  			{rightSubtries, expectedNodesInRightSubtries},
   405  		}
   406  
   407  		for _, tc := range testcases {
   408  			visitedNodes := make(map[*node.Node]uint64)
   409  			i := 0
   410  			for _, n := range tc.roots {
   411  				for itr := flattener.NewUniqueNodeIterator(n, visitedNodes); itr.Next(); {
   412  					n := itr.Value()
   413  					visitedNodes[n] = uint64(i)
   414  
   415  					require.True(t, i < len(tc.expectedNodes))
   416  					require.Equal(t, tc.expectedNodes[i], n)
   417  					i++
   418  				}
   419  			}
   420  			require.Equal(t, i, len(tc.expectedNodes))
   421  		}
   422  	})
   423  }