From 79987722bb3d6e599f9a416ff93e59e403fe4ddb Mon Sep 17 00:00:00 2001
From: Brannon King <countprimes@gmail.com>
Date: Fri, 27 Aug 2021 10:54:26 -0400
Subject: [PATCH] fixed bug with block rollback; re-orgs should work now

---
 blockchain/chain.go                        | 6 ++++--
 claimtrie/claimtrie.go                     | 3 ++-
 claimtrie/claimtrie_test.go                | 7 ++++++-
 claimtrie/merkletrie/collapsedtrie_test.go | 2 +-
 claimtrie/merkletrie/ramtrie.go            | 8 +-------
 claimtrie/node/noderepo/pebble.go          | 2 +-
 goclean.sh                                 | 1 +
 7 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/blockchain/chain.go b/blockchain/chain.go
index 1511b770..8ffc4046 100644
--- a/blockchain/chain.go
+++ b/blockchain/chain.go
@@ -773,8 +773,10 @@ func (b *BlockChain) disconnectBlock(node *blockNode, block *btcutil.Block, view
 		return err
 	}
 
-	if err = b.claimTrie.ResetHeight(node.parent.height); err != nil {
-		return err
+	if b.claimTrie != nil {
+		if err = b.claimTrie.ResetHeight(node.parent.height); err != nil {
+			return err
+		}
 	}
 
 	// Prune fully spent entries and mark all entries in the view unmodified
diff --git a/claimtrie/claimtrie.go b/claimtrie/claimtrie.go
index 5cdd44e1..78fc40aa 100644
--- a/claimtrie/claimtrie.go
+++ b/claimtrie/claimtrie.go
@@ -318,12 +318,13 @@ func (ct *ClaimTrie) ResetHeight(height int32) error {
 	}
 
 	passedHashFork := ct.height >= param.ActiveParams.AllClaimsInMerkleForkHeight && height < param.ActiveParams.AllClaimsInMerkleForkHeight
-	ct.height = height
 	hash, err := ct.blockRepo.Get(height)
 	if err != nil {
 		return err
 	}
 
+	ct.height = height // keep this before the rebuild
+
 	if passedHashFork {
 		names = nil // force them to reconsider all names
 	}
diff --git a/claimtrie/claimtrie_test.go b/claimtrie/claimtrie_test.go
index 7b56372b..fe61f4f7 100644
--- a/claimtrie/claimtrie_test.go
+++ b/claimtrie/claimtrie_test.go
@@ -338,9 +338,14 @@ func BenchmarkClaimTrie_AppendBlock(b *testing.B) {
 	b.StopTimer()
 	ht := ct.height
 	h1 = *ct.MerkleHash()
-	ct.Close()
 	b.Logf("Running AppendBlock bench with %d names in %f sec. Height: %d, Hash: %s",
 		b.N, time.Since(start).Seconds(), ht, h1.String())
+
+	// a very important test of the functionality:
+	for ct.height > 0 {
+		err = ct.ResetHeight(ct.height - 1)
+		r.NoError(err)
+	}
 }
 
 func randomName() []byte {
diff --git a/claimtrie/merkletrie/collapsedtrie_test.go b/claimtrie/merkletrie/collapsedtrie_test.go
index 8efc3375..ce50b5d6 100644
--- a/claimtrie/merkletrie/collapsedtrie_test.go
+++ b/claimtrie/merkletrie/collapsedtrie_test.go
@@ -66,7 +66,7 @@ func TestNilNameHandling(t *testing.T) {
 }
 
 func TestCollapsedTriePerformance(t *testing.T) {
-	inserts := 10000 // increase this to 1M for more interesting results
+	inserts := 100000 // increase this to 1M for more interesting results
 	data := make([][]byte, inserts)
 	rand.Seed(42)
 	for i := 0; i < inserts; i++ {
diff --git a/claimtrie/merkletrie/ramtrie.go b/claimtrie/merkletrie/ramtrie.go
index c86aa407..5022b335 100644
--- a/claimtrie/merkletrie/ramtrie.go
+++ b/claimtrie/merkletrie/ramtrie.go
@@ -42,13 +42,7 @@ func (rt *RamTrie) SetRoot(h *chainhash.Hash) error {
 		return nil
 	}
 
-	// should technically clear the old trie first:
-	if rt.Nodes > 1 {
-		rt.Root = &collapsedVertex{key: make(KeyType, 0)}
-		rt.Nodes = 1
-		runtime.GC()
-	}
-
+	// should technically clear the old trie first, but this is abused for partial rebuilds so don't
 	return ErrFullRebuildRequired
 }
 
diff --git a/claimtrie/node/noderepo/pebble.go b/claimtrie/node/noderepo/pebble.go
index 29f5e05f..a988c486 100644
--- a/claimtrie/node/noderepo/pebble.go
+++ b/claimtrie/node/noderepo/pebble.go
@@ -86,7 +86,7 @@ func (repo *Pebble) DropChanges(name []byte, finalHeight int32) error {
 		return errors.Wrapf(err, "in load changes for %s", name)
 	}
 	i := 0
-	for ; i < len(changes); i++ {
+	for ; i < len(changes); i++ { // assuming changes are ordered by height
 		if changes[i].Height > finalHeight {
 			break
 		}
diff --git a/goclean.sh b/goclean.sh
index dad9f8f1..397407dd 100755
--- a/goclean.sh
+++ b/goclean.sh
@@ -10,6 +10,7 @@
 set -ex
 
 env GORACE="halt_on_error=1" go test -race -tags="rpctest" -covermode atomic -coverprofile=profile.cov ./...
+go test -bench=. -benchtime=4000x ./claimtrie/
 
 # Automatic checks
 golangci-lint run --deadline=10m --disable-all \