From bda174f38fc332eb37e5210cc5bf5e677e9c8d76 Mon Sep 17 00:00:00 2001
From: Dave Collins <davec@conformal.com>
Date: Mon, 23 Aug 2021 06:29:07 -0500
Subject: [PATCH] chaingen: Add spendable coinbase out snapshots.

This adds the ability to snapshot the current spendable coinbase output
state of the generator with a name and restore it later.  This will
allow more efficient tests to be created since they will be able to use
invalidation along with restoring the spendable outputs state to avoid
the need to go through the entire vote activation process multiple
times.
---
 blockchain/chaingen/generator.go | 52 ++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/blockchain/chaingen/generator.go b/blockchain/chaingen/generator.go
index b361e364..8f534082 100644
--- a/blockchain/chaingen/generator.go
+++ b/blockchain/chaingen/generator.go
@@ -155,6 +155,13 @@ func (t stakeTicketSorter) Less(i, j int) bool {
 	return bytes.Compare(iHash, jHash) < 0
 }
 
+// spendableOutsSnap houses a snapshot of the spendable outputs state.  It is
+// useful to allow callers to reset back to known good points.
+type spendableOutsSnap struct {
+	spendableOuts     [][]SpendableOut
+	prevCollectedHash chainhash.Hash
+}
+
 // Generator houses state used to ease the process of generating test blocks
 // that build from one another along with housing other useful things such as
 // available spendable outputs and generic payment scripts used throughout the
@@ -174,6 +181,7 @@ type Generator struct {
 	// Used for tracking spendable coinbase outputs.
 	spendableOuts     [][]SpendableOut
 	prevCollectedHash chainhash.Hash
+	spendableOutSnaps map[string]spendableOutsSnap
 
 	// Used for tracking the live ticket pool and revocations.
 	originalParents map[chainhash.Hash]chainhash.Hash
@@ -207,6 +215,7 @@ func MakeGenerator(params *chaincfg.Params) (Generator, error) {
 		blockHeights:        map[chainhash.Hash]uint32{genesisHash: 0},
 		blocksByName:        map[string]*wire.MsgBlock{"genesis": genesis},
 		blockNames:          map[chainhash.Hash]string{genesisHash: "genesis"},
+		spendableOutSnaps:   make(map[string]spendableOutsSnap),
 		p2shOpTrueAddr:      p2shOpTrueAddr,
 		p2shOpTrueScriptVer: p2shOpTrueScriptVer,
 		p2shOpTrueScript:    p2shOpTrueScript,
@@ -2506,6 +2515,49 @@ func (g *Generator) SaveSpendableCoinbaseOuts() {
 	}
 }
 
+// dupSpendableOuts returns a deep copy of the passed spendable outputs.
+func dupSpendableOuts(spendableOuts [][]SpendableOut) [][]SpendableOut {
+	result := make([][]SpendableOut, 0, len(spendableOuts))
+	for _, outs := range spendableOuts {
+		outsCopy := make([]SpendableOut, len(outs))
+		copy(outsCopy, outs)
+		result = append(result, outsCopy)
+	}
+	return result
+}
+
+// SnapshotCoinbaseOuts saves the current state of the spendable coinbase
+// outputs with the given snapshot name so that it can be restored later by
+// name.  This is primarily useful for tests that want to be able to undo a
+// bunch of blocks via invalidation and reset the spendable outputs back to a
+// matching known good state.
+//
+// This will panic if the specified snapshot name already exists.
+func (g *Generator) SnapshotCoinbaseOuts(snapName string) {
+	if _, ok := g.spendableOutSnaps[snapName]; ok {
+		panic(fmt.Sprintf("snapshot name %s already exists", snapName))
+	}
+	g.spendableOutSnaps[snapName] = spendableOutsSnap{
+		spendableOuts:     dupSpendableOuts(g.spendableOuts),
+		prevCollectedHash: g.prevCollectedHash,
+	}
+}
+
+// RestoreCoinbaseOutsSnapshot restores the state of the spendable coinbase
+// outputs saved with the given snapshot name.  This is primarily useful for
+// tests that want to be able to undo a bunch of blocks via invalidation and
+// reset the spendable outputs back to a matching known good state.
+//
+// This will panic if the specified snapshot name does not exist.
+func (g *Generator) RestoreCoinbaseOutsSnapshot(snapName string) {
+	snap, ok := g.spendableOutSnaps[snapName]
+	if !ok {
+		panic(fmt.Sprintf("snapshot name %s does not exist", snapName))
+	}
+	g.spendableOuts = dupSpendableOuts(snap.spendableOuts)
+	g.prevCollectedHash = snap.prevCollectedHash
+}
+
 // AssertTipHeight panics if the current tip block associated with the generator
 // does not have the specified height.
 func (g *Generator) AssertTipHeight(expected uint32) {