From a174827a38e602ab756bd44481819f4aa1a6b685 Mon Sep 17 00:00:00 2001
From: David Kazlauskas <davidkazlauskas@gmail.com>
Date: Fri, 7 Jun 2024 08:06:15 +0300
Subject: [PATCH] Move coprocessor routine processing to defer block

---
 core/vm/interpreter.go | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/core/vm/interpreter.go b/core/vm/interpreter.go
index 62957bf4c..107a63395 100644
--- a/core/vm/interpreter.go
+++ b/core/vm/interpreter.go
@@ -168,8 +168,23 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 	}
 
 	var coprocSegmentId fhevm.SegmentId
-	if !readOnly && in.evm.CoprocessorSession != nil {
+	useCoprocessor := !readOnly && in.evm.CoprocessorSession != nil
+	isValidCoprocessorSegment := false
+	if useCoprocessor {
 		coprocSegmentId = in.evm.CoprocessorSession.NextSegment()
+		defer func() {
+			if isValidCoprocessorSegment {
+				if input != nil && res != nil && contract.Address() == in.evm.CoprocessorSession.ContractAddress() {
+					log.Info("Executing coprocessor payload", "input", common.Bytes2Hex(input), "output", common.Bytes2Hex(res))
+					coprocErr := in.evm.CoprocessorSession.Execute(input, res)
+					if coprocErr != nil {
+						log.Error("Error executing coprocessor payload", "error", coprocErr)
+					}
+				}
+			} else {
+				_ = in.evm.CoprocessorSession.InvalidateSinceSegment(coprocSegmentId)
+			}
+		}()
 	}
 
 	// The Interpreter main run loop (contextual). This loop runs until either an
@@ -245,18 +260,9 @@ func (in *EVMInterpreter) Run(contract *Contract, input []byte, readOnly bool) (
 		err = nil // clear stop token error
 	}
 
-	if !readOnly && in.evm.CoprocessorSession != nil {
-		if err != nil {
-			_ = in.evm.CoprocessorSession.InvalidateSinceSegment(coprocSegmentId)
-		} else {
-			if input != nil && res != nil && contract.Address() == in.evm.CoprocessorSession.ContractAddress() {
-				log.Info("Executing coprocessor payload", "input", common.Bytes2Hex(input), "output", common.Bytes2Hex(res))
-				coprocErr := in.evm.CoprocessorSession.Execute(input, res)
-				if coprocErr != nil {
-					log.Error("Error executing coprocessor payload", "error", coprocErr)
-				}
-			}
-		}
+	if useCoprocessor && err == nil {
+		// only if this point is reached and there is no error process coprocessor payload in defer block
+		isValidCoprocessorSegment = true
 	}
 
 	return res, err