Merge pull request #2889 from o1-labs/dw/riscv32-alias-register

o1vm/riscv32: add alias for registers

o1vm/src/interpreters/riscv32im/registers.rs

@@ -64,3 +64,102 @@ impl<T: Clone> IndexMut<usize> for Registers<T> {
         }
     }
 }
+
+/// This enum provides aliases for the registers.
+/// This is useful for debugging and for providing a more readable interface.
+/// It can be used to index the registers in the witness.
+pub enum RegisterAlias {
+    Zero,
+    /// Return address
+    Ra,
+    /// Stack pointer
+    Sp,
+    /// Global pointer
+    Gp,
+    /// Thread pointer
+    Tp,
+    /// Temporary/alternate register
+    T0,
+    /// Temporaries
+    T1,
+    T2,
+    /// Frame pointer/saved register. This is the same register.
+    Fp,
+    S0,
+    /// Saved registers
+    S1,
+    /// Function arguments/results
+    A0,
+    A1,
+    A2,
+    A3,
+    A4,
+    A5,
+    A6,
+    A7,
+    S2,
+    S3,
+    S4,
+    S5,
+    S6,
+    S7,
+    S8,
+    S9,
+    S10,
+    S11,
+    T3,
+    T4,
+    T5,
+    T6,
+    /// Current instruction pointer
+    Ip,
+    /// Next instruction pointer
+    NextIp,
+    HeapPointer,
+}
+
+impl<T: Clone> Index<RegisterAlias> for Registers<T> {
+    type Output = T;
+
+    fn index(&self, index: RegisterAlias) -> &Self::Output {
+        match index {
+            RegisterAlias::Zero => &self.general_purpose[0],
+            RegisterAlias::Ra => &self.general_purpose[1],
+            RegisterAlias::Sp => &self.general_purpose[2],
+            RegisterAlias::Gp => &self.general_purpose[3],
+            RegisterAlias::Tp => &self.general_purpose[4],
+            RegisterAlias::T0 => &self.general_purpose[5],
+            RegisterAlias::T1 => &self.general_purpose[6],
+            RegisterAlias::T2 => &self.general_purpose[7],
+            // Frame pointer and first saved register are the same register.
+            RegisterAlias::Fp => &self.general_purpose[8],
+            RegisterAlias::S0 => &self.general_purpose[8],
+            RegisterAlias::S1 => &self.general_purpose[9],
+            RegisterAlias::A0 => &self.general_purpose[10],
+            RegisterAlias::A1 => &self.general_purpose[11],
+            RegisterAlias::A2 => &self.general_purpose[12],
+            RegisterAlias::A3 => &self.general_purpose[13],
+            RegisterAlias::A4 => &self.general_purpose[14],
+            RegisterAlias::A5 => &self.general_purpose[15],
+            RegisterAlias::A6 => &self.general_purpose[16],
+            RegisterAlias::A7 => &self.general_purpose[17],
+            RegisterAlias::S2 => &self.general_purpose[18],
+            RegisterAlias::S3 => &self.general_purpose[19],
+            RegisterAlias::S4 => &self.general_purpose[20],
+            RegisterAlias::S5 => &self.general_purpose[21],
+            RegisterAlias::S6 => &self.general_purpose[22],
+            RegisterAlias::S7 => &self.general_purpose[23],
+            RegisterAlias::S8 => &self.general_purpose[24],
+            RegisterAlias::S9 => &self.general_purpose[25],
+            RegisterAlias::S10 => &self.general_purpose[26],
+            RegisterAlias::S11 => &self.general_purpose[27],
+            RegisterAlias::T3 => &self.general_purpose[28],
+            RegisterAlias::T4 => &self.general_purpose[29],
+            RegisterAlias::T5 => &self.general_purpose[30],
+            RegisterAlias::T6 => &self.general_purpose[31],
+            RegisterAlias::Ip => &self.current_instruction_pointer,
+            RegisterAlias::NextIp => &self.next_instruction_pointer,
+            RegisterAlias::HeapPointer => &self.heap_pointer,
+        }
+    }
+}

o1vm/tests/test_riscv_elf.rs

@@ -1,10 +1,24 @@
 use mina_curves::pasta::Fp;
 use o1vm::interpreters::riscv32im::{
     interpreter::{IInstruction, Instruction, RInstruction},
+    registers::RegisterAlias::*,
     witness::Env,
     PAGE_SIZE,
 };
 
+#[test]
+fn test_registers_indexed_by_alias() {
+    let curr_dir = std::env::current_dir().unwrap();
+    let path = curr_dir.join(std::path::PathBuf::from(
+        "resources/programs/riscv32im/bin/sll",
+    ));
+    let state = o1vm::elf_loader::parse_riscv32(&path).unwrap();
+    let witness = Env::<Fp>::create(PAGE_SIZE.try_into().unwrap(), state);
+
+    assert_eq!(witness.registers[Ip], 65688);
+    assert_eq!(witness.registers[NextIp], 65692);
+}
+
 #[test]
 // Checking an instruction can be converted into a string.
 // It is mostly because we would want to use it to debug or write better error