From 107c6164994c3b60b90fe3aaa48b0fb789e114f5 Mon Sep 17 00:00:00 2001
From: Jerry <jerrycgh@gmail.com>
Date: Wed, 24 Jan 2024 03:00:55 -0800
Subject: [PATCH] Add actual code in SMT witness (#68)

---
 smt/pkg/db/mdbx.go     | 13 +++++++++++++
 smt/pkg/db/mem-db.go   |  4 ++++
 smt/pkg/smt/smt.go     |  1 +
 smt/pkg/smt/witness.go | 24 ++++++++++++++++--------
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/smt/pkg/db/mdbx.go b/smt/pkg/db/mdbx.go
index a00eae83bb8..1c649cb9fa1 100644
--- a/smt/pkg/db/mdbx.go
+++ b/smt/pkg/db/mdbx.go
@@ -235,6 +235,19 @@ func (m *EriDb) GetHashKey(key utils.NodeKey) (utils.NodeKey, error) {
 	return utils.NodeKey{na[0], na[1], na[2], na[3]}, nil
 }
 
+func (m *EriDb) GetCode(codeHash []byte) ([]byte, error) {
+	data, err := m.tx.GetOne(kv.Code, codeHash)
+	if err != nil {
+		return nil, err
+	}
+
+	if data == nil {
+		return nil, fmt.Errorf("code hash %x not found", codeHash)
+	}
+
+	return data, nil
+}
+
 func (m *EriDb) Delete(key string) error {
 	return m.tx.Delete(TableSmt, []byte(key))
 }
diff --git a/smt/pkg/db/mem-db.go b/smt/pkg/db/mem-db.go
index bc95a14957c..f44c4f7af98 100644
--- a/smt/pkg/db/mem-db.go
+++ b/smt/pkg/db/mem-db.go
@@ -176,6 +176,10 @@ func (m *MemDb) GetHashKey(key utils.NodeKey) (utils.NodeKey, error) {
 	return utils.NodeKey{na[0], na[1], na[2], na[3]}, nil
 }
 
+func (m *MemDb) GetCode(codeHash []byte) ([]byte, error) {
+	return nil, fmt.Errorf("not implemented")
+}
+
 func (m *MemDb) Delete(key string) error {
 	m.lock.Lock()         // Lock for writing
 	defer m.lock.Unlock() // Make sure to unlock when done
diff --git a/smt/pkg/smt/smt.go b/smt/pkg/smt/smt.go
index 5ae5f0878f4..0e05916b944 100644
--- a/smt/pkg/smt/smt.go
+++ b/smt/pkg/smt/smt.go
@@ -29,6 +29,7 @@ type DB interface {
 	InsertHashKey(key utils.NodeKey, value utils.NodeKey) error
 	GetHashKey(key utils.NodeKey) (utils.NodeKey, error)
 	Delete(string) error
+	GetCode(codeHash []byte) ([]byte, error)
 
 	SetLastRoot(lr *big.Int) error
 	GetLastRoot() (*big.Int, error)
diff --git a/smt/pkg/smt/witness.go b/smt/pkg/smt/witness.go
index 4b52724218e..ec2d7c3c80f 100644
--- a/smt/pkg/smt/witness.go
+++ b/smt/pkg/smt/witness.go
@@ -18,11 +18,9 @@ func BuildWitness(s *SMT, rd trie.RetainDecider, ctx context.Context) (*trie.Wit
 
 	action := func(prefix []byte, k utils.NodeKey, v utils.NodeValue12) (bool, error) {
 		if rd != nil && !rd.Retain(prefix) {
-			if !v.IsFinalNode() {
-				h := libcommon.BigToHash(k.ToBigInt())
-				hNode := trie.OperatorHash{Hash: h}
-				operands = append(operands, &hNode)
-			}
+			h := libcommon.BigToHash(k.ToBigInt())
+			hNode := trie.OperatorHash{Hash: h}
+			operands = append(operands, &hNode)
 			return false, nil
 		}
 
@@ -46,21 +44,31 @@ func BuildWitness(s *SMT, rd trie.RetainDecider, ctx context.Context) (*trie.Wit
 
 			valHash := v.Get4to8()
 
-			value, err := s.Db.Get(*valHash)
+			v, err := s.Db.Get(*valHash)
 
-			value8 := utils.BigIntArrayFromNodeValue8(value.GetNodeValue8())
+			vInBytes := utils.ArrayBigToScalar(utils.BigIntArrayFromNodeValue8(v.GetNodeValue8())).Bytes()
 
 			if err != nil {
 				return false, err
 			}
 
+			if t == utils.SC_CODE {
+				code, err := s.Db.GetCode(vInBytes)
+
+				if err != nil {
+					return false, err
+				} else {
+					operands = append(operands, &trie.OperatorCode{Code: code})
+				}
+			}
+
 			// fmt.Printf("Node hash: %s, Node type: %d, address %x, storage %x, value %x\n", utils.ConvertBigIntToHex(k.ToBigInt()), t, addr, storage, utils.ArrayBigToScalar(value8).Bytes())
 
 			operands = append(operands, &trie.OperatorSMTLeafValue{
 				NodeType:   uint8(t),
 				Address:    addr.Bytes(),
 				StorageKey: storage.Bytes(),
-				Value:      utils.ArrayBigToScalar(value8).Bytes(),
+				Value:      vInBytes,
 			})
 			return false, nil
 		}