From e60e5801daeab86af8b6775e27bf397b0cc47ae1 Mon Sep 17 00:00:00 2001
From: web3-developer <51288821+web3-developer@users.noreply.github.com>
Date: Tue, 9 Jan 2024 23:15:19 +0800
Subject: [PATCH] Fix vmflags for witness generation 1934. (#1961)

* Fix issue causing vmflags to be reset during call to processBlocks and enable witness generation in test_blockchain_json test.

* Fix copyright notice on updated files.
---
 nimbus/evm/state.nim           | 18 +++++++++++++-----
 tests/test_blockchain_json.nim | 12 +++++++++---
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/nimbus/evm/state.nim b/nimbus/evm/state.nim
index 5f51dedc3..f10a86541 100644
--- a/nimbus/evm/state.nim
+++ b/nimbus/evm/state.nim
@@ -1,5 +1,5 @@
 # Nimbus
-# Copyright (c) 2018-2023 Status Research & Development GmbH
+# Copyright (c) 2018-2024 Status Research & Development GmbH
 # Licensed under either of
 #  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
 #    http://www.apache.org/licenses/LICENSE-2.0)
@@ -27,7 +27,8 @@ proc init(
       blockCtx:     BlockContext;
       com:          CommonRef;
       tracer:       TracerRef,
-      asyncFactory: AsyncOperationFactory = AsyncOperationFactory(maybeDataSource: none[AsyncDataSource]()))
+      asyncFactory: AsyncOperationFactory = AsyncOperationFactory(maybeDataSource: none[AsyncDataSource]()),
+      flags:        set[VMFlag] = self.flags)
     {.gcsafe.} =
   ## Initialisation helper
   self.parent = parent
@@ -37,6 +38,7 @@ proc init(
   self.tracer = tracer
   self.stateDB = ac
   self.asyncFactory = asyncFactory
+  self.flags = flags
 
 func blockCtx(com: CommonRef, header: BlockHeader):
                 BlockContext {.gcsafe, raises: [CatchableError].} =
@@ -102,13 +104,15 @@ proc reinit*(self:     BaseVMState;     ## Object descriptor
       db     = com.db
       ac     = if self.stateDB.rootHash == parent.stateRoot: self.stateDB
                else: com.ledgerType.init(db, parent.stateRoot, com.pruneTrie)
+      flags  = self.flags
     self[].reset
     self.init(
       ac       = ac,
       parent   = parent,
       blockCtx = blockCtx,
       com      = com,
-      tracer   = tracer)
+      tracer   = tracer,
+      flags    = flags)
     return true
   # else: false
 
@@ -283,8 +287,8 @@ proc generateWitness*(vmState: BaseVMState): bool =
   GenerateWitness in vmState.flags
 
 proc `generateWitness=`*(vmState: BaseVMState, status: bool) =
- if status: vmState.flags.incl GenerateWitness
- else: vmState.flags.excl GenerateWitness
+  if status: vmState.flags.incl GenerateWitness
+  else: vmState.flags.excl GenerateWitness
 
 proc buildWitness*(vmState: BaseVMState): seq[byte]
     {.raises: [CatchableError].} =
@@ -292,6 +296,10 @@ proc buildWitness*(vmState: BaseVMState): seq[byte]
   let mkeys = vmState.stateDB.makeMultiKeys()
   let flags = if vmState.fork >= FkSpurious: {wfEIP170} else: {}
 
+  # A valid block having no transactions should return an empty witness
+  if mkeys.keys.len() == 0:
+    return @[]
+
   # build witness from tree
   var wb = initWitnessBuilder(vmState.com.db, rootHash, flags)
   wb.buildWitness(mkeys)
diff --git a/tests/test_blockchain_json.nim b/tests/test_blockchain_json.nim
index 74e3d6f1d..7c22da64d 100644
--- a/tests/test_blockchain_json.nim
+++ b/tests/test_blockchain_json.nim
@@ -1,5 +1,5 @@
 # Nimbus
-# Copyright (c) 2018-2023 Status Research & Development GmbH
+# Copyright (c) 2018-2024 Status Research & Development GmbH
 # Licensed under either of
 #  * Apache License, version 2.0, ([LICENSE-APACHE](LICENSE-APACHE) or
 #    http://www.apache.org/licenses/LICENSE-2.0)
@@ -176,6 +176,12 @@ proc parseTestCtx(fixture: JsonNode, testStatusIMPL: var TestStatus): TestCtx =
 proc blockWitness(vmState: BaseVMState, chainDB: CoreDbRef) =
   let rootHash = vmState.stateDB.rootHash
   let witness = vmState.buildWitness()
+
+  if witness.len() == 0:
+    if vmState.stateDB.makeMultiKeys().keys.len() != 0:
+      raise newException(ValidationError, "Invalid trie generated from block witness")
+    return
+
   let fork = vmState.fork
   let flags = if fork >= FKSpurious: {wfEIP170} else: {}
 
@@ -221,6 +227,7 @@ proc importBlock(ctx: var TestCtx, com: CommonRef,
       com,
       tracerInst,
     )
+    ctx.vmState.generateWitness = true # Enable saving witness data
 
   let
     chain = newChain(com, extraValidation = true, ctx.vmState)
@@ -229,8 +236,7 @@ proc importBlock(ctx: var TestCtx, com: CommonRef,
   if res == ValidationResult.Error:
     raise newException(ValidationError, "persistBlocks validation")
   else:
-    if ctx.vmState.generateWitness():
-      blockWitness(ctx.vmState, com.db)
+    blockWitness(chain.vmState, com.db)
 
 proc applyFixtureBlockToChain(ctx: var TestCtx, tb: var TestBlock,
                               com: CommonRef, checkSeal: bool) =