Refactor cfi decoder

mohanson · mohanson · commit 11b4782101f6 · 2025-11-06T21:02:40.000+08:00
diff --git a/examples/fib b/examples/fib
diff --git a/src/decoder.rs b/src/decoder.rs
@@ -1,6 +1,7 @@
 use ckb_vm_definitions::instructions::{self as insts};
 use ckb_vm_definitions::registers::{RA, ZERO};
 
+use crate::elf::CFI;
 use crate::error::OutOfBoundKind;
 use crate::instructions::{
     Instruction, InstructionFactory, Itype, R4type, R5type, Register, Rtype, Utype, a, b, cfi,
@@ -14,7 +15,7 @@ const RISCV_PAGESIZE_MASK: u64 = RISCV_PAGESIZE as u64 - 1;
 const INSTRUCTION_CACHE_SIZE: usize = 4096;
 
 pub trait InstDecoder {
-    fn new<R: Register>(isa: u8, version: u32) -> Self;
+    fn new<R: Register>(isa: u8, version: u32, cfi: CFI) -> Self;
     fn decode<M: Memory>(&mut self, memory: &mut M, pc: u64) -> Result<Instruction, Error>;
     fn reset_instructions_cache(&mut self) -> Result<(), Error>;
 }
@@ -23,6 +24,7 @@ pub struct DefaultDecoder {
     factories: Vec<InstructionFactory>,
     mop: bool,
     version: u32,
+    cfi: CFI,
     // use a cache of instructions to avoid decoding the same instruction twice, pc is the key and the instruction is the value
     instructions_cache: [(u64, u64); INSTRUCTION_CACHE_SIZE],
 }
@@ -34,6 +36,7 @@ impl DefaultDecoder {
             factories: vec![],
             mop,
             version,
+            cfi: CFI::default(),
             instructions_cache: [(u64::MAX, 0); INSTRUCTION_CACHE_SIZE],
         }
     }
