Add push/pop instructions (#498)

* Add push/pop instructions * Add tests to stack push/pop operations * clippy * more clippy * fmt after merge * Update changelog * Fix bug on edge case when sp == hp * Clean up casts * Fix errors after merge
FuelLabs · Aug 8, 2023 · 11dc5c9 · 11dc5c9
1 parent e24f188
commit 11dc5c9
Show file tree

Hide file tree

Showing 9 changed files with 403 additions and 13 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -63,6 +63,8 @@ The release mostly fixes funding during the audit and integration with the bridg
 
 - [#486](https://github.com/FuelLabs/fuel-vm/pull/486/): Adds `ed25519` signature verification and `secp256r1` signature recovery to `fuel-crypto`, and corresponding opcodes `ED19` and `ECR1` to `fuel-vm`.
 
+- [#486](https://github.com/FuelLabs/fuel-vm/pull/498): Adds `PSHL`, `PSHH`, `POPH` and `POPL` instructions, which allow cheap push and pop stack operations with multiple registers.
+
 - [#500](https://github.com/FuelLabs/fuel-vm/pull/500): Introduced `ParallelExecutor` trait
     and made available async versions of verify and estimate predicates.
     Updated tests to test for both parallel and sequential execution.

diff --git a/fuel-asm/src/lib.rs b/fuel-asm/src/lib.rs
@@ -298,6 +298,14 @@ impl_instructions! {
     0x93 CFE cfe [amount: RegId]
     "Shrink the current call frame's stack"
     0x94 CFS cfs [amount: RegId]
+    "Push a bitmask-selected set of registers in range 16..40 to the stack."
+    0x95 PSHL pshl [bitmask: Imm24]
+    "Push a bitmask-selected set of registers in range 40..64 to the stack."
+    0x96 PSHH pshh [bitmask: Imm24]
+    "Pop a bitmask-selected set of registers in range 16..40 to the stack."
+    0x97 POPL popl [bitmask: Imm24]
+    "Pop a bitmask-selected set of registers in range 40..64 to the stack."
+    0x98 POPH poph [bitmask: Imm24]
 
     "Compare 128bit integers"
     0xa0 WDCM wdcm [dst: RegId lhs: RegId rhs: RegId flags: Imm06]
@@ -678,12 +686,12 @@ impl Opcode {
         match self {
             ADD | AND | DIV | EQ | EXP | GT | LT | MLOG | MROO | MOD | MOVE | MUL
             | NOT | OR | SLL | SRL | SUB | XOR | WDCM | WQCM | WDOP | WQOP | WDML
-            | WQML | WDDV | WQDV | WDMD | WQMD | WDAM | WQAM | WDMM | WQMM | RET
-            | ALOC | MCL | MCP | MEQ | ECK1 | ECR1 | ED19 | K256 | S256 | NOOP | FLAG
-            | ADDI | ANDI | DIVI | EXPI | MODI | MULI | MLDV | ORI | SLLI | SRLI
-            | SUBI | XORI | JNEI | LB | LW | SB | SW | MCPI | MCLI | GM | MOVI | JNZI
-            | JI | JMP | JNE | JMPF | JMPB | JNZF | JNZB | JNEF | JNEB | CFEI | CFSI
-            | CFE | CFS | GTF => true,
+            | WQML | WDDV | WQDV | WDMD | WQMD | WDAM | WQAM | WDMM | WQMM | PSHH
+            | PSHL | POPH | POPL | RET | ALOC | MCL | MCP | MEQ | ECK1 | ECR1 | ED19
+            | K256 | S256 | NOOP | FLAG | ADDI | ANDI | DIVI | EXPI | MODI | MULI
+            | MLDV | ORI | SLLI | SRLI | SUBI | XORI | JNEI | LB | LW | SB | SW
+            | MCPI | MCLI | GM | MOVI | JNZI | JI | JMP | JNE | JMPF | JMPB | JNZF
+            | JNZB | JNEF | JNEB | CFEI | CFSI | CFE | CFS | GTF => true,
             _ => false,
         }
     }

diff --git a/fuel-tx/src/transaction/consensus_parameters/gas.rs b/fuel-tx/src/transaction/consensus_parameters/gas.rs
@@ -254,6 +254,10 @@ pub struct GasCostsValues {
     pub not: Word,
     pub or: Word,
     pub ori: Word,
+    pub poph: Word,
+    pub popl: Word,
+    pub pshh: Word,
+    pub pshl: Word,
     #[cfg_attr(feature = "serde", serde(rename = "ret_contract"))]
     pub ret: Word,
     #[cfg_attr(feature = "serde", serde(rename = "rvrt_contract"))]
@@ -394,6 +398,10 @@ impl GasCostsValues {
             not: 0,
             or: 0,
             ori: 0,
+            poph: 0,
+            popl: 0,
+            pshh: 0,
+            pshl: 0,
             ret: 0,
             rvrt: 0,
             s256: 0,
@@ -503,6 +511,10 @@ impl GasCostsValues {
             or: 1,
             ori: 1,
             ret: 1,
+            poph: 1,
+            popl: 1,
+            pshh: 1,
+            pshl: 1,
             rvrt: 1,
             s256: 1,
             sb: 1,

diff --git a/fuel-tx/src/transaction/consensus_parameters/gas/default_gas_costs.rs b/fuel-tx/src/transaction/consensus_parameters/gas/default_gas_costs.rs
@@ -63,6 +63,10 @@ pub fn default_gas_costs() -> GasCostsValues {
         not: 1,
         or: 1,
         ori: 1,
+        poph: 2,
+        popl: 2,
+        pshh: 2,
+        pshl: 2,
         move_op: 1,
         ret: 13,
         s256: 2,

diff --git a/fuel-vm/src/constraints/reg_key.rs b/fuel-vm/src/constraints/reg_key.rs
@@ -426,3 +426,32 @@ impl<'a> From<&SystemRegistersRef<'a>> for [Word; VM_REGISTER_SYSTEM_COUNT] {
         ]
     }
 }
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) enum ProgramRegistersSegment {
+    /// Registers 16..40
+    Low,
+    /// Registers 40..64
+    High,
+}
+
+impl<'r> ProgramRegisters<'r> {
+    /// Returns the registers corresponding to the segment, always 24 elements.
+    pub(crate) fn segment(&self, segment: ProgramRegistersSegment) -> &[Word] {
+        match segment {
+            ProgramRegistersSegment::Low => &self.0[..24],
+            ProgramRegistersSegment::High => &self.0[24..],
+        }
+    }
+
+    /// Returns the registers corresponding to the segment, always 24 elements.
+    pub(crate) fn segment_mut(
+        &mut self,
+        segment: ProgramRegistersSegment,
+    ) -> &mut [Word] {
+        match segment {
+            ProgramRegistersSegment::Low => &mut self.0[..24],
+            ProgramRegistersSegment::High => &mut self.0[24..],
+        }
+    }
+}
diff --git a/fuel-vm/src/interpreter/executors/instruction.rs b/fuel-vm/src/interpreter/executors/instruction.rs
@@ -1,4 +1,5 @@
 use crate::{
+    constraints::reg_key::ProgramRegistersSegment,
     error::{
         InterpreterError,
         RuntimeError,
@@ -618,6 +619,30 @@ where
                 self.stack_pointer_overflow(Word::overflowing_sub, r!(a))?;
             }
 
+            Instruction::PSHL(pshl) => {
+                self.gas_charge(self.gas_costs().pshl)?;
+                let bitmask = pshl.unpack();
+                self.push_selected_registers(ProgramRegistersSegment::Low, bitmask)?;
+            }
+
+            Instruction::PSHH(pshh) => {
+                self.gas_charge(self.gas_costs().pshh)?;
+                let bitmask = pshh.unpack();
+                self.push_selected_registers(ProgramRegistersSegment::High, bitmask)?;
+            }
+
+            Instruction::POPL(popl) => {
+                self.gas_charge(self.gas_costs().popl)?;
+                let bitmask = popl.unpack();
+                self.pop_selected_registers(ProgramRegistersSegment::Low, bitmask)?;
+            }
+
+            Instruction::POPH(poph) => {
+                self.gas_charge(self.gas_costs().poph)?;
+                let bitmask = poph.unpack();
+                self.pop_selected_registers(ProgramRegistersSegment::High, bitmask)?;
+            }
+
             Instruction::LB(lb) => {
                 self.gas_charge(self.gas_costs().lb)?;
                 let (a, b, imm) = lb.unpack();

diff --git a/fuel-vm/src/interpreter/executors/instruction/tests/reserved_registers.rs b/fuel-vm/src/interpreter/executors/instruction/tests/reserved_registers.rs
@@ -138,6 +138,10 @@ fn writes_to_ra(opcode: Opcode) -> bool {
         Opcode::NOT => true,
         Opcode::OR => true,
         Opcode::ORI => true,
+        Opcode::POPH => false,
+        Opcode::POPL => false,
+        Opcode::PSHH => false,
+        Opcode::PSHL => false,
         Opcode::SLL => true,
         Opcode::SLLI => true,
         Opcode::SRL => true,
@@ -248,6 +252,10 @@ fn writes_to_rb(opcode: Opcode) -> bool {
         Opcode::NOT => false,
         Opcode::OR => false,
         Opcode::ORI => false,
+        Opcode::POPH => false,
+        Opcode::POPL => false,
+        Opcode::PSHH => false,
+        Opcode::PSHL => false,
         Opcode::SLL => false,
         Opcode::SLLI => false,
         Opcode::SRL => false,

diff --git a/fuel-vm/src/interpreter/memory.rs b/fuel-vm/src/interpreter/memory.rs
@@ -11,6 +11,7 @@ use crate::{
 };
 
 use fuel_asm::{
+    Imm24,
     PanicReason,
     RegId,
 };
@@ -32,6 +33,9 @@ mod tests;
 #[cfg(test)]
 mod allocation_tests;
 
+#[cfg(test)]
+mod stack_tests;
+
 /// Used to handle `Word` to `usize` conversions for memory addresses,
 /// as well as checking that the resulting value is withing the VM ram boundaries.
 pub trait ToAddr {
@@ -214,6 +218,52 @@ where
         stack_pointer_overflow(sp, ssp.as_ref(), hp.as_ref(), pc, f, v)
     }
 
+    pub(crate) fn push_selected_registers(
+        &mut self,
+        segment: ProgramRegistersSegment,
+        bitmask: Imm24,
+    ) -> Result<(), RuntimeError> {
+        let (
+            SystemRegisters {
+                sp, ssp, hp, pc, ..
+            },
+            program_regs,
+        ) = split_registers(&mut self.registers);
+        push_selected_registers(
+            &mut self.memory,
+            sp,
+            ssp.as_ref(),
+            hp.as_ref(),
+            pc,
+            &program_regs,
+            segment,
+            bitmask,
+        )
+    }
+
+    pub(crate) fn pop_selected_registers(
+        &mut self,
+        segment: ProgramRegistersSegment,
+        bitmask: Imm24,
+    ) -> Result<(), RuntimeError> {
+        let (
+            SystemRegisters {
+                sp, ssp, hp, pc, ..
+            },
+            mut program_regs,
+        ) = split_registers(&mut self.registers);
+        pop_selected_registers(
+            &self.memory,
+            sp,
+            ssp.as_ref(),
+            hp.as_ref(),
+            pc,
+            &mut program_regs,
+            segment,
+            bitmask,
+        )
+    }
+
     pub(crate) fn load_byte(
         &mut self,
         ra: RegisterId,
@@ -290,26 +340,107 @@ where
     }
 }
 
-pub(crate) fn stack_pointer_overflow<F>(
+/// Update stack pointer, checking for validity first.
+pub(crate) fn try_update_stack_pointer(
     mut sp: RegMut<SP>,
     ssp: Reg<SSP>,
     hp: Reg<HP>,
+    new_sp: Word,
+) -> Result<(), RuntimeError> {
+    if new_sp >= *hp || new_sp < *ssp {
+        Err(PanicReason::MemoryOverflow.into())
+    } else {
+        *sp = new_sp;
+        Ok(())
+    }
+}
+
+pub(crate) fn stack_pointer_overflow<F>(
+    sp: RegMut<SP>,
+    ssp: Reg<SSP>,
+    hp: Reg<HP>,
     pc: RegMut<PC>,
     f: F,
     v: Word,
 ) -> Result<(), RuntimeError>
 where
     F: FnOnce(Word, Word) -> (Word, bool),
 {
-    let (result, overflow) = f(*sp, v);
+    let (new_sp, overflow) = f(*sp, v);
 
-    if overflow || result >= *hp || result < *ssp {
-        Err(PanicReason::MemoryOverflow.into())
-    } else {
-        *sp = result;
+    if overflow {
+        return Err(PanicReason::MemoryOverflow.into())
+    }
 
-        inc_pc(pc)
+    try_update_stack_pointer(sp, ssp, hp, new_sp)?;
+    inc_pc(pc)
+}
+
+#[allow(clippy::too_many_arguments)]
+pub(crate) fn push_selected_registers(
+    memory: &mut [u8; MEM_SIZE],
+    sp: RegMut<SP>,
+    ssp: Reg<SSP>,
+    hp: Reg<HP>,
+    pc: RegMut<PC>,
+    program_regs: &ProgramRegisters,
+    segment: ProgramRegistersSegment,
+    bitmask: Imm24,
+) -> Result<(), RuntimeError> {
+    let bitmask = bitmask.to_u32();
+
+    // First update the new stack pointer, as that's the only error condition
+    let count: u64 = bitmask.count_ones().into();
+    let stack_range = MemoryRange::new(*sp, count * 8)?;
+    try_update_stack_pointer(sp, ssp, hp, stack_range.words().end)?;
+
+    // Write the registers to the stack
+    let mut it = memory[stack_range.usizes()].chunks_exact_mut(8);
+    for (i, reg) in program_regs.segment(segment).iter().enumerate() {
+        if (bitmask & (1 << i)) != 0 {
+            let item = it
+                .next()
+                .expect("Memory range mismatched with register count");
+            item.copy_from_slice(&reg.to_be_bytes());
+        }
     }
+
+    inc_pc(pc)
+}
+
+#[allow(clippy::too_many_arguments)]
+pub(crate) fn pop_selected_registers(
+    memory: &[u8; MEM_SIZE],
+    sp: RegMut<SP>,
+    ssp: Reg<SSP>,
+    hp: Reg<HP>,
+    pc: RegMut<PC>,
+    program_regs: &mut ProgramRegisters,
+    segment: ProgramRegistersSegment,
+    bitmask: Imm24,
+) -> Result<(), RuntimeError> {
+    let bitmask = bitmask.to_u32();
+
+    // First update the stack pointer, as that's the only error condition
+    let count: u64 = bitmask.count_ones().into();
+    let size_in_stack = count * 8;
+    let new_sp = sp
+        .checked_sub(size_in_stack)
+        .ok_or(PanicReason::MemoryOverflow)?;
+    try_update_stack_pointer(sp, ssp, hp, new_sp)?;
+    let stack_range = MemoryRange::new(new_sp, size_in_stack)?.usizes();
+
+    // Restore registers from the stack
+    let mut it = memory[stack_range].chunks_exact(8);
+    for (i, reg) in program_regs.segment_mut(segment).iter_mut().enumerate() {
+        if (bitmask & (1 << i)) != 0 {
+            let mut buf = [0u8; 8];
+            buf.copy_from_slice(it.next().expect("Count mismatch"));
+            *reg = Word::from_be_bytes(buf);
+        }
+    }
+
+    inc_pc(pc)
 }
 
 pub(crate) fn load_byte(