From 5948c664e2e4f205932d7db254d45bf6d7f6c9b1 Mon Sep 17 00:00:00 2001
From: EclesioMeloJunior <eclesiomelo.1@gmail.com>
Date: Thu, 11 Jan 2024 08:44:32 -0400
Subject: [PATCH] chore: include tests for nested transactions

---
 lib/runtime/storage/trie.go      |   7 +--
 lib/runtime/storage/trie_test.go | 103 +++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 6 deletions(-)

diff --git a/lib/runtime/storage/trie.go b/lib/runtime/storage/trie.go
index 1dbc016fcf8..568d2491663 100644
--- a/lib/runtime/storage/trie.go
+++ b/lib/runtime/storage/trie.go
@@ -62,7 +62,7 @@ func (t *TrieState) Rollback() {
 	defer t.mtx.Unlock()
 
 	if len(t.transactions) < 1 {
-		panic("no transactions to commit")
+		panic("no transactions to rollback")
 	}
 
 	t.transactions = t.transactions[:len(t.transactions)-1]
@@ -112,11 +112,6 @@ func (t *TrieState) Put(key, value []byte) (err error) {
 func (t *TrieState) Get(key []byte) []byte {
 	t.mtx.RLock()
 	defer t.mtx.RUnlock()
-
-	if t.getCurrentTrie() == nil {
-		panic("trie cannot be nil at trie state")
-	}
-
 	return t.getCurrentTrie().Get(key)
 }
 
diff --git a/lib/runtime/storage/trie_test.go b/lib/runtime/storage/trie_test.go
index ba498c1ad98..2f17b08b871 100644
--- a/lib/runtime/storage/trie_test.go
+++ b/lib/runtime/storage/trie_test.go
@@ -223,6 +223,109 @@ func TestTrieState_RollbackStorageTransaction(t *testing.T) {
 	require.Equal(t, []byte(testCases[0]), val)
 }
 
+func TestTrieState_NestedTransactions(t *testing.T) {
+	cases := map[string]struct {
+		createTrieState func() *TrieState
+		assert          func(*testing.T, *TrieState)
+	}{
+		"commiting_and_rollback_on_nested_transactions": {
+			createTrieState: func() *TrieState {
+				ts := NewTrieState(trie.NewEmptyTrie())
+
+				ts.Put([]byte("key-1"), []byte("value-1"))
+				ts.Put([]byte("key-2"), []byte("value-2"))
+				ts.Put([]byte("key-3"), []byte("value-3"))
+
+				{
+					ts.StartTransaction()
+					ts.Put([]byte("key-4"), []byte("value-4"))
+
+					{
+						ts.StartTransaction()
+						ts.Delete([]byte("key-3"))
+						ts.Commit() // commit the most nested transaction
+					}
+
+					// rollback this transaction will discard the modifications
+					// made by the most nested transactions so this original trie
+					// should not be affected
+					ts.Rollback()
+				}
+				return ts
+			},
+			assert: func(t *testing.T, ts *TrieState) {
+				require.NotNil(t, ts.Get([]byte("key-1")))
+				require.NotNil(t, ts.Get([]byte("key-2")))
+				require.NotNil(t, ts.Get([]byte("key-3")))
+
+				require.Nil(t, ts.Get([]byte("key-4")))
+				require.Zero(t, len(ts.transactions))
+			},
+		},
+		"commiting_all_nested_transactions": {
+			createTrieState: func() *TrieState {
+				ts := NewTrieState(trie.NewEmptyTrie())
+				{
+					ts.StartTransaction()
+					ts.Put([]byte("key-1"), []byte("value-1"))
+					{
+						ts.StartTransaction()
+						ts.Put([]byte("key-2"), []byte("value-2"))
+						{
+							ts.StartTransaction()
+							ts.Put([]byte("key-3"), []byte("value-3"))
+							{
+								ts.StartTransaction()
+								ts.Put([]byte("key-4"), []byte("value-4"))
+								{
+									ts.StartTransaction()
+									ts.Delete([]byte("key-3"))
+									ts.Commit()
+								}
+								ts.Commit()
+							}
+							ts.Commit()
+						}
+						ts.Commit()
+					}
+					ts.Commit()
+				}
+				return ts
+			},
+			assert: func(t *testing.T, ts *TrieState) {
+				require.NotNil(t, ts.Get([]byte("key-1")))
+				require.NotNil(t, ts.Get([]byte("key-2")))
+				require.NotNil(t, ts.Get([]byte("key-4")))
+				require.Zero(t, len(ts.transactions))
+			},
+		},
+		"rollback_without_transaction_should_panic": {
+			createTrieState: func() *TrieState {
+				return NewTrieState(trie.NewEmptyTrie())
+			},
+			assert: func(t *testing.T, ts *TrieState) {
+				require.PanicsWithValue(t, "no transactions to rollback", func() { ts.Rollback() })
+			},
+		},
+		"commit_without_transaction_should_panic": {
+			createTrieState: func() *TrieState {
+				return NewTrieState(trie.NewEmptyTrie())
+			},
+			assert: func(t *testing.T, ts *TrieState) {
+				require.PanicsWithValue(t, "no transactions to commit", func() { ts.Commit() })
+			},
+		},
+	}
+
+	for tname, tt := range cases {
+		tt := tt
+		t.Run(tname, func(t *testing.T) {
+			ts := tt.createTrieState()
+			tt.assert(t, ts)
+		})
+	}
+}
+
 func TestTrieState_DeleteChildLimit(t *testing.T) {
 	ts := &TrieState{currentTrie: trie.NewEmptyTrie()}
 	child := trie.NewEmptyTrie()