o1vm/interpreters/riscv32im/
registers.rs

1use std::ops::{Index, IndexMut};
2
3use serde::{Deserialize, Serialize};
4
5pub const N_GP_REGISTERS: usize = 32;
6// FIXME:
7pub const REGISTER_CURRENT_IP: usize = N_GP_REGISTERS + 1;
8pub const REGISTER_NEXT_IP: usize = N_GP_REGISTERS + 2;
9pub const REGISTER_HEAP_POINTER: usize = N_GP_REGISTERS + 3;
10
11/// This represents the internal state of the virtual machine.
12#[derive(Clone, Default, Debug, Serialize, Deserialize)]
13pub struct Registers<T> {
14    /// There are 32 general purpose registers.
15    /// - x0: hard-wired zero
16    /// - x1: return address
17    /// - x2: stack pointer
18    /// - x3: global pointer
19    /// - x4: thread pointer
20    /// - x5: temporary/alternate register
21    /// - x6-x7: temporaries
22    /// - x8: saved register/frame pointer
23    /// - x9: saved register
24    /// - x10-x11: function arguments/results
25    /// - x12-x17: function arguments
26    /// - x18-x27: saved registers
27    /// - x28-x31: temporaries
28    pub general_purpose: [T; N_GP_REGISTERS],
29    pub current_instruction_pointer: T,
30    pub next_instruction_pointer: T,
31    pub heap_pointer: T,
32}
33
34impl<T: Clone> Index<usize> for Registers<T> {
35    type Output = T;
36
37    fn index(&self, index: usize) -> &Self::Output {
38        if index < N_GP_REGISTERS {
39            &self.general_purpose[index]
40        } else if index == REGISTER_CURRENT_IP {
41            &self.current_instruction_pointer
42        } else if index == REGISTER_NEXT_IP {
43            &self.next_instruction_pointer
44        } else if index == REGISTER_HEAP_POINTER {
45            &self.heap_pointer
46        } else {
47            panic!("Index out of bounds");
48        }
49    }
50}
51
52impl<T: Clone> IndexMut<usize> for Registers<T> {
53    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
54        if index < N_GP_REGISTERS {
55            &mut self.general_purpose[index]
56        } else if index == REGISTER_CURRENT_IP {
57            &mut self.current_instruction_pointer
58        } else if index == REGISTER_NEXT_IP {
59            &mut self.next_instruction_pointer
60        } else if index == REGISTER_HEAP_POINTER {
61            &mut self.heap_pointer
62        } else {
63            panic!("Index out of bounds");
64        }
65    }
66}
67
68/// This enum provides aliases for the registers.
69/// This is useful for debugging and for providing a more readable interface.
70/// It can be used to index the registers in the witness.
71pub enum RegisterAlias {
72    Zero,
73    /// Return address
74    Ra,
75    /// Stack pointer
76    Sp,
77    /// Global pointer
78    Gp,
79    /// Thread pointer
80    Tp,
81    /// Temporary/alternate register
82    T0,
83    /// Temporaries
84    T1,
85    T2,
86    /// Frame pointer/saved register. This is the same register.
87    Fp,
88    S0,
89    /// Saved registers
90    S1,
91    /// Function arguments/results
92    A0,
93    A1,
94    A2,
95    A3,
96    A4,
97    A5,
98    A6,
99    A7,
100    S2,
101    S3,
102    S4,
103    S5,
104    S6,
105    S7,
106    S8,
107    S9,
108    S10,
109    S11,
110    T3,
111    T4,
112    T5,
113    T6,
114    /// Current instruction pointer
115    Ip,
116    /// Next instruction pointer
117    NextIp,
118    HeapPointer,
119}
120
121impl<T: Clone> Index<RegisterAlias> for Registers<T> {
122    type Output = T;
123
124    fn index(&self, index: RegisterAlias) -> &Self::Output {
125        match index {
126            RegisterAlias::Zero => &self.general_purpose[0],
127            RegisterAlias::Ra => &self.general_purpose[1],
128            RegisterAlias::Sp => &self.general_purpose[2],
129            RegisterAlias::Gp => &self.general_purpose[3],
130            RegisterAlias::Tp => &self.general_purpose[4],
131            RegisterAlias::T0 => &self.general_purpose[5],
132            RegisterAlias::T1 => &self.general_purpose[6],
133            RegisterAlias::T2 => &self.general_purpose[7],
134            // Frame pointer and first saved register are the same register.
135            RegisterAlias::Fp => &self.general_purpose[8],
136            RegisterAlias::S0 => &self.general_purpose[8],
137            RegisterAlias::S1 => &self.general_purpose[9],
138            RegisterAlias::A0 => &self.general_purpose[10],
139            RegisterAlias::A1 => &self.general_purpose[11],
140            RegisterAlias::A2 => &self.general_purpose[12],
141            RegisterAlias::A3 => &self.general_purpose[13],
142            RegisterAlias::A4 => &self.general_purpose[14],
143            RegisterAlias::A5 => &self.general_purpose[15],
144            RegisterAlias::A6 => &self.general_purpose[16],
145            RegisterAlias::A7 => &self.general_purpose[17],
146            RegisterAlias::S2 => &self.general_purpose[18],
147            RegisterAlias::S3 => &self.general_purpose[19],
148            RegisterAlias::S4 => &self.general_purpose[20],
149            RegisterAlias::S5 => &self.general_purpose[21],
150            RegisterAlias::S6 => &self.general_purpose[22],
151            RegisterAlias::S7 => &self.general_purpose[23],
152            RegisterAlias::S8 => &self.general_purpose[24],
153            RegisterAlias::S9 => &self.general_purpose[25],
154            RegisterAlias::S10 => &self.general_purpose[26],
155            RegisterAlias::S11 => &self.general_purpose[27],
156            RegisterAlias::T3 => &self.general_purpose[28],
157            RegisterAlias::T4 => &self.general_purpose[29],
158            RegisterAlias::T5 => &self.general_purpose[30],
159            RegisterAlias::T6 => &self.general_purpose[31],
160            RegisterAlias::Ip => &self.current_instruction_pointer,
161            RegisterAlias::NextIp => &self.next_instruction_pointer,
162            RegisterAlias::HeapPointer => &self.heap_pointer,
163        }
164    }
165}