Skip to main content

fuel_vm/constraints/
reg_key.rs

1//! Utilities for accessing register values and proving at compile time that
2//! the register index is valid.
3//!
4//! This module also provides utilities for mutably accessing multiple registers.
5use core::ops::{
6    Deref,
7    DerefMut,
8};
9
10use fuel_asm::{
11    PanicReason,
12    RegId,
13    Word,
14};
15
16use crate::consts::{
17    VM_REGISTER_COUNT,
18    VM_REGISTER_PROGRAM_COUNT,
19    VM_REGISTER_SYSTEM_COUNT,
20};
21
22#[cfg(test)]
23mod tests;
24
25#[derive(Debug, PartialEq, Eq)]
26/// Mutable reference to a register value at a given index.
27pub struct RegMut<'r, const INDEX: u8>(&'r mut Word);
28
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30/// Immutable reference to a register value at a given index.
31pub struct Reg<'r, const INDEX: u8>(&'r Word);
32
33#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
34/// A key to a writable register that is within
35/// the bounds of the writable registers.
36pub struct WriteRegKey(RegId);
37
38impl WriteRegKey {
39    /// Create a new writable register key if the index is within the bounds
40    /// of the writable registers.
41    pub fn new(k: impl Into<RegId>) -> Result<Self, PanicReason> {
42        let k = k.into();
43
44        if k >= RegId::WRITABLE {
45            Ok(Self(k))
46        } else {
47            Err(PanicReason::ReservedRegisterNotWritable)
48        }
49    }
50
51    /// Translate this key from an absolute register index
52    /// to a program register index.
53    ///
54    /// This subtracts the number of system registers from the key.
55    #[allow(clippy::arithmetic_side_effects)] // Safety: checked in constructor
56    fn translate(self) -> usize {
57        self.0.to_u8() as usize - VM_REGISTER_SYSTEM_COUNT
58    }
59}
60
61impl<'r, const INDEX: u8> RegMut<'r, INDEX> {
62    /// Create a new mutable register reference.
63    pub fn new(reg: &'r mut Word) -> Self {
64        Self(reg)
65    }
66}
67
68impl<'r, const INDEX: u8> Reg<'r, INDEX> {
69    /// Create a new immutable register reference.
70    pub fn new(reg: &'r Word) -> Self {
71        Self(reg)
72    }
73}
74
75impl<const INDEX: u8> Deref for Reg<'_, INDEX> {
76    type Target = Word;
77
78    fn deref(&self) -> &Self::Target {
79        self.0
80    }
81}
82
83impl<const INDEX: u8> Deref for RegMut<'_, INDEX> {
84    type Target = Word;
85
86    fn deref(&self) -> &Self::Target {
87        self.0
88    }
89}
90
91impl<const INDEX: u8> DerefMut for RegMut<'_, INDEX> {
92    fn deref_mut(&mut self) -> &mut Self::Target {
93        self.0
94    }
95}
96
97impl<'a, const INDEX: u8> From<RegMut<'a, INDEX>> for Reg<'a, INDEX> {
98    fn from(reg: RegMut<'a, INDEX>) -> Self {
99        Self(reg.0)
100    }
101}
102
103impl<const INDEX: u8> RegMut<'_, INDEX> {
104    /// Re-borrow the register as an immutable reference.
105    pub fn as_ref(&self) -> Reg<'_, INDEX> {
106        Reg(self.0)
107    }
108}
109
110impl<const INDEX: u8> RegMut<'_, INDEX> {
111    /// Re-borrow the register as a mutable reference.
112    pub fn as_mut(&mut self) -> RegMut<'_, INDEX> {
113        RegMut(self.0)
114    }
115}
116
117macro_rules! impl_keys {
118    ( $($i:ident, $f:ident $(,$f_mut:ident)?)* ) => {
119        $(
120            #[doc = "Register index key for use with Reg and RegMut."]
121            pub const $i: u8 = RegId::$i.to_u8();
122        )*
123        #[doc = "Get register reference by name."]
124        pub trait GetReg {
125        $(
126            #[doc = "Get register reference for this key."]
127            fn $f(&self) -> Reg<'_, $i>;
128        )*
129        }
130        #[doc = "Get register mutable reference by name."]
131        pub trait GetRegMut {
132        $(
133            $(
134            #[doc = "Get mutable register reference for this key."]
135            fn $f_mut(&mut self) -> RegMut<'_, $i>;
136            )?
137        )*
138        }
139        impl GetReg for [Word; VM_REGISTER_COUNT] {
140        $(
141            fn $f(&self) -> Reg<'_, $i> {
142                Reg(&self[$i as usize])
143            }
144        )*
145        }
146        impl GetRegMut for [Word; VM_REGISTER_COUNT] {
147        $(
148            $(
149            fn $f_mut(&mut self) -> RegMut<'_, $i> {
150                RegMut(&mut self[$i as usize])
151            }
152            )?
153        )*
154        }
155    };
156}
157
158impl_keys! {
159    ZERO, zero
160    ONE, one
161    OF, of, of_mut
162    PC, pc, pc_mut
163    SSP, ssp, ssp_mut
164    SP, sp, sp_mut
165    FP, fp, fp_mut
166    HP, hp, hp_mut
167    ERR, err, err_mut
168    GGAS, ggas, ggas_mut
169    CGAS, cgas, cgas_mut
170    BAL, bal, bal_mut
171    IS, is, is_mut
172    RET, ret, ret_mut
173    RETL, retl, retl_mut
174    FLAG, flag, flag_mut
175}
176
177/// The set of system registers split into
178/// individual mutable references.
179pub(crate) struct SystemRegisters<'a> {
180    pub(crate) zero: RegMut<'a, ZERO>,
181    pub(crate) one: RegMut<'a, ONE>,
182    pub(crate) of: RegMut<'a, OF>,
183    pub(crate) pc: RegMut<'a, PC>,
184    pub(crate) ssp: RegMut<'a, SSP>,
185    pub(crate) sp: RegMut<'a, SP>,
186    pub(crate) fp: RegMut<'a, FP>,
187    pub(crate) hp: RegMut<'a, HP>,
188    pub(crate) err: RegMut<'a, ERR>,
189    pub(crate) ggas: RegMut<'a, GGAS>,
190    pub(crate) cgas: RegMut<'a, CGAS>,
191    pub(crate) bal: RegMut<'a, BAL>,
192    pub(crate) is: RegMut<'a, IS>,
193    pub(crate) ret: RegMut<'a, RET>,
194    pub(crate) retl: RegMut<'a, RETL>,
195    pub(crate) flag: RegMut<'a, FLAG>,
196}
197
198/// Same as `SystemRegisters` but with immutable references.
199pub(crate) struct SystemRegistersRef<'a> {
200    pub(crate) zero: Reg<'a, ZERO>,
201    pub(crate) one: Reg<'a, ONE>,
202    pub(crate) of: Reg<'a, OF>,
203    pub(crate) pc: Reg<'a, PC>,
204    pub(crate) ssp: Reg<'a, SSP>,
205    pub(crate) sp: Reg<'a, SP>,
206    pub(crate) fp: Reg<'a, FP>,
207    pub(crate) hp: Reg<'a, HP>,
208    pub(crate) err: Reg<'a, ERR>,
209    pub(crate) ggas: Reg<'a, GGAS>,
210    pub(crate) cgas: Reg<'a, CGAS>,
211    pub(crate) bal: Reg<'a, BAL>,
212    pub(crate) is: Reg<'a, IS>,
213    pub(crate) ret: Reg<'a, RET>,
214    pub(crate) retl: Reg<'a, RETL>,
215    pub(crate) flag: Reg<'a, FLAG>,
216}
217
218/// The set of program registers split from the system registers.
219pub(crate) struct ProgramRegisters<'a>(pub &'a mut [Word; VM_REGISTER_PROGRAM_COUNT]);
220
221/// Same as `ProgramRegisters` but with immutable references.
222pub(crate) struct ProgramRegistersRef<'a>(pub &'a [Word; VM_REGISTER_PROGRAM_COUNT]);
223
224/// Split the registers into system and program registers.
225///
226/// This allows multiple mutable references to registers.
227pub(crate) fn split_registers(
228    registers: &mut [Word; VM_REGISTER_COUNT],
229) -> (SystemRegisters<'_>, ProgramRegisters<'_>) {
230    let [
231        zero,
232        one,
233        of,
234        pc,
235        ssp,
236        sp,
237        fp,
238        hp,
239        err,
240        ggas,
241        cgas,
242        bal,
243        is,
244        ret,
245        retl,
246        flag,
247        rest @ ..,
248    ] = registers;
249    let r = SystemRegisters {
250        zero: RegMut(zero),
251        one: RegMut(one),
252        of: RegMut(of),
253        pc: RegMut(pc),
254        ssp: RegMut(ssp),
255        sp: RegMut(sp),
256        fp: RegMut(fp),
257        hp: RegMut(hp),
258        err: RegMut(err),
259        ggas: RegMut(ggas),
260        cgas: RegMut(cgas),
261        bal: RegMut(bal),
262        is: RegMut(is),
263        ret: RegMut(ret),
264        retl: RegMut(retl),
265        flag: RegMut(flag),
266    };
267    (r, ProgramRegisters(rest))
268}
269
270/// Copy the system and program registers into a single array.
271pub(crate) fn copy_registers(
272    system_registers: &SystemRegistersRef<'_>,
273    program_registers: &ProgramRegistersRef<'_>,
274) -> [Word; VM_REGISTER_COUNT] {
275    let mut out = [0u64; VM_REGISTER_COUNT];
276    out[..VM_REGISTER_SYSTEM_COUNT]
277        .copy_from_slice(&<[Word; VM_REGISTER_SYSTEM_COUNT]>::from(system_registers));
278    out[VM_REGISTER_SYSTEM_COUNT..].copy_from_slice(program_registers.0);
279    out
280}
281
282impl<'a> From<&'a SystemRegisters<'_>> for SystemRegistersRef<'a> {
283    fn from(value: &'a SystemRegisters<'_>) -> Self {
284        Self {
285            zero: Reg(value.zero.0),
286            one: Reg(value.one.0),
287            of: Reg(value.of.0),
288            pc: Reg(value.pc.0),
289            ssp: Reg(value.ssp.0),
290            sp: Reg(value.sp.0),
291            fp: Reg(value.fp.0),
292            hp: Reg(value.hp.0),
293            err: Reg(value.err.0),
294            ggas: Reg(value.ggas.0),
295            cgas: Reg(value.cgas.0),
296            bal: Reg(value.bal.0),
297            is: Reg(value.is.0),
298            ret: Reg(value.ret.0),
299            retl: Reg(value.retl.0),
300            flag: Reg(value.flag.0),
301        }
302    }
303}
304
305impl<'a> From<SystemRegisters<'a>> for SystemRegistersRef<'a> {
306    fn from(value: SystemRegisters<'a>) -> Self {
307        Self {
308            zero: Reg(value.zero.0),
309            one: Reg(value.one.0),
310            of: Reg(value.of.0),
311            pc: Reg(value.pc.0),
312            ssp: Reg(value.ssp.0),
313            sp: Reg(value.sp.0),
314            fp: Reg(value.fp.0),
315            hp: Reg(value.hp.0),
316            err: Reg(value.err.0),
317            ggas: Reg(value.ggas.0),
318            cgas: Reg(value.cgas.0),
319            bal: Reg(value.bal.0),
320            is: Reg(value.is.0),
321            ret: Reg(value.ret.0),
322            retl: Reg(value.retl.0),
323            flag: Reg(value.flag.0),
324        }
325    }
326}
327
328impl<'a> From<&'a ProgramRegisters<'_>> for ProgramRegistersRef<'a> {
329    fn from(value: &'a ProgramRegisters<'_>) -> Self {
330        Self(value.0)
331    }
332}
333
334impl<'a> From<ProgramRegisters<'a>> for ProgramRegistersRef<'a> {
335    fn from(value: ProgramRegisters<'a>) -> Self {
336        Self(value.0)
337    }
338}
339
340impl TryFrom<RegId> for WriteRegKey {
341    type Error = PanicReason;
342
343    fn try_from(r: RegId) -> Result<Self, Self::Error> {
344        Self::new(r)
345    }
346}
347
348impl core::ops::Index<WriteRegKey> for ProgramRegisters<'_> {
349    type Output = Word;
350
351    fn index(&self, index: WriteRegKey) -> &Self::Output {
352        &self.0[index.translate()]
353    }
354}
355
356impl core::ops::IndexMut<WriteRegKey> for ProgramRegisters<'_> {
357    fn index_mut(&mut self, index: WriteRegKey) -> &mut Self::Output {
358        &mut self.0[index.translate()]
359    }
360}
361
362impl<'a> From<&SystemRegistersRef<'a>> for [Word; VM_REGISTER_SYSTEM_COUNT] {
363    fn from(value: &SystemRegistersRef<'a>) -> Self {
364        let SystemRegistersRef {
365            zero,
366            one,
367            of,
368            pc,
369            ssp,
370            sp,
371            fp,
372            hp,
373            err,
374            ggas,
375            cgas,
376            bal,
377            is,
378            ret,
379            retl,
380            flag,
381        } = value;
382        [
383            *zero.0, *one.0, *of.0, *pc.0, *ssp.0, *sp.0, *fp.0, *hp.0, *err.0, *ggas.0,
384            *cgas.0, *bal.0, *is.0, *ret.0, *retl.0, *flag.0,
385        ]
386    }
387}
388
389#[derive(Debug, Clone, Copy)]
390pub(crate) enum ProgramRegistersSegment {
391    /// Registers 16..40
392    Low,
393    /// Registers 40..64
394    High,
395}
396
397impl ProgramRegisters<'_> {
398    /// Returns the registers corresponding to the segment, always 24 elements.
399    pub(crate) fn segment(&self, segment: ProgramRegistersSegment) -> &[Word] {
400        match segment {
401            ProgramRegistersSegment::Low => &self.0[..24],
402            ProgramRegistersSegment::High => &self.0[24..],
403        }
404    }
405
406    /// Returns the registers corresponding to the segment, always 24 elements.
407    pub(crate) fn segment_mut(
408        &mut self,
409        segment: ProgramRegistersSegment,
410    ) -> &mut [Word] {
411        match segment {
412            ProgramRegistersSegment::Low => &mut self.0[..24],
413            ProgramRegistersSegment::High => &mut self.0[24..],
414        }
415    }
416}