polkavm_assembler/
misc.rs

1use core::num::NonZeroU32;
2
3#[cfg(feature = "alloc")]
4use alloc::vec::Vec;
5
6#[derive(Copy, Clone, PartialEq, Eq, Debug)]
7#[repr(transparent)]
8pub struct Label(NonZeroU32);
9
10impl Label {
11    #[inline]
12    pub fn raw(self) -> u32 {
13        self.0.get() - 1
14    }
15
16    #[inline]
17    pub fn from_raw(value: u32) -> Self {
18        Label(NonZeroU32::new(value + 1).unwrap())
19    }
20}
21
22impl core::fmt::Display for Label {
23    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
24        fmt.write_fmt(core::format_args!("<{}>", self.0))
25    }
26}
27
28#[derive(Copy, Clone)]
29pub struct Instruction<T> {
30    pub(crate) instruction: T,
31    pub(crate) bytes: InstBuf,
32
33    #[cfg_attr(not(feature = "alloc"), allow(dead_code))]
34    pub(crate) fixup: Option<(Label, FixupKind)>,
35}
36
37impl<T> core::fmt::Debug for Instruction<T>
38where
39    T: core::fmt::Debug,
40{
41    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
42        self.instruction.fmt(fmt)
43    }
44}
45
46impl<T> core::fmt::Display for Instruction<T>
47where
48    T: core::fmt::Display,
49{
50    fn fmt(&self, fmt: &mut core::fmt::Formatter) -> core::fmt::Result {
51        self.instruction.fmt(fmt)
52    }
53}
54
55impl<T> Instruction<T> {
56    #[allow(clippy::len_without_is_empty)]
57    #[inline]
58    pub fn len(&self) -> usize {
59        self.bytes.len()
60    }
61}
62
63#[derive(Copy, Clone)]
64#[repr(transparent)]
65pub(crate) struct FixupKind(pub u32);
66
67impl FixupKind {
68    #[cfg_attr(not(feature = "alloc"), allow(dead_code))]
69    #[inline]
70    pub const fn offset(self) -> u32 {
71        (self.0 >> 24) & 0b11
72    }
73
74    #[cfg_attr(not(feature = "alloc"), allow(dead_code))]
75    #[inline]
76    pub const fn length(self) -> u32 {
77        self.0 >> 28
78    }
79
80    #[inline]
81    pub const fn new_1(opcode: u32, length: u32) -> Self {
82        FixupKind((1 << 24) | (length << 28) | opcode)
83    }
84
85    #[inline]
86    pub const fn new_2(opcode: [u32; 2], length: u32) -> Self {
87        let opcode = opcode[0] | (opcode[1] << 8);
88        FixupKind((2 << 24) | (length << 28) | opcode)
89    }
90
91    #[inline]
92    pub const fn new_3(opcode: [u32; 3], length: u32) -> Self {
93        let opcode = opcode[0] | (opcode[1] << 8) | (opcode[2] << 16);
94        FixupKind((3 << 24) | (length << 28) | opcode)
95    }
96}
97
98const MAXIMUM_INSTRUCTION_SIZE: usize = 16;
99
100#[derive(Copy, Clone)]
101pub struct InstBuf {
102    out: u128,
103    length: u32,
104}
105
106#[allow(clippy::new_without_default)]
107impl InstBuf {
108    #[inline]
109    pub fn new() -> Self {
110        Self { out: 0, length: 0 }
111    }
112
113    #[inline]
114    pub fn len(&self) -> usize {
115        (self.length >> 3) as usize
116    }
117
118    #[inline]
119    pub fn append(&mut self, byte: u8) {
120        self.out |= u128::from(byte).wrapping_shl(self.length);
121        self.length += 8;
122    }
123
124    #[inline]
125    pub fn append_packed_bytes(&mut self, value: u32, length: u32) {
126        self.out |= u128::from(value).wrapping_shl(self.length);
127        self.length += length;
128    }
129
130    #[cfg(feature = "alloc")]
131    #[inline]
132    unsafe fn encode_into_raw(self, output: *mut u8) {
133        core::ptr::write_unaligned(output.cast::<u64>(), u64::from_le(self.out as u64));
134        core::ptr::write_unaligned(output.add(8).cast::<u64>(), u64::from_le((self.out >> 64) as u64));
135    }
136
137    #[cfg(feature = "alloc")]
138    #[allow(clippy::debug_assert_with_mut_call)]
139    #[inline]
140    pub unsafe fn encode_into_vec_unsafe(self, output: &mut Vec<u8>) {
141        debug_assert!(output.spare_capacity_mut().len() >= MAXIMUM_INSTRUCTION_SIZE);
142
143        self.encode_into_raw(output.spare_capacity_mut().as_mut_ptr().cast());
144        let new_length = output.len() + (self.length as usize >> 3);
145        output.set_len(new_length);
146    }
147
148    #[cfg(feature = "alloc")]
149    #[cold]
150    #[inline(never)]
151    fn reserve_impl(output: &mut Vec<u8>, length: usize) {
152        output.reserve(length);
153    }
154
155    #[cfg(feature = "alloc")]
156    #[inline(always)]
157    pub fn reserve_const<const INSTRUCTIONS: usize>(output: &mut Vec<u8>) {
158        Self::reserve(output, INSTRUCTIONS);
159    }
160
161    #[cfg(feature = "alloc")]
162    #[inline(always)]
163    pub fn reserve(output: &mut Vec<u8>, count: usize) {
164        let count = count.checked_mul(MAXIMUM_INSTRUCTION_SIZE).unwrap();
165        if output.spare_capacity_mut().len() < count {
166            Self::reserve_impl(output, count);
167            if output.spare_capacity_mut().len() < count {
168                // SAFETY: `reserve` made sure that we have this much capacity, so this is safe.
169                unsafe {
170                    core::hint::unreachable_unchecked();
171                }
172            }
173        }
174    }
175
176    #[inline]
177    pub fn from_array<const N: usize>(array: [u8; N]) -> Self {
178        if N > MAXIMUM_INSTRUCTION_SIZE {
179            panic!();
180        }
181
182        let mut out = Self::new();
183        for value in array {
184            out.append(value);
185        }
186        out
187    }
188
189    #[cfg(feature = "alloc")]
190    pub fn to_vec(self) -> Vec<u8> {
191        let mut vec = Vec::with_capacity(MAXIMUM_INSTRUCTION_SIZE);
192
193        // SAFETY: We've reserved space for at least one instruction.
194        unsafe {
195            self.encode_into_vec_unsafe(&mut vec);
196        }
197
198        vec
199    }
200}
201
202#[cfg(feature = "alloc")]
203#[test]
204fn test_inst_buf() {
205    assert_eq!(InstBuf::from_array([0x01]).to_vec(), [0x01]);
206    assert_eq!(InstBuf::from_array([0x01, 0x02]).to_vec(), [0x01, 0x02]);
207    assert_eq!(InstBuf::from_array([0x01, 0x02, 0x03]).to_vec(), [0x01, 0x02, 0x03]);
208    assert_eq!(InstBuf::from_array([0x01, 0x02, 0x03, 0x04]).to_vec(), [0x01, 0x02, 0x03, 0x04]);
209    assert_eq!(
210        InstBuf::from_array([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]).to_vec(),
211        [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08]
212    );
213    assert_eq!(
214        InstBuf::from_array([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]).to_vec(),
215        [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]
216    );
217    assert_eq!(
218        InstBuf::from_array([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A]).to_vec(),
219        [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A]
220    );
221
222    let mut buf = InstBuf::from_array([0x01]);
223    assert_eq!(buf.to_vec(), [0x01]);
224    buf.append_packed_bytes(0x05040302, 32);
225    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03, 0x04, 0x05]);
226    buf.append_packed_bytes(0x09080706, 32);
227    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]);
228
229    let mut buf = InstBuf::from_array([0x01]);
230    assert_eq!(buf.to_vec(), [0x01]);
231    buf.append_packed_bytes(0x0302, 16);
232    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03]);
233    buf.append_packed_bytes(0x0504, 16);
234    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03, 0x04, 0x05]);
235    buf.append_packed_bytes(0x0706, 16);
236    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07]);
237    buf.append_packed_bytes(0x0908, 16);
238    assert_eq!(buf.to_vec(), [0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09]);
239}
polkavm_assembler/misc.rs

polkavm_assembler/
misc.rs