@@ -118,7 +121,7 @@ impl DefaultDecoder {
         }
         let instruction_bits = self.decode_bits(memory, pc)?;
         for factory in &self.factories {
-            if let Some(instruction) = factory(instruction_bits, self.version) {
+            if let Some(instruction) = factory(instruction_bits, self.version, self.cfi) {
                 self.instructions_cache[instruction_cache_key] = (pc, instruction);
                 return Ok(instruction);
             }
@@ -859,8 +862,9 @@ impl DefaultDecoder {
 }
 
 impl InstDecoder for DefaultDecoder {
-    fn new<R: Register>(isa: u8, version: u32) -> Self {
+    fn new<R: Register>(isa: u8, version: u32, cfi: CFI) -> Self {
         let mut decoder = Self::empty(isa & ISA_MOP != 0, version);
+        decoder.cfi = cfi;
         if isa & ISA_CFI != 0 {
             decoder.add_instruction_factory(cfi::factory::<R>);
         }
diff --git a/src/elf.rs b/src/elf.rs
@@ -136,7 +136,7 @@ pub struct LoadingAction {
     pub offset_from_addr: u64,
 }
 
-#[derive(Clone, Debug, Default, Eq, PartialEq)]
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
 pub struct CFI {
     pub lp_unlabeled: bool,
     pub ss: bool,
diff --git a/src/instructions/a.rs b/src/instructions/a.rs
@@ -2,8 +2,9 @@ use ckb_vm_definitions::instructions as insts;
 
 use super::utils::{funct3, funct7, opcode, rd, rs1, rs2};
 use super::{Instruction, Register, Rtype, set_instruction_length_4};
+use crate::elf::CFI;
 
-pub fn factory<R: Register>(instruction_bits: u32, _: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, _: u32, _: CFI) -> Option<Instruction> {
     let bit_length = R::BITS;
     if bit_length != 32 && bit_length != 64 {
         return None;
diff --git a/src/instructions/b.rs b/src/instructions/b.rs
@@ -5,8 +5,9 @@ use ckb_vm_definitions::instructions as insts;
 
 use super::utils::{self, funct3, funct7, opcode, rd, rs1, rs2};
 use super::{Instruction, Itype, Register, Rtype, set_instruction_length_4};
+use crate::elf::CFI;
 
-pub fn factory<R: Register>(instruction_bits: u32, _: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, _: u32, _: CFI) -> Option<Instruction> {
     let bit_length = R::BITS;
     if bit_length != 32 && bit_length != 64 {
         return None;
diff --git a/src/instructions/cfi.rs b/src/instructions/cfi.rs
@@ -1,9 +1,29 @@
 use super::utils::{funct3, funct7, opcode, rd, rs1, rs2};
 use super::{Instruction, Register, set_instruction_length_2, set_instruction_length_4};
-use crate::instructions::{Itype, Rtype, Utype};
+use crate::elf::CFI;
+use crate::instructions::i::nop;
+use crate::instructions::{Itype, Rtype, Utype, extract_opcode, instruction_length};
 use ckb_vm_definitions::instructions as insts;
 
-pub fn factory<R: Register>(instruction_bits: u32, _: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, _: u32, cfi: CFI) -> Option<Instruction> {
+    let inst = factory_bare::<R>(instruction_bits, 0, cfi);
+    if let Some(i) = inst {
+        if extract_opcode(i) == insts::OP_LPAD && cfi.lp_unlabeled {
+            return Some(i);
+        }
+        if cfi.ss {
+            return Some(i);
+        }
+        match instruction_length(i) {
+            2 => return Some(set_instruction_length_4(nop())),
+            4 => return Some(set_instruction_length_2(nop())),
+            _ => return None, // Should not happen.
+        }
+    }
+    None
+}
+
+pub fn factory_bare<R: Register>(instruction_bits: u32, _: u32, _: CFI) -> Option<Instruction> {
     match opcode(instruction_bits) {
         0b_0010111 => {
             if rd(instruction_bits) == 0 {
diff --git a/src/instructions/execute.rs b/src/instructions/execute.rs
@@ -555,8 +555,10 @@ pub fn handle_jalr_version1<Mac: Machine>(
     next_pc = next_pc & (!Mac::REG::one());
     update_register(machine, i.rd(), link);
     machine.update_pc(next_pc);
-    if i.rs1() != RA && i.rs1() != T0 && i.rs1() != T2 {
-        machine.set_elp(1);
+    if machine.cfi().lp_unlabeled {
+        if i.rs1() != RA && i.rs1() != T0 && i.rs1() != T2 {
+            machine.set_elp(1);
+        }
     }
     Ok(())
 }
@@ -1510,7 +1512,6 @@ pub fn handle_lpad<Mac: Machine>(machine: &mut Mac, inst: Instruction) -> Result
     if machine.elp() == 0 {
         return Ok(());
     }
-    println!("Handling LPAD instruction");
     // If PC not 4-byte aligned then software-check exception.
     if machine.pc().to_u64() % 4 != 0 {
         return Err(Error::ShadowStackSoftwareCheckException);
diff --git a/src/instructions/i.rs b/src/instructions/i.rs
@@ -7,6 +7,7 @@ use super::utils::{
 use super::{
     Instruction, Itype, Register, Rtype, Stype, Utype, blank_instruction, set_instruction_length_4,
 };
+use crate::elf::CFI;
 
 // The FENCE instruction is used to order device I/O and memory accesses
 // as viewed by other RISC- V harts and external devices or coprocessors.
@@ -31,7 +32,7 @@ impl FenceType {
     }
 }
 
-pub fn factory<R: Register>(instruction_bits: u32, version: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, version: u32, _: CFI) -> Option<Instruction> {
     let bit_length = R::BITS;
     if bit_length != 32 && bit_length != 64 {
         return None;
diff --git a/src/instructions/m.rs b/src/instructions/m.rs
@@ -2,8 +2,9 @@ use ckb_vm_definitions::instructions as insts;
 
 use super::utils::{funct3, funct7, opcode, rd, rs1, rs2};
 use super::{Instruction, Register, Rtype, set_instruction_length_4};
+use crate::elf::CFI;
 
-pub fn factory<R: Register>(instruction_bits: u32, _: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, _: u32, _: CFI) -> Option<Instruction> {
     let bit_length = R::BITS;
     if bit_length != 32 && bit_length != 64 {
         return None;
diff --git a/src/instructions/mod.rs b/src/instructions/mod.rs
@@ -14,6 +14,7 @@ pub mod tagged;
 
 pub use self::register::Register;
 use super::Error;
+use super::elf::CFI;
 pub use ckb_vm_definitions::{
     instructions::{
         self as insts, Instruction, InstructionOpcode, MAXIMUM_BASIC_BLOCK_END_OPCODE,
@@ -35,7 +36,8 @@ pub fn extract_opcode(i: Instruction) -> InstructionOpcode {
     (((i >> 8) & 0xff00) | (i & 0x00ff)) as u16
 }
 
-pub type InstructionFactory = fn(instruction_bits: u32, version: u32) -> Option<Instruction>;
+pub type InstructionFactory =
+    fn(instruction_bits: u32, version: u32, cfi: CFI) -> Option<Instruction>;
 
 // Blank instructions need no register indices nor immediates, they only have opcode
 // and module bit set.
@@ -458,14 +460,14 @@ mod tests {
     fn test_stype_display() {
         // This is "sd	a5,568(sp)"
         let sd_inst = 0x22f13c23;
-        let decoded = factory::<u64>(sd_inst, u32::MAX).expect("decoding");
+        let decoded = factory::<u64>(sd_inst, u32::MAX, CFI::default()).expect("decoding");
         let stype = Stype(decoded);
 
         assert_eq!("sd a5,568(sp)", format!("{}", stype));
 
         // This is "beq	a0,a5,1012e"
         let sd_inst = 0xf4f500e3;
-        let decoded = factory::<u64>(sd_inst, u32::MAX).expect("decoding");
+        let decoded = factory::<u64>(sd_inst, u32::MAX, CFI::default()).expect("decoding");
         let stype = Stype(decoded);
 
         assert_eq!("beq a0,a5,-192", format!("{}", stype));
diff --git a/src/instructions/rvc.rs b/src/instructions/rvc.rs
@@ -5,6 +5,7 @@ use super::i::nop;
 use super::register::Register;
 use super::utils::{jalr, ld, lw, rd, x, xs};
 use super::{Instruction, Itype, Rtype, Stype, Utype, blank_instruction, set_instruction_length_2};
+use crate::elf::CFI;
 
 // Notice the location of rs2 in RVC encoding is different from full encoding
 #[inline(always)]
@@ -90,7 +91,7 @@ fn b_immediate(instruction_bits: u32) -> i32 {
 }
 
 #[allow(clippy::cognitive_complexity)]
-pub fn factory<R: Register>(instruction_bits: u32, version: u32) -> Option<Instruction> {
+pub fn factory<R: Register>(instruction_bits: u32, version: u32, _: CFI) -> Option<Instruction> {
     let bit_length = R::BITS;
     if bit_length != 32 && bit_length != 64 {
         return None;
diff --git a/src/machine/mod.rs b/src/machine/mod.rs
@@ -11,7 +11,7 @@ use bytes::Bytes;
 
 use super::debugger::Debugger;
 use super::decoder::{DefaultDecoder, InstDecoder};
-use super::elf::{LoadingAction, ProgramMetadata, parse_elf};
+use super::elf::{CFI, LoadingAction, ProgramMetadata, parse_elf};
 use super::instructions::{Instruction, Register, execute};
 use super::memory::{Memory, load_c_string_byte_by_byte};
 use super::syscalls::Syscalls;
@@ -53,6 +53,8 @@ pub trait CoreMachine {
     fn isa(&self) -> u8;
 
     // CFI specific field.
+    fn cfi(&self) -> CFI;
+    fn set_cfi(&mut self, cfi: CFI);
     fn elp(&self) -> u32;
     fn set_elp(&mut self, elp: u32);
     fn ssp(&self) -> &Self::REG;
@@ -181,6 +183,7 @@ pub trait SupportMachine: CoreMachine {
             self.update_pc(Self::REG::from_u64(metadata.entry));
             self.commit_pc();
         }
+        self.set_cfi(metadata.cfi);
         Ok(bytes)
     }
 
@@ -311,6 +314,7 @@ pub trait DefaultMachineRunner {
         let mut decoder = Self::Decoder::new::<<Self::Inner as CoreMachine>::REG>(
             self.machine().isa(),
             self.machine().version(),
+            self.machine().cfi(),
         );
         self.run_with_decoder(&mut decoder)
     }
@@ -353,6 +357,7 @@ pub struct DefaultCoreMachine<R, M> {
     isa: u8,
     version: u32,
     // CFI specific field.
+    cfi: CFI,
     elp: u32,
     shadow_stack: [u8; DEFAULT_SHADOW_STACK_SIZE],
     ssp: R,
@@ -373,6 +378,7 @@ impl<R: Default, M: Default> Default for DefaultCoreMachine<R, M> {
             running: Default::default(),
             isa: Default::default(),
             version: Default::default(),
+            cfi: CFI::default(),
             elp: Default::default(),
             shadow_stack: [0; DEFAULT_SHADOW_STACK_SIZE],
             ssp: Default::default(),
@@ -421,6 +427,14 @@ impl<R: Register, M: Memory<REG = R>> CoreMachine for DefaultCoreMachine<R, M> {
         self.version
     }
 
+    fn cfi(&self) -> CFI {
+        self.cfi
+    }
+
+    fn set_cfi(&mut self, cfi: CFI) {
+        self.cfi = cfi;
+    }
+
     fn elp(&self) -> u32 {
         self.elp
     }
@@ -500,6 +514,7 @@ impl<R: Register, M: Memory<REG = R>> SupportMachine for DefaultCoreMachine<R, M
             running: Default::default(),
             isa,
             version,
+            cfi: Default::default(),
             elp: Default::default(),
             shadow_stack: [0; DEFAULT_SHADOW_STACK_SIZE],
             ssp: Default::default(),
@@ -643,6 +658,14 @@ impl<Inner: CoreMachine, Decoder> CoreMachine for DefaultMachine<Inner, Decoder>
         self.inner.version()
     }
 
+    fn cfi(&self) -> CFI {
+        self.inner.cfi()
+    }
+
+    fn set_cfi(&mut self, cfi: CFI) {
+        self.inner.set_cfi(cfi);
+    }
+
     fn elp(&self) -> u32 {
         self.inner.elp()
     }
diff --git a/src/machine/trace.rs b/src/machine/trace.rs
@@ -8,7 +8,7 @@ use super::{
             handle_invalid_op, instruction_length, is_basic_block_end_instruction,
         },
     },
-    CoreMachine, DEFAULT_SHADOW_STACK_SIZE, DefaultMachine, DefaultMachineRunner, Machine,
+    CFI, CoreMachine, DEFAULT_SHADOW_STACK_SIZE, DefaultMachine, DefaultMachineRunner, Machine,
     SupportMachine, VERSION2,
 };
 use bytes::Bytes;
@@ -96,6 +96,14 @@ impl<Inner: SupportMachine, Decoder> CoreMachine for AbstractTraceMachine<Inner,
         self.machine.version()
     }
 
+    fn cfi(&self) -> CFI {
+        self.machine.cfi()
+    }
+
+    fn set_cfi(&mut self, cfi: CFI) {
+        self.machine.set_cfi(cfi);
+    }
+
     fn elp(&self) -> u32 {
         self.machine.elp()
     }
diff --git a/tests/test_auipc_fusion.rs b/tests/test_auipc_fusion.rs
@@ -1,4 +1,5 @@
 use ckb_vm::decoder::{DefaultDecoder, InstDecoder};
+use ckb_vm::elf::CFI;
 use ckb_vm::instructions::{
     Instruction, Utype, extract_opcode, instruction_length, set_instruction_length_n,
 };
@@ -18,9 +19,9 @@ pub struct AuxDecoder {
 }
 
 impl InstDecoder for AuxDecoder {
-    fn new<R: Register>(isa: u8, version: u32) -> Self {
+    fn new<R: Register>(isa: u8, version: u32, cfi: CFI) -> Self {
         Self {
-            inner: DefaultDecoder::new::<R>(isa, version),
+            inner: DefaultDecoder::new::<R>(isa, version, cfi),
         }
     }
 
@@ -62,7 +63,7 @@ pub fn test_rust_auipc_fusion() {
         .load_program(&buffer, [Ok("auipc_no_sign_extend".into())].into_iter())
         .unwrap();
 
-    let mut decoder = AuxDecoder::new::<u64>(machine.isa(), machine.version());
+    let mut decoder = AuxDecoder::new::<u64>(machine.isa(), machine.version(), CFI::default());
     let result = machine.run_with_decoder(&mut decoder).unwrap();
     assert_eq!(result, 0);
 }
diff --git a/tests/test_cfi.rs b/tests/test_cfi.rs
@@ -0,0 +1,15 @@
+use ckb_vm::machine::VERSION3;
+use ckb_vm::{DefaultMachineRunner, ISA_B, ISA_CFI, ISA_IMC, ISA_MOP};
+pub mod machine_build;
+
+#[test]
+pub fn test_simple_instructions_64() {
+    let mut machine = machine_build::int(
+        "tests/programs/simple64",
+        vec![],
+        VERSION3,
+        ISA_IMC | ISA_B | ISA_MOP | ISA_CFI,
+    );
+    let ret = machine.run();
+    assert!(ret.is_ok());
+}

Original file line number	Diff line number	Diff line change
`@@ -136,7 +136,7 @@ pub struct LoadingAction {`
`136`	`136`	`pub offset_from_addr: u64,`
`137`	`137`	`}`
`138`	`138`
`139`		`-#[derive(Clone, Debug, Default, Eq, PartialEq)]`
	`139`	`+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]`
`140`	`140`	`pub struct CFI {`
`141`	`141`	`pub lp_unlabeled: bool,`
`142`	`142`	`pub ss: bool,`
Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ use super::utils::{`
`7`	`7`	`use super::{`
`8`	`8`	`Instruction, Itype, Register, Rtype, Stype, Utype, blank_instruction, set_instruction_length_4,`
`9`	`9`	`};`
	`10`	`+use crate::elf::CFI;`
`10`	`11`
`11`	`12`	`// The FENCE instruction is used to order device I/O and memory accesses`
`12`	`13`	`// as viewed by other RISC- V harts and external devices or coprocessors.`
`@@ -31,7 +32,7 @@ impl FenceType {`
`31`	`32`	`}`
`32`	`33`	`}`
`33`	`34`
`34`		`-pub fn factory<R: Register>(instruction_bits: u32, version: u32) -> Option<Instruction> {`
	`35`	`+pub fn factory<R: Register>(instruction_bits: u32, version: u32, _: CFI) -> Option<Instruction> {`
`35`	`36`	`let bit_length = R::BITS;`
`36`	`37`	`if bit_length != 32 && bit_length != 64 {`
`37`	`38`	`return None;`