From 4945162dd790a142a9400d2256d5fe40ae814e5d Mon Sep 17 00:00:00 2001
From: Alex Sharov <AskAlexSharov@gmail.com>
Date: Mon, 15 Aug 2022 10:27:08 +0700
Subject: [PATCH] erigon22: unwind code #587

---
 state/aggregator22.go | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/state/aggregator22.go b/state/aggregator22.go
index 9d2ce7a4f..fd2afbdb7 100644
--- a/state/aggregator22.go
+++ b/state/aggregator22.go
@@ -22,7 +22,6 @@ import (
 	math2 "math"
 
 	"github.com/RoaringBitmap/roaring/roaring64"
-	libcommon "github.com/ledgerwatch/erigon-lib/common"
 	"github.com/ledgerwatch/erigon-lib/etl"
 	"github.com/ledgerwatch/erigon-lib/kv"
 )
@@ -37,6 +36,7 @@ type Aggregator22 struct {
 	tracesFrom      *InvertedIndex
 	tracesTo        *InvertedIndex
 	txNum           uint64
+	logPrefix       string
 	rwTx            kv.RwTx
 }
 
@@ -103,6 +103,8 @@ func (a *Aggregator22) Close() {
 	}
 }
 
+func (a *Aggregator22) SetLogPrefix(v string) { a.logPrefix = v }
+
 func (a *Aggregator22) SetTx(tx kv.RwTx) {
 	a.rwTx = tx
 	a.accounts.SetTx(tx)
@@ -285,11 +287,11 @@ func (a *Aggregator22) Unwind(ctx context.Context, txUnwindTo uint64, stateLoad
 		return nil
 	}
 
-	changes := etl.NewCollector("", "", etl.NewOldestEntryBuffer(etl.BufferOptimalSize))
-	defer changes.Close()
+	stateChanges := etl.NewCollector(a.logPrefix, "", etl.NewOldestEntryBuffer(etl.BufferOptimalSize))
+	defer stateChanges.Close()
 
 	if err := a.accounts.pruneF(step, txUnwindTo, math2.MaxUint64, func(txNum uint64, k, v []byte) error {
-		if err := changes.Collect(libcommon.Copy(k), libcommon.Copy(v)); err != nil {
+		if err := stateChanges.Collect(k, v); err != nil {
 			return err
 		}
 		return nil
@@ -297,7 +299,7 @@ func (a *Aggregator22) Unwind(ctx context.Context, txUnwindTo uint64, stateLoad
 		return err
 	}
 	if err := a.storage.pruneF(step, txUnwindTo, math2.MaxUint64, func(txNu uint64, k, v []byte) error {
-		if err := changes.Collect(libcommon.Copy(k), libcommon.Copy(v)); err != nil {
+		if err := stateChanges.Collect(k, v); err != nil {
 			return err
 		}
 		return nil
@@ -305,13 +307,25 @@ func (a *Aggregator22) Unwind(ctx context.Context, txUnwindTo uint64, stateLoad
 		return err
 	}
 
-	if err := changes.Load(a.rwTx, kv.PlainState, stateLoad, etl.TransformArgs{Quit: ctx.Done()}); err != nil {
+	if err := stateChanges.Load(a.rwTx, kv.PlainState, stateLoad, etl.TransformArgs{Quit: ctx.Done()}); err != nil {
 		return err
 	}
 
-	if err := a.code.prune(step, txUnwindTo, math2.MaxUint64); err != nil {
+	codeChanges := etl.NewCollector(a.logPrefix, "", etl.NewOldestEntryBuffer(etl.BufferOptimalSize))
+	defer codeChanges.Close()
+
+	if err := a.code.pruneF(step, txUnwindTo, math2.MaxUint64, func(txNum uint64, k, v []byte) error {
+		if err := codeChanges.Collect(k, v); err != nil {
+			return err
+		}
+		return nil
+	}); err != nil {
 		return err
 	}
+	if err := codeChanges.Load(a.rwTx, kv.PlainContractCode, etl.IdentityLoadFunc, etl.TransformArgs{Quit: ctx.Done()}); err != nil {
+		return err
+	}
+
 	if err := a.logAddrs.prune(txUnwindTo, math2.MaxUint64); err != nil {
 		return err
 	}
@@ -689,6 +703,7 @@ func (a *Aggregator22) AddStoragePrev(addr []byte, loc []byte, prev []byte) erro
 	return nil
 }
 
+// AddCodePrev - addr+inc => code
 func (a *Aggregator22) AddCodePrev(addr []byte, prev []byte) error {
 	if err := a.code.AddPrevValue(addr, nil, prev); err != nil {
 		return err