From a76b23809ce236d621c9ed74f5118d056297c8dd Mon Sep 17 00:00:00 2001
From: Philip Morlier <philipmorlier@gmail.com>
Date: Thu, 24 Oct 2024 11:23:40 -0700
Subject: [PATCH] Follow up modifications to wrappers

---
 .../wrappers/backendwrapper/backendwrapper.go   | 17 +++++++++++++----
 plugins/wrappers/backendwrapper/triewrapper.go  |  4 +---
 plugins/wrappers/wrappers.go                    |  3 +--
 3 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/plugins/wrappers/backendwrapper/backendwrapper.go b/plugins/wrappers/backendwrapper/backendwrapper.go
index 99f5de39b..cc6f61cef 100644
--- a/plugins/wrappers/backendwrapper/backendwrapper.go
+++ b/plugins/wrappers/backendwrapper/backendwrapper.go
@@ -3,6 +3,7 @@ package backendwrapper
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"math/big"
 	"reflect"
@@ -44,10 +45,16 @@ type Backend struct {
 	removedLogsOnce sync.Once
 	chainConfig     *params.ChainConfig
 }
-// TODO AR review addition of db to backend type
 
 func NewBackend(b ethapi.Backend) *Backend {
-	return &Backend{b: b}
+	state, _, err := b.StateAndHeaderByNumber(context.Background(), 0)
+	if err != nil {
+		panic(err.Error())
+	}
+	return &Backend{
+		b: b,
+		db: state.Database(),
+	}
 }
 
 func (b *Backend) SuggestGasTipCap(ctx context.Context) (*big.Int, error) {
@@ -125,14 +132,16 @@ func (b *Backend) SendTx(ctx context.Context, signedTx []byte) error {
 	return b.b.SendTx(ctx, tx)
 }
 func (b *Backend) GetTransaction(ctx context.Context, txHash core.Hash) ([]byte, core.Hash, uint64, uint64, error) { // RLP Encoded transaction {
-	_, tx, blockHash, blockNumber, index, err := b.b.GetTransaction(ctx, common.Hash(txHash))
+	found, tx, blockHash, blockNumber, index, err := b.b.GetTransaction(ctx, common.Hash(txHash))
 	if err != nil {
 		return nil, core.Hash(blockHash), blockNumber, index, err
 	}
+	if !found {
+		return nil, core.Hash(blockHash), blockNumber, index, errors.New("not found returned from GetTransaction")
+	}
 	enc, err := tx.MarshalBinary()
 	return enc, core.Hash(blockHash), blockNumber, index, err
 }
-// TODO AR the above internal function signature needs review
 func (b *Backend) GetPoolTransactions() ([][]byte, error) {
 	txs, err := b.b.GetPoolTransactions()
 	if err != nil {
diff --git a/plugins/wrappers/backendwrapper/triewrapper.go b/plugins/wrappers/backendwrapper/triewrapper.go
index 6171b9c79..f71a932f8 100644
--- a/plugins/wrappers/backendwrapper/triewrapper.go
+++ b/plugins/wrappers/backendwrapper/triewrapper.go
@@ -1,8 +1,6 @@
 package backendwrapper
 
 import (
-	"math/big"
-
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/state"
 	"github.com/ethereum/go-ethereum/log"
@@ -30,7 +28,7 @@ func (t *WrappedTrie) GetAccount(address core.Address) (*core.StateAccount, erro
 	}
 	return &core.StateAccount{
 		Nonce: act.Nonce,
-		Balance: new(big.Int).SetBytes(act.Balance.Bytes()),
+		Balance: act.Balance.ToBig(),
 		Root: core.Hash(act.Root),
 		CodeHash: act.CodeHash,
 	}, nil
diff --git a/plugins/wrappers/wrappers.go b/plugins/wrappers/wrappers.go
index 6eec53d9c..f690e5bb5 100644
--- a/plugins/wrappers/wrappers.go
+++ b/plugins/wrappers/wrappers.go
@@ -21,7 +21,7 @@ func NewWrappedStateDB(d *state.StateDB) *WrappedStateDB {
 
 // GetBalance(Address) *big.Int
 func (w *WrappedStateDB) GetBalance(addr core.Address) *big.Int {
-	return new(big.Int).SetBytes(w.s.GetBalance(common.Address(addr)).Bytes())
+	return w.s.GetBalance(common.Address(addr)).ToBig()
 }
 
 // GetNonce(Address) uint64
@@ -97,7 +97,6 @@ func (w *WrappedStateDB) AddBalance(addr core.Address, amount *big.Int) {
 	castAmount := new(uint256.Int)
 	w.s.AddBalance(common.Address(addr), castAmount.SetBytes(amount.Bytes()), 0)
 }
-// TODO AR the above internal function signature needs review
 
 type Node struct {
 	n *node.Node