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 [zero, one, of, pc, ssp, sp, fp, hp, err, ggas, cgas, bal, is, ret, retl, flag, rest @ ..] =
231        registers;
232    let r = SystemRegisters {
233        zero: RegMut(zero),
234        one: RegMut(one),
235        of: RegMut(of),
236        pc: RegMut(pc),
237        ssp: RegMut(ssp),
238        sp: RegMut(sp),
239        fp: RegMut(fp),
240        hp: RegMut(hp),
241        err: RegMut(err),
242        ggas: RegMut(ggas),
243        cgas: RegMut(cgas),
244        bal: RegMut(bal),
245        is: RegMut(is),
246        ret: RegMut(ret),
247        retl: RegMut(retl),
248        flag: RegMut(flag),
249    };
250    (r, ProgramRegisters(rest))
251}
252
253/// Copy the system and program registers into a single array.
254pub(crate) fn copy_registers(
255    system_registers: &SystemRegistersRef<'_>,
256    program_registers: &ProgramRegistersRef<'_>,
257) -> [Word; VM_REGISTER_COUNT] {
258    let mut out = [0u64; VM_REGISTER_COUNT];
259    out[..VM_REGISTER_SYSTEM_COUNT]
260        .copy_from_slice(&<[Word; VM_REGISTER_SYSTEM_COUNT]>::from(system_registers));
261    out[VM_REGISTER_SYSTEM_COUNT..].copy_from_slice(program_registers.0);
262    out
263}
264
265impl ProgramRegisters<'_> {
266    /// Get two mutable references to program registers.
267    /// Note they cannot be the same register.
268    pub fn get_mut_two(
269        &mut self,
270        a: WriteRegKey,
271        b: WriteRegKey,
272    ) -> Option<(&mut Word, &mut Word)> {
273        if a == b {
274            // Cannot mutably borrow the same register twice.
275            return None
276        }
277
278        // Order registers
279        let swap = a > b;
280        let (a, b) = if swap { (b, a) } else { (a, b) };
281
282        // Translate the absolute register indices to a program register indeces.
283        let a = a.translate();
284
285        // Subtract a + 1 because because we split the array at `a`.
286        let b = b
287            .translate()
288            .checked_sub(a.saturating_add(1))
289            .expect("Cannot underflow as the values are ordered");
290
291        // Split the array at the first register which is a.
292        let [i, rest @ ..] = &mut self.0[a..] else {
293            return None
294        };
295
296        // Translate the higher absolute register index to a program register index.
297        // Get the `b` register.
298        let j = &mut rest[b];
299
300        Some(if swap { (j, i) } else { (i, j) })
301    }
302}
303
304impl<'a> From<&'a SystemRegisters<'_>> for SystemRegistersRef<'a> {
305    fn from(value: &'a SystemRegisters<'_>) -> Self {
306        Self {
307            zero: Reg(value.zero.0),
308            one: Reg(value.one.0),
309            of: Reg(value.of.0),
310            pc: Reg(value.pc.0),
311            ssp: Reg(value.ssp.0),
312            sp: Reg(value.sp.0),
313            fp: Reg(value.fp.0),
314            hp: Reg(value.hp.0),
315            err: Reg(value.err.0),
316            ggas: Reg(value.ggas.0),
317            cgas: Reg(value.cgas.0),
318            bal: Reg(value.bal.0),
319            is: Reg(value.is.0),
320            ret: Reg(value.ret.0),
321            retl: Reg(value.retl.0),
322            flag: Reg(value.flag.0),
323        }
324    }
325}
326
327impl<'a> From<SystemRegisters<'a>> for SystemRegistersRef<'a> {
328    fn from(value: SystemRegisters<'a>) -> Self {
329        Self {
330            zero: Reg(value.zero.0),
331            one: Reg(value.one.0),
332            of: Reg(value.of.0),
333            pc: Reg(value.pc.0),
334            ssp: Reg(value.ssp.0),
335            sp: Reg(value.sp.0),
336            fp: Reg(value.fp.0),
337            hp: Reg(value.hp.0),
338            err: Reg(value.err.0),
339            ggas: Reg(value.ggas.0),
340            cgas: Reg(value.cgas.0),
341            bal: Reg(value.bal.0),
342            is: Reg(value.is.0),
343            ret: Reg(value.ret.0),
344            retl: Reg(value.retl.0),
345            flag: Reg(value.flag.0),
346        }
347    }
348}
349
350impl<'a> From<&'a ProgramRegisters<'_>> for ProgramRegistersRef<'a> {
351    fn from(value: &'a ProgramRegisters<'_>) -> Self {
352        Self(value.0)
353    }
354}
355
356impl<'a> From<ProgramRegisters<'a>> for ProgramRegistersRef<'a> {
357    fn from(value: ProgramRegisters<'a>) -> Self {
358        Self(value.0)
359    }
360}
361
362impl TryFrom<RegId> for WriteRegKey {
363    type Error = PanicReason;
364
365    fn try_from(r: RegId) -> Result<Self, Self::Error> {
366        Self::new(r)
367    }
368}
369
370impl core::ops::Index<WriteRegKey> for ProgramRegisters<'_> {
371    type Output = Word;
372
373    fn index(&self, index: WriteRegKey) -> &Self::Output {
374        &self.0[index.translate()]
375    }
376}
377
378impl core::ops::IndexMut<WriteRegKey> for ProgramRegisters<'_> {
379    fn index_mut(&mut self, index: WriteRegKey) -> &mut Self::Output {
380        &mut self.0[index.translate()]
381    }
382}
383
384impl<'a> From<&SystemRegistersRef<'a>> for [Word; VM_REGISTER_SYSTEM_COUNT] {
385    fn from(value: &SystemRegistersRef<'a>) -> Self {
386        let SystemRegistersRef {
387            zero,
388            one,
389            of,
390            pc,
391            ssp,
392            sp,
393            fp,
394            hp,
395            err,
396            ggas,
397            cgas,
398            bal,
399            is,
400            ret,
401            retl,
402            flag,
403        } = value;
404        [
405            *zero.0, *one.0, *of.0, *pc.0, *ssp.0, *sp.0, *fp.0, *hp.0, *err.0, *ggas.0,
406            *cgas.0, *bal.0, *is.0, *ret.0, *retl.0, *flag.0,
407        ]
408    }
409}
410
411#[derive(Debug, Clone, Copy)]
412pub(crate) enum ProgramRegistersSegment {
413    /// Registers 16..40
414    Low,
415    /// Registers 40..64
416    High,
417}
418
419impl ProgramRegisters<'_> {
420    /// Returns the registers corresponding to the segment, always 24 elements.
421    pub(crate) fn segment(&self, segment: ProgramRegistersSegment) -> &[Word] {
422        match segment {
423            ProgramRegistersSegment::Low => &self.0[..24],
424            ProgramRegistersSegment::High => &self.0[24..],
425        }
426    }
427
428    /// Returns the registers corresponding to the segment, always 24 elements.
429    pub(crate) fn segment_mut(
430        &mut self,
431        segment: ProgramRegistersSegment,
432    ) -> &mut [Word] {
433        match segment {
434            ProgramRegistersSegment::Low => &mut self.0[..24],
435            ProgramRegistersSegment::High => &mut self.0[24..],
436        }
437    }
438}
fuel_vm/constraints/reg_key.rs

fuel_vm/constraints/
reg_key.rs
