From 103f33f809f381e0393eae3ee263c593264c3bae Mon Sep 17 00:00:00 2001
From: Ivan Kalinin <i.kalinin@dexpa.io>
Date: Tue, 21 Jan 2025 01:12:32 +0100
Subject: [PATCH] fix(vm): fix saving `ret` on deep jump

---
 crypto/vm/vm.cpp | 9 +++++++--
 crypto/vm/vm.h   | 9 +++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/crypto/vm/vm.cpp b/crypto/vm/vm.cpp
index fb774f80..2924f081 100644
--- a/crypto/vm/vm.cpp
+++ b/crypto/vm/vm.cpp
@@ -257,6 +257,11 @@ int VmState::jump(Ref<Continuation> cont) {
 
 // general jump to continuation cont
 int VmState::jump(Ref<Continuation> cont, int pass_args) {
+  cont = adjust_jump_cont(std::move(cont), pass_args);
+  return jump_to(std::move(cont));
+}
+
+Ref<Continuation> VmState::adjust_jump_cont(Ref<Continuation> cont, int pass_args) {
   const ControlData* cont_data = cont->get_cdata();
   if (cont_data) {
     // first do the checks
@@ -297,7 +302,7 @@ int VmState::jump(Ref<Continuation> cont, int pass_args) {
         consume_stack_gas(copy);
       }
     }
-    return jump_to(std::move(cont));
+    return cont;
   } else {
     // have no continuation data, situation is somewhat simpler
     if (pass_args >= 0) {
@@ -309,7 +314,7 @@ int VmState::jump(Ref<Continuation> cont, int pass_args) {
         consume_stack_gas(pass_args);
       }
     }
-    return jump_to(std::move(cont));
+    return cont;
   }
 }
 
diff --git a/crypto/vm/vm.h b/crypto/vm/vm.h
index cf532293..534a46d3 100644
--- a/crypto/vm/vm.h
+++ b/crypto/vm/vm.h
@@ -352,6 +352,7 @@ class VmState final : public VmStateInterface {
   int call(Ref<Continuation> cont, int pass_args, int ret_args = -1);
   int jump(Ref<Continuation> cont);
   int jump(Ref<Continuation> cont, int pass_args);
+  Ref<Continuation> adjust_jump_cont(Ref<Continuation> cont, int pass_args);
   int ret();
   int ret(int ret_args);
   int ret_alt();
@@ -379,6 +380,14 @@ class VmState final : public VmStateInterface {
       if (cnt > free_nested_cont_jump && global_version >= 9) {
         consume_gas(1);
       }
+
+      if (cont.not_null()) {
+        const ControlData* cont_data = cont->get_cdata();
+        if (cont_data && (cont_data->stack.not_null() || cont_data->nargs >= 0)) {
+          // if cont has non-empty stack or expects fixed number of arguments, jump is not simple
+          cont = adjust_jump_cont(std::move(cont), -1);
+        }
+      }
     }
     return res;
   }