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 ProgramRegisters<'_> {
283    /// Get two mutable references to program registers.
284    /// Note they cannot be the same register.
285    pub fn get_mut_two(
286        &mut self,
287        a: WriteRegKey,
288        b: WriteRegKey,
289    ) -> Option<(&mut Word, &mut Word)> {
290        if a == b {
291            // Cannot mutably borrow the same register twice.
292            return None
293        }
294
295        // Order registers
296        let swap = a > b;
297        let (a, b) = if swap { (b, a) } else { (a, b) };
298
299        // Translate the absolute register indices to a program register indeces.
300        let a = a.translate();
301
302        // Subtract a + 1 because because we split the array at `a`.
303        let b = b
304            .translate()
305            .checked_sub(a.saturating_add(1))
306            .expect("Cannot underflow as the values are ordered");
307
308        // Split the array at the first register which is a.
309        let [i, rest @ ..] = &mut self.0[a..] else {
310            return None
311        };
312
313        // Translate the higher absolute register index to a program register index.
314        // Get the `b` register.
315        let j = &mut rest[b];
316
317        Some(if swap { (j, i) } else { (i, j) })
318    }
319}
320
321impl<'a> From<&'a SystemRegisters<'_>> for SystemRegistersRef<'a> {
322    fn from(value: &'a SystemRegisters<'_>) -> Self {
323        Self {
324            zero: Reg(value.zero.0),
325            one: Reg(value.one.0),
326            of: Reg(value.of.0),
327            pc: Reg(value.pc.0),
328            ssp: Reg(value.ssp.0),
329            sp: Reg(value.sp.0),
330            fp: Reg(value.fp.0),
331            hp: Reg(value.hp.0),
332            err: Reg(value.err.0),
333            ggas: Reg(value.ggas.0),
334            cgas: Reg(value.cgas.0),
335            bal: Reg(value.bal.0),
336            is: Reg(value.is.0),
337            ret: Reg(value.ret.0),
338            retl: Reg(value.retl.0),
339            flag: Reg(value.flag.0),
340        }
341    }
342}
343
344impl<'a> From<SystemRegisters<'a>> for SystemRegistersRef<'a> {
345    fn from(value: SystemRegisters<'a>) -> Self {
346        Self {
347            zero: Reg(value.zero.0),
348            one: Reg(value.one.0),
349            of: Reg(value.of.0),
350            pc: Reg(value.pc.0),
351            ssp: Reg(value.ssp.0),
352            sp: Reg(value.sp.0),
353            fp: Reg(value.fp.0),
354            hp: Reg(value.hp.0),
355            err: Reg(value.err.0),
356            ggas: Reg(value.ggas.0),
357            cgas: Reg(value.cgas.0),
358            bal: Reg(value.bal.0),
359            is: Reg(value.is.0),
360            ret: Reg(value.ret.0),
361            retl: Reg(value.retl.0),
362            flag: Reg(value.flag.0),
363        }
364    }
365}
366
367impl<'a> From<&'a ProgramRegisters<'_>> for ProgramRegistersRef<'a> {
368    fn from(value: &'a ProgramRegisters<'_>) -> Self {
369        Self(value.0)
370    }
371}
372
373impl<'a> From<ProgramRegisters<'a>> for ProgramRegistersRef<'a> {
374    fn from(value: ProgramRegisters<'a>) -> Self {
375        Self(value.0)
376    }
377}
378
379impl TryFrom<RegId> for WriteRegKey {
380    type Error = PanicReason;
381
382    fn try_from(r: RegId) -> Result<Self, Self::Error> {
383        Self::new(r)
384    }
385}
386
387impl core::ops::Index<WriteRegKey> for ProgramRegisters<'_> {
388    type Output = Word;
389
390    fn index(&self, index: WriteRegKey) -> &Self::Output {
391        &self.0[index.translate()]
392    }
393}
394
395impl core::ops::IndexMut<WriteRegKey> for ProgramRegisters<'_> {
396    fn index_mut(&mut self, index: WriteRegKey) -> &mut Self::Output {
397        &mut self.0[index.translate()]
398    }
399}
400
401impl<'a> From<&SystemRegistersRef<'a>> for [Word; VM_REGISTER_SYSTEM_COUNT] {
402    fn from(value: &SystemRegistersRef<'a>) -> Self {
403        let SystemRegistersRef {
404            zero,
405            one,
406            of,
407            pc,
408            ssp,
409            sp,
410            fp,
411            hp,
412            err,
413            ggas,
414            cgas,
415            bal,
416            is,
417            ret,
418            retl,
419            flag,
420        } = value;
421        [
422            *zero.0, *one.0, *of.0, *pc.0, *ssp.0, *sp.0, *fp.0, *hp.0, *err.0, *ggas.0,
423            *cgas.0, *bal.0, *is.0, *ret.0, *retl.0, *flag.0,
424        ]
425    }
426}
427
428#[derive(Debug, Clone, Copy)]
429pub(crate) enum ProgramRegistersSegment {
430    /// Registers 16..40
431    Low,
432    /// Registers 40..64
433    High,
434}
435
436impl ProgramRegisters<'_> {
437    /// Returns the registers corresponding to the segment, always 24 elements.
438    pub(crate) fn segment(&self, segment: ProgramRegistersSegment) -> &[Word] {
439        match segment {
440            ProgramRegistersSegment::Low => &self.0[..24],
441            ProgramRegistersSegment::High => &self.0[24..],
442        }
443    }
444
445    /// Returns the registers corresponding to the segment, always 24 elements.
446    pub(crate) fn segment_mut(
447        &mut self,
448        segment: ProgramRegistersSegment,
449    ) -> &mut [Word] {
450        match segment {
451            ProgramRegistersSegment::Low => &mut self.0[..24],
452            ProgramRegistersSegment::High => &mut self.0[24..],
453        }
454    }
455}
fuel_vm/constraints/reg_key.rs

fuel_vm/constraints/
reg_key.rs
