From 32d1715b922451adbd528fe72ed39fe85fd783af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?P=C3=A9ter=20Szil=C3=A1gyi?= <peterke@gmail.com>
Date: Fri, 22 Nov 2019 16:26:33 +0200
Subject: [PATCH] core/state/snapshot: switch quadratic cleanup to linear
 algorithm

---
 core/state/snapshot/difflayer.go |  9 ---------
 core/state/snapshot/snapshot.go  | 29 +++++++++++++++--------------
 2 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/core/state/snapshot/difflayer.go b/core/state/snapshot/difflayer.go
index b3bfb20630f4..7e8487ea840d 100644
--- a/core/state/snapshot/difflayer.go
+++ b/core/state/snapshot/difflayer.go
@@ -98,15 +98,6 @@ func (dl *diffLayer) Root() common.Hash {
 // Stale return whether this layer has become stale (was flattened across) or if
 // it's still live.
 func (dl *diffLayer) Stale() bool {
-	// If the parent is stale, mark this layer as stale too
-	if stale := dl.parent.Stale(); stale {
-		dl.lock.Lock()
-		defer dl.lock.Unlock()
-
-		dl.stale = true
-		return true
-	}
-	// Parent not stale, return only whether we are
 	dl.lock.RLock()
 	defer dl.lock.RUnlock()
 
diff --git a/core/state/snapshot/snapshot.go b/core/state/snapshot/snapshot.go
index d4dc1f55c3fb..edca5781d670 100644
--- a/core/state/snapshot/snapshot.go
+++ b/core/state/snapshot/snapshot.go
@@ -228,24 +228,25 @@ func (t *Tree) Cap(root common.Hash, layers int, memory uint64) error {
 		// Many layers requested to be retained, cap normally
 		t.cap(diff, layers, memory)
 	}
-	// Layers have been capped and paths invalidated, remove stales
+	// Remove any layer that is stale or links into a stale layer
+	children := make(map[common.Hash][]common.Hash)
 	for root, snap := range t.layers {
-		if snap.Stale() {
-			delete(t.layers, root)
+		if diff, ok := snap.(*diffLayer); ok {
+			parent := diff.parent.Root()
+			children[parent] = append(children[parent], root)
 		}
 	}
-	// Remove anything that links into a stale.
-	// TODO(karalabe): this is super suboptimal
-	for {
-		done := true
-		for root, snap := range t.layers {
-			if snap.Stale() {
-				delete(t.layers, root)
-				done = false
-			}
+	var remove func(root common.Hash)
+	remove = func(root common.Hash) {
+		delete(t.layers, root)
+		for _, child := range children[root] {
+			remove(child)
 		}
-		if done {
-			break
+		delete(children, root)
+	}
+	for root, snap := range t.layers {
+		if snap.Stale() {
+			remove(root)
 		}
 	}
 	return nil