From e04447554e31c570da0107b56c2b0fee6ba31eba Mon Sep 17 00:00:00 2001
From: Julius Andrikonis <julius.andrikonis@ext.iota.org>
Date: Mon, 23 Oct 2023 19:49:06 +0300
Subject: [PATCH] Some more tests and BugFixes for picking blocks for pruning
 algorythm

---
 .../statemanager/sm_gpa/state_manager_gpa.go  | 15 ++++++++---
 .../sm_gpa/state_manager_gpa_cob_test.go      | 26 ++++++++++++++++---
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/packages/chain/statemanager/sm_gpa/state_manager_gpa.go b/packages/chain/statemanager/sm_gpa/state_manager_gpa.go
index 01910181d3..5692237a82 100644
--- a/packages/chain/statemanager/sm_gpa/state_manager_gpa.go
+++ b/packages/chain/statemanager/sm_gpa/state_manager_gpa.go
@@ -586,7 +586,7 @@ func (smT *stateManagerGPA) commitStateDraft(stateDraft state.StateDraft) state.
 	stateIndex := block.StateIndex()
 	smT.metrics.BlockIndexCommitted(stateIndex)
 	if smT.pruningNeeded() {
-		smT.pruneStore(block.PreviousL1Commitment(), stateIndex)
+		smT.pruneStore(block.PreviousL1Commitment(), stateIndex-1)
 	}
 	smT.output.addBlockCommitted(stateIndex, block.L1Commitment())
 	return block
@@ -697,6 +697,12 @@ func (smT *stateManagerGPA) updateChainOfBlocks(commitment *state.L1Commitment,
 			bi, err = GetPreviousBlockInfoFun(bi)
 		}
 	}
+	if err == nil && bi != nil {
+		for lastKnownBi != nil && lastKnownBi.blockIndex > bi.blockIndex {
+			_ = smT.chainOfBlocks.RemoveEnd()
+			lastKnownBi = GetLastKnownBlockInfoFun()
+		}
+	}
 	// Try to find a place to merge newest block chain with currently known block chain: `bi.trieRoot.Equals(lastKnownBi.trieRoot)``
 	for err == nil && bi != nil && lastKnownBi != nil && !bi.trieRoot.Equals(lastKnownBi.trieRoot) && smT.store.HasTrieRoot(bi.trieRoot) {
 		// Normally, no iteration of this cycle should occur: once a common index
@@ -732,10 +738,12 @@ func (smT *stateManagerGPA) updateChainOfBlocks(commitment *state.L1Commitment,
 			return
 		}
 		smT.chainOfBlocks = cob
+	} else if bi == nil { // origin block has been reached
+		smT.chainOfBlocks = cob
 	} else if bi.trieRoot.Equals(lastKnownBi.trieRoot) { // Here is the the place to merge newest block chain with currently known block chain
 		// Normally newest blocks chain should contain only several (usually, 1)
 		// block indexes and currently known block chain should contain at least
-		// `PruningMinStatesToKeep` indexes and on a sudden enabling of pruning
+		// `PruningMinStatesToKeep` indexes, but on a sudden enabling of pruning
 		// might contain millions of them. Therefore it is more effective to copy
 		// newest block chain to the currently known one compared to doing it
 		// the other way round. Let's merge them this way.
@@ -760,8 +768,7 @@ func (smT *stateManagerGPA) updateChainOfBlocks(commitment *state.L1Commitment,
 			smT.log.Errorf("Failed to obtain previous block info: %v", err)
 			return
 		}
-	} else { // bi == nil, which means that origin block has been reached or
-		// smT.store.HasTrieRoot(bi.trieRoot), which means that this block
+	} else { // !smT.store.HasTrieRoot(bi.trieRoot), which means that this block
 		// has already been pruned from the store.
 		smT.chainOfBlocks = cob
 	}
diff --git a/packages/chain/statemanager/sm_gpa/state_manager_gpa_cob_test.go b/packages/chain/statemanager/sm_gpa/state_manager_gpa_cob_test.go
index 863d248125..bad308cbd0 100644
--- a/packages/chain/statemanager/sm_gpa/state_manager_gpa_cob_test.go
+++ b/packages/chain/statemanager/sm_gpa/state_manager_gpa_cob_test.go
@@ -130,10 +130,7 @@ func TestChainOfBlocksMergeAllSomeHistory(t *testing.T) {
 	runTestChainOfBlocks(t, log, bf, store, sm, blocksToCommit, blocksToPrune, blocksInChain, blocksExpected)
 }
 
-func TestChainOfBlocksMergeMiddleFullHistory(t *testing.T) {
-	totalBlocks := 15
-	branchFrom := 9
-	branchLength := 5
+func testChainOfBlocksMergeMiddleFullHistory(t *testing.T, totalBlocks, branchFrom, branchLength int) {
 	log, bf, store, sm := initTestChainOfBlocks(t)
 	originalBlocks := bf.GetBlocks(totalBlocks, 1)
 	branchBlocks := bf.GetBlocksFrom(branchLength, 1, originalBlocks[branchFrom].L1Commitment(), 2)
@@ -143,6 +140,27 @@ func TestChainOfBlocksMergeMiddleFullHistory(t *testing.T) {
 	runTestChainOfBlocks(t, log, bf, store, sm, blocksToCommit, nil, blocksInChain, prependOriginBlock(bf, blocksToCommit))
 }
 
+func TestChainOfBlocksMergeMiddleFullHistory(t *testing.T) {
+	totalBlocks := 15
+	branchFrom := 9
+	branchLength := 5
+	testChainOfBlocksMergeMiddleFullHistory(t, totalBlocks, branchFrom, branchLength)
+}
+
+func TestChainOfBlocksMergeMiddleFullHistoryLonger(t *testing.T) {
+	totalBlocks := 15
+	branchFrom := 9
+	branchLength := 10
+	testChainOfBlocksMergeMiddleFullHistory(t, totalBlocks, branchFrom, branchLength)
+}
+
+func TestChainOfBlocksMergeMiddleFullHistoryShorter(t *testing.T) {
+	totalBlocks := 15
+	branchFrom := 9
+	branchLength := 3
+	testChainOfBlocksMergeMiddleFullHistory(t, totalBlocks, branchFrom, branchLength)
+}
+
 func TestChainOfBlocksMergeMiddleSomeHistory(t *testing.T) {
 	totalBlocks := 15
 	branchFrom := 9