From 5758d1fb1162f14ebb2420017d83181e576e3d88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Tue, 23 Aug 2022 21:17:12 +0300
Subject: [PATCH] core/state, trie: fix trie flush order for proper pruning

---
 core/state/statedb_test.go | 40 ++++++++++++++++++++++++++++++++++++++
 trie/database.go           | 19 +++++++++++++++---
 2 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/core/state/statedb_test.go b/core/state/statedb_test.go
index 092a4fb87..6fe36a7ec 100644
--- a/core/state/statedb_test.go
+++ b/core/state/statedb_test.go
@@ -914,3 +914,43 @@ func TestStateDBAccessList(t *testing.T) {
 		t.Fatalf("expected empty, got %d", got)
 	}
 }
+
+// Tests that account and storage tries are flushed in the correct order and that
+// no data loss occurs.
+func TestFlushOrderDataLoss(t *testing.T) {
+	// Create a state trie with many accounts and slots
+	var (
+		memdb    = rawdb.NewMemoryDatabase()
+		statedb  = NewDatabase(memdb)
+		state, _ = New(common.Hash{}, statedb, nil)
+	)
+	for a := byte(0); a < 10; a++ {
+		state.CreateAccount(common.Address{a})
+		for s := byte(0); s < 10; s++ {
+			state.SetState(common.Address{a}, common.Hash{a, s}, common.Hash{a, s})
+		}
+	}
+	root, err := state.Commit(false)
+	if err != nil {
+		t.Fatalf("failed to commit state trie: %v", err)
+	}
+	statedb.TrieDB().Reference(root, common.Hash{})
+	if err := statedb.TrieDB().Cap(1024); err != nil {
+		t.Fatalf("failed to cap trie dirty cache: %v", err)
+	}
+	if err := statedb.TrieDB().Commit(root, false, nil); err != nil {
+		t.Fatalf("failed to commit state trie: %v", err)
+	}
+	// Reopen the state trie from flushed disk and verify it
+	state, err = New(root, NewDatabase(memdb), nil)
+	if err != nil {
+		t.Fatalf("failed to reopen state trie: %v", err)
+	}
+	for a := byte(0); a < 10; a++ {
+		for s := byte(0); s < 10; s++ {
+			if have := state.GetState(common.Address{a}, common.Hash{a, s}); have != (common.Hash{a, s}) {
+				t.Errorf("account %d: slot %d: state mismatch: have %x, want %x", a, s, have, common.Hash{a, s})
+			}
+		}
+	}
+}
diff --git a/trie/database.go b/trie/database.go
index 8c154ba96..b10bbca9b 100644
--- a/trie/database.go
+++ b/trie/database.go
@@ -776,9 +776,22 @@ func (db *Database) Update(nodes *MergedNodeSet) error {
 
 	// Insert dirty nodes into the database. In the same tree, it must be
 	// ensured that children are inserted first, then parent so that children
-	// can be linked with their parent correctly. The order of writing between
-	// different tries(account trie, storage tries) is not required.
-	for owner, subset := range nodes.sets {
+	// can be linked with their parent correctly.
+	//
+	// Note, the storage tries must be flushed before the account trie to
+	// retain the invariant that children go into the dirty cache first.
+	var order []common.Hash
+	for owner := range nodes.sets {
+		if owner == (common.Hash{}) {
+			continue
+		}
+		order = append(order, owner)
+	}
+	if _, ok := nodes.sets[common.Hash{}]; ok {
+		order = append(order, common.Hash{})
+	}
+	for _, owner := range order {
+		subset := nodes.sets[owner]
 		for _, path := range subset.paths {
 			n, ok := subset.nodes[path]
 			if !ok {