From c7663a17af3d2249652cdefd256c66017386305b Mon Sep 17 00:00:00 2001
From: Andrei Kashin <kashin.andrej@gmail.com>
Date: Tue, 23 Jan 2024 12:39:20 +0000
Subject: [PATCH] Account for virtual sp adjustments

When we use NominalSPOffset, we need to adjust it using `virtual_sp_offset`.
---
 cranelift/codegen/src/isa/zkasm/inst/emit.rs  | 42 +++++++++----------
 .../sha256/generated/from_rust.zkasm          |  4 +-
 2 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/cranelift/codegen/src/isa/zkasm/inst/emit.rs b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
index 92bb22cc2d19..cb5cb04a8269 100644
--- a/cranelift/codegen/src/isa/zkasm/inst/emit.rs
+++ b/cranelift/codegen/src/isa/zkasm/inst/emit.rs
@@ -589,34 +589,34 @@ impl MachInstEmit for Inst {
                 let offset = from.get_offset_with_state(state);
                 let rd = allocs.next_writable(rd);
                 match from {
-                    AMode::RegOffset(r, off, _) => {
+                    AMode::RegOffset(r, ..) => {
                         debug_assert_eq!(r, e0());
                         put_string(
                             &format!(
                                 "$ => {} :MLOAD(MEM:{})\n",
                                 reg_name(rd.to_reg()),
-                                access_reg_with_offset(r, off)
+                                access_reg_with_offset(r, offset)
                             ),
                             sink,
                         );
                     }
-                    AMode::SPOffset(off, _) | AMode::NominalSPOffset(off, _) => {
-                        assert_eq!(off % 8, 0);
+                    AMode::SPOffset(..) | AMode::NominalSPOffset(..) => {
+                        assert_eq!(offset % 8, 0);
                         put_string(
                             &format!(
                                 "$ => {} :MLOAD({})\n",
                                 reg_name(rd.to_reg()),
-                                access_reg_with_offset(stack_reg(), off / 8),
+                                access_reg_with_offset(stack_reg(), offset / 8),
                             ),
                             sink,
                         );
                     }
-                    AMode::FPOffset(off, _) => {
+                    AMode::FPOffset(..) => {
                         put_string(
                             &format!(
                                 "$ => {} :MLOAD({})\n",
                                 reg_name(rd.to_reg()),
-                                access_reg_with_offset(fp_reg(), off),
+                                access_reg_with_offset(fp_reg(), offset),
                             ),
                             sink,
                         );
@@ -634,36 +634,37 @@ impl MachInstEmit for Inst {
             &Inst::Store { op, src, flags, to } => {
                 let to = to.clone().with_allocs(&mut allocs);
                 let src = allocs.next(src);
+                let offset = to.get_offset_with_state(state);
 
                 match to {
-                    AMode::RegOffset(r, off, _) => {
+                    AMode::RegOffset(r, ..) => {
                         debug_assert_eq!(r, e0());
                         put_string(
                             &format!(
                                 "{} :MSTORE(MEM:{})\n",
                                 reg_name(src),
-                                access_reg_with_offset(r, off)
+                                access_reg_with_offset(r, offset)
                             ),
                             sink,
                         );
                     }
-                    AMode::SPOffset(off, _) | AMode::NominalSPOffset(off, _) => {
-                        assert_eq!(off % 8, 0);
+                    AMode::SPOffset(..) | AMode::NominalSPOffset(..) => {
+                        assert_eq!(offset % 8, 0);
                         put_string(
                             &format!(
                                 "{} :MSTORE({})\n",
                                 reg_name(src),
-                                access_reg_with_offset(stack_reg(), off / 8),
+                                access_reg_with_offset(stack_reg(), offset / 8),
                             ),
                             sink,
                         );
                     }
-                    AMode::FPOffset(off, _) => {
+                    AMode::FPOffset(..) => {
                         put_string(
                             &format!(
                                 "{} :MSTORE({})\n",
                                 reg_name(src),
-                                access_reg_with_offset(fp_reg(), off),
+                                access_reg_with_offset(fp_reg(), offset),
                             ),
                             sink,
                         );
@@ -1182,13 +1183,12 @@ impl MachInstEmit for Inst {
             }
 
             &Inst::VirtualSPOffsetAdj { amount } => {
-                println!("virtual_sp_offset_adj {amount}");
-                // crate::trace!(
-                //     "virtual sp offset adjusted by {} -> {}",
-                //     amount,
-                //     state.virtual_sp_offset + amount
-                //     );
-                // state.virtual_sp_offset += amount;
+                crate::trace!(
+                    "virtual sp offset adjusted by {} -> {}",
+                    amount,
+                    state.virtual_sp_offset + amount
+                );
+                state.virtual_sp_offset += amount;
             }
 
             &Inst::LoadAddr { rd, mem } => {
diff --git a/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm b/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm
index e24bf5a78c15..df2edfd9d786 100644
--- a/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm
+++ b/cranelift/zkasm_data/benchmarks/sha256/generated/from_rust.zkasm
@@ -557,7 +557,7 @@ label_1_3:
   1n => B  ;; LoadConst32
   SP - 1 => SP
   B :MSTORE(SP)
-  $ => B :MLOAD(SP)
+  $ => B :MLOAD(SP + 1)
   zkPC + 2 => RR
   :JMP(function_2)
   SP + 1 => SP
@@ -572,7 +572,7 @@ label_1_5:
   1n => C  ;; LoadConst32
   SP - 1 => SP
   C :MSTORE(SP)
-  $ => B :MLOAD(SP)
+  $ => B :MLOAD(SP + 1)
   zkPC + 2 => RR
   :JMP(function_2)
   SP + 1 => SP