Changes based on PR comments

Author: amsanghi

core/state/statedb.go

@@ -1326,7 +1326,7 @@ func (s *StateDB) commit(deleteEmptyObjects bool) (*stateUpdate, error) {
 
 	origin := s.originalRoot
 	s.originalRoot = root
-	return newStateUpdate(origin, root, deletes, updates, nodes), nil
+	return newStateUpdate(origin, root, deletes, updates, nodes, s.arbExtraData.activatedWasms), nil
 }
 
 // commitAndFlush is a wrapper of commit which also commits the state mutations
@@ -1347,10 +1347,10 @@ func (s *StateDB) commitAndFlush(block uint64, deleteEmptyObjects bool) (*stateU
 		}
 	}
 
-	if db := s.db.WasmStore(); db != nil && len(s.arbExtraData.activatedWasms) > 0 {
+	if db := s.db.WasmStore(); db != nil && len(ret.activatedWasms) > 0 {
 		batch := db.NewBatch()
 		// Arbitrum: write Stylus programs to disk
-		for moduleHash, asmMap := range s.arbExtraData.activatedWasms {
+		for moduleHash, asmMap := range ret.activatedWasms {
 			rawdb.WriteActivation(batch, moduleHash, asmMap)
 		}
 		s.arbExtraData.activatedWasms = make(map[common.Hash]ActivatedWasm)

core/state/stateupdate.go

@@ -57,6 +57,7 @@ type stateUpdate struct {
 	storages       map[common.Hash]map[common.Hash][]byte    // storages stores mutated slots in 'prefix-zero-trimmed' RLP format
 	storagesOrigin map[common.Address]map[common.Hash][]byte // storagesOrigin stores the original values of mutated slots in 'prefix-zero-trimmed' RLP format
 	codes          map[common.Address]contractCode           // codes contains the set of dirty codes
+	activatedWasms map[common.Hash]ActivatedWasm             // newly activated WASMs
 	nodes          *trienode.MergedNodeSet                   // Aggregated dirty nodes caused by state changes
 }
 
@@ -68,7 +69,7 @@ func (sc *stateUpdate) empty() bool {
 // newStateUpdate constructs a state update object, representing the differences
 // between two states by performing state execution. It aggregates the given
 // account deletions and account updates to form a comprehensive state update.
-func newStateUpdate(originRoot common.Hash, root common.Hash, deletes map[common.Hash]*accountDelete, updates map[common.Hash]*accountUpdate, nodes *trienode.MergedNodeSet) *stateUpdate {
+func newStateUpdate(originRoot common.Hash, root common.Hash, deletes map[common.Hash]*accountDelete, updates map[common.Hash]*accountUpdate, nodes *trienode.MergedNodeSet, activatedWasms map[common.Hash]ActivatedWasm) *stateUpdate {
 	var (
 		destructs      = make(map[common.Hash]struct{})
 		accounts       = make(map[common.Hash][]byte)
@@ -128,6 +129,7 @@ func newStateUpdate(originRoot common.Hash, root common.Hash, deletes map[common
 		storages:       storages,
 		storagesOrigin: storagesOrigin,
 		codes:          codes,
+		activatedWasms: activatedWasms,
 		nodes:          nodes,
 	}
 }