juicebox_asm/
asm.rs

1// SPDX-License-Identifier: MIT
2//
3// Copyright (c) 2023, Johannes Stoelp <dev@memzero.de>
4
5//! The `x64` jit assembler.
6
7use crate::imm::Imm;
8use crate::mem::{AddrMode, Mem, Mem16, Mem32, Mem64, Mem8};
9use crate::reg::{Reg, Reg16, Reg32, Reg64, Reg8};
10use crate::Label;
11
12/// Encode the `REX` byte.
13pub(crate) const fn rex(w: bool, r: u8, x: u8, b: u8) -> u8 {
14    let w = if w { 1 } else { 0 };
15    let r = (r >> 3) & 1;
16    let x = (x >> 3) & 1;
17    let b = (b >> 3) & 1;
18    0b0100_0000 | ((w & 1) << 3) | (r << 2) | (x << 1) | b
19}
20
21/// Encode the `ModR/M` byte.
22pub(crate) const fn modrm(mod_: u8, reg: u8, rm: u8) -> u8 {
23    ((mod_ & 0b11) << 6) | ((reg & 0b111) << 3) | (rm & 0b111)
24}
25
26/// Encode the `SIB` byte.
27const fn sib(scale: u8, index: u8, base: u8) -> u8 {
28    ((scale & 0b11) << 6) | ((index & 0b111) << 3) | (base & 0b111)
29}
30
31/// `x64` jit assembler.
32pub struct Asm {
33    buf: Vec<u8>,
34}
35
36impl Asm {
37    /// Create a new `x64` jit assembler.
38    pub fn new() -> Asm {
39        // Some random default capacity.
40        let buf = Vec::with_capacity(1024);
41        Asm { buf }
42    }
43
44    /// Consume the assembler and get the emitted code.
45    pub fn into_code(self) -> Vec<u8> {
46        self.buf
47    }
48
49    /// Disassemble the code currently added to the assembler, using
50    /// [`ndisasm`](https://nasm.us/index.php) and print it to _stdout_. If
51    /// `ndisasm` is not available on the system this prints a warning and
52    /// becomes a nop.
53    ///
54    /// # Panics
55    ///
56    /// Panics if anything goes wrong with spawning, writing to or reading from
57    /// the `ndisasm` child process.
58    pub fn disasm(&self) {
59        crate::disasm::disasm(&self.buf);
60    }
61
62    /// Emit a slice of bytes.
63    pub(crate) fn emit(&mut self, bytes: &[u8]) {
64        self.buf.extend_from_slice(bytes);
65    }
66
67    /// Emit a slice of optional bytes.
68    pub(crate) fn emit_optional(&mut self, bytes: &[Option<u8>]) {
69        for byte in bytes.iter().filter_map(|&b| b) {
70            self.buf.push(byte);
71        }
72    }
73
74    /// Emit a slice of bytes at `pos`.
75    ///
76    /// # Panics
77    ///
78    /// Panics if [pos..pos+len] indexes out of bound of the underlying code buffer.
79    fn emit_at(&mut self, pos: usize, bytes: &[u8]) {
80        if let Some(buf) = self.buf.get_mut(pos..pos + bytes.len()) {
81            buf.copy_from_slice(bytes);
82        } else {
83            unimplemented!();
84        }
85    }
86
87    /// Bind the [Label] to the current location.
88    pub fn bind(&mut self, label: &mut Label) {
89        // Bind the label to the current offset.
90        label.bind(self.buf.len());
91
92        // Resolve any pending relocations for the label.
93        self.resolve(label);
94    }
95
96    /// If the [Label] is bound, patch any pending relocation.
97    fn resolve(&mut self, label: &mut Label) {
98        if let Some(loc) = label.location() {
99            // Resolve any pending relocations for the label.
100            for off in label.offsets_mut().drain() {
101                // Displacement is relative to the next instruction following the jump.
102                let disp = {
103                    let loc = isize::try_from(loc).expect("loc does not fit into isize");
104                    let off = isize::try_from(off).expect("off does not fit into isize");
105
106                    // We record the offset to patch at the first byte of the disp32
107                    // therefore we need to account for that in the disp computation.
108                    loc - (off + 4/* account for the disp32 */)
109                };
110
111                // For now we only support disp32 as label location.
112                let disp32 = i32::try_from(disp).expect("Label offset did not fit into i32");
113
114                // Patch the relocation with the disp32.
115                self.emit_at(off, &disp32.to_ne_bytes());
116            }
117        }
118    }
119
120    // -- Encode utilities.
121
122    /// Encode an offset-immediate instruction.
123    /// Register idx is encoded in the opcode.
124    pub(crate) fn encode_oi<T: Reg, U: Imm>(&mut self, opc: u8, op1: T, op2: U)
125    where
126        Self: EncodeR<T>,
127    {
128        let opc = opc + (op1.idx() & 0b111);
129        let prefix = <Self as EncodeR<T>>::legacy_prefix();
130        let rex = <Self as EncodeR<T>>::rex(op1);
131
132        self.emit_optional(&[prefix, rex]);
133        self.emit(&[opc]);
134        self.emit(op2.bytes());
135    }
136
137    /// Encode a register instruction.
138    pub(crate) fn encode_r<T: Reg>(&mut self, opc: &[u8], opc_ext: u8, op1: T)
139    where
140        Self: EncodeR<T>,
141    {
142        // M operand encoding.
143        //   op1           -> modrm.rm
144        //   opc extension -> modrm.reg
145        let modrm = modrm(
146            0b11,      /* mod */
147            opc_ext,   /* reg */
148            op1.idx(), /* rm */
149        );
150
151        let prefix = <Self as EncodeR<T>>::legacy_prefix();
152        let rex = <Self as EncodeR<T>>::rex(op1);
153
154        self.emit_optional(&[prefix, rex]);
155        self.emit(opc);
156        self.emit(&[modrm]);
157    }
158
159    /// Encode a register-register instruction with MR operand encoding,
160    /// for example "<ins> r/m64, r64".
161    pub(crate) fn encode_rr_mr<T: Reg, U: Reg>(&mut self, opc: &[u8], op1: T, op2: U)
162    where
163        Self: EncodeRR<T, U>,
164    {
165        // RR as MR operand encoding, eg for "<ins> r/m64, r64".
166        //   op1 -> modrm.rm
167        //   op2 -> modrm.reg
168        //
169        // NOTE: There are also RR with RM encoding "<ins> r64 r/m64", caution!
170        let modrm = modrm(
171            0b11,      /* mod */
172            op2.idx(), /* reg */
173            op1.idx(), /* rm */
174        );
175
176        let prefix = <Self as EncodeRR<T, U>>::legacy_prefix();
177        let rex = <Self as EncodeRR<T, U>>::rex(op1, op2);
178
179        self.emit_optional(&[prefix, rex]);
180        self.emit(opc);
181        self.emit(&[modrm]);
182    }
183
184    /// Encode a register-register instruction with RM operand encoding,
185    /// for example "<ins> r64, r/m64".
186    pub(crate) fn encode_rr_rm<T: Reg, U: Reg>(&mut self, opc: &[u8], op1: T, op2: U)
187    where
188        Self: EncodeRR<U, T>,
189    {
190        // Flip operands to read as MR.
191        self.encode_rr_mr::<U, T>(opc, op2, op1);
192    }
193
194    /// Encode a register-immediate instruction.
195    pub(crate) fn encode_ri<T: Reg, U: Imm>(&mut self, opc: u8, opc_ext: u8, op1: T, op2: U)
196    where
197        Self: EncodeR<T>,
198    {
199        let modrm = modrm(
200            0b11,      /* mode */
201            opc_ext,   /* reg */
202            op1.idx(), /* rm */
203        );
204
205        let prefix = <Self as EncodeR<T>>::legacy_prefix();
206        let rex = <Self as EncodeR<T>>::rex(op1);
207
208        self.emit_optional(&[prefix, rex]);
209        self.emit(&[opc, modrm]);
210        self.emit(op2.bytes());
211    }
212
213    /// Encode a memory operand instruction.
214    pub(crate) fn encode_m<T: Mem>(&mut self, opc: u8, opc_ext: u8, op1: T)
215    where
216        Self: EncodeM<T>,
217    {
218        // M operand encoding.
219        //   op1 -> modrm.rm
220        let (mode, rm) = match op1.mode() {
221            AddrMode::Indirect => {
222                assert!(!op1.base().need_sib() && !op1.base().is_pc_rel());
223                (0b00, op1.base().idx())
224            }
225            AddrMode::IndirectDisp => {
226                assert!(!op1.base().need_sib());
227                (0b10, op1.base().idx())
228            }
229            AddrMode::IndirectBaseIndex => {
230                assert!(!op1.base().is_pc_rel());
231                // Using rsp as index register is interpreted as just base w/o offset.
232                //   https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing_2
233                // Disallow this case, as guard for the user.
234                assert!(!matches!(op1.index(), Reg64::rsp));
235                (0b00, 0b100)
236            }
237        };
238
239        let modrm = modrm(
240            mode,    /* mode */
241            opc_ext, /* reg */
242            rm,      /* rm */
243        );
244
245        let prefix = <Self as EncodeM<T>>::legacy_prefix();
246        let rex = <Self as EncodeM<T>>::rex(&op1);
247
248        self.emit_optional(&[prefix, rex]);
249        self.emit(&[opc, modrm]);
250        match op1.mode() {
251            AddrMode::Indirect => {}
252            AddrMode::IndirectDisp => self.emit(&op1.disp().to_ne_bytes()),
253            AddrMode::IndirectBaseIndex => {
254                self.emit(&[sib(0, op1.index().idx(), op1.base().idx())])
255            }
256        }
257    }
258
259    /// Encode a memory-immediate instruction.
260    pub(crate) fn encode_mi<M: Mem, T: Imm>(&mut self, opc: u8, opc_ext: u8, op1: M, op2: T)
261    where
262        Self: EncodeM<M>,
263    {
264        // MI operand encoding.
265        //   op1 -> modrm.rm
266        //   op2 -> imm
267        let (mode, rm) = match op1.mode() {
268            AddrMode::Indirect => {
269                assert!(!op1.base().need_sib() && !op1.base().is_pc_rel());
270                (0b00, op1.base().idx())
271            }
272            AddrMode::IndirectDisp => {
273                assert!(!op1.base().need_sib());
274                (0b10, op1.base().idx())
275            }
276            AddrMode::IndirectBaseIndex => {
277                assert!(!op1.base().is_pc_rel());
278                // Using rsp as index register is interpreted as just base w/o offset.
279                //   https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing_2
280                // Disallow this case, as guard for the user.
281                assert!(!matches!(op1.index(), Reg64::rsp));
282                (0b00, 0b100)
283            }
284        };
285
286        let modrm = modrm(
287            mode,    /* mode */
288            opc_ext, /* reg */
289            rm,      /* rm */
290        );
291
292        let prefix = <Self as EncodeM<M>>::legacy_prefix();
293        let rex = <Self as EncodeM<M>>::rex(&op1);
294
295        self.emit_optional(&[prefix, rex]);
296        self.emit(&[opc, modrm]);
297        match op1.mode() {
298            AddrMode::Indirect => {}
299            AddrMode::IndirectDisp => self.emit(&op1.disp().to_ne_bytes()),
300            AddrMode::IndirectBaseIndex => {
301                self.emit(&[sib(0, op1.index().idx(), op1.base().idx())])
302            }
303        }
304        self.emit(op2.bytes());
305    }
306
307    /// Encode a memory-register instruction.
308    pub(crate) fn encode_mr<M: Mem, T: Reg>(&mut self, opc: &[u8], op1: M, op2: T)
309    where
310        Self: EncodeMR<M>,
311    {
312        // MR operand encoding.
313        //   op1 -> modrm.rm
314        //   op2 -> modrm.reg
315        let (mode, rm) = match op1.mode() {
316            AddrMode::Indirect => {
317                assert!(!op1.base().need_sib() && !op1.base().is_pc_rel());
318                (0b00, op1.base().idx())
319            }
320            AddrMode::IndirectDisp => {
321                assert!(!op1.base().need_sib());
322                (0b10, op1.base().idx())
323            }
324            AddrMode::IndirectBaseIndex => {
325                assert!(!op1.base().is_pc_rel());
326                // Using rsp as index register is interpreted as just base w/o offset.
327                //   https://wiki.osdev.org/X86-64_Instruction_Encoding#32.2F64-bit_addressing_2
328                // Disallow this case, as guard for the user.
329                assert!(!matches!(op1.index(), Reg64::rsp));
330                (0b00, 0b100)
331            }
332        };
333
334        let modrm = modrm(
335            mode,      /* mode */
336            op2.idx(), /* reg */
337            rm,        /* rm */
338        );
339
340        let prefix = <Self as EncodeMR<M>>::legacy_prefix();
341        let rex = <Self as EncodeMR<M>>::rex(&op1, op2);
342
343        self.emit_optional(&[prefix, rex]);
344        self.emit(opc);
345        self.emit(&[modrm]);
346        match op1.mode() {
347            AddrMode::Indirect => {}
348            AddrMode::IndirectDisp => self.emit(&op1.disp().to_ne_bytes()),
349            AddrMode::IndirectBaseIndex => {
350                self.emit(&[sib(0, op1.index().idx(), op1.base().idx())])
351            }
352        }
353    }
354
355    /// Encode a register-memory instruction.
356    pub(crate) fn encode_rm<T: Reg, M: Mem>(&mut self, opc: &[u8], op1: T, op2: M)
357    where
358        Self: EncodeMR<M>,
359    {
360        // RM operand encoding.
361        //   op1 -> modrm.reg
362        //   op2 -> modrm.rm
363        self.encode_mr(opc, op2, op1);
364    }
365
366    /// Encode a jump to label instruction.
367    pub(crate) fn encode_jmp_label(&mut self, opc: &[u8], op1: &mut Label) {
368        // Emit the opcode.
369        self.emit(opc);
370
371        // Record relocation offset starting at the first byte of the disp32.
372        op1.record_offset(self.buf.len());
373
374        // Emit a zeroed disp32, which serves as placeholder for the relocation.
375        // We currently only support disp32 jump targets.
376        self.emit(&[0u8; 4]);
377
378        // Resolve any pending relocations for the label.
379        self.resolve(op1);
380    }
381}
382
383// -- Encoder helper.
384
385/// Encode helper for register-register instructions.
386pub(crate) trait EncodeRR<T: Reg, U: Reg> {
387    fn legacy_prefix() -> Option<u8> {
388        None
389    }
390
391    fn rex(op1: T, op2: U) -> Option<u8> {
392        if op1.need_rex() || op2.need_rex() {
393            Some(rex(op1.rexw(), op2.idx(), 0, op1.idx()))
394        } else {
395            None
396        }
397    }
398}
399
400impl EncodeRR<Reg8, Reg8> for Asm {}
401impl EncodeRR<Reg16, Reg16> for Asm {
402    fn legacy_prefix() -> Option<u8> {
403        Some(0x66)
404    }
405}
406impl EncodeRR<Reg32, Reg32> for Asm {}
407impl EncodeRR<Reg64, Reg64> for Asm {}
408
409impl EncodeRR<Reg8, Reg32> for Asm {}
410impl EncodeRR<Reg32, Reg8> for Asm {}
411
412/// Encode helper for register instructions.
413pub(crate) trait EncodeR<T: Reg> {
414    fn legacy_prefix() -> Option<u8> {
415        None
416    }
417
418    fn rex(op1: T) -> Option<u8> {
419        if op1.need_rex() {
420            Some(rex(op1.rexw(), 0, 0, op1.idx()))
421        } else {
422            None
423        }
424    }
425}
426
427impl EncodeR<Reg8> for Asm {}
428impl EncodeR<Reg16> for Asm {
429    fn legacy_prefix() -> Option<u8> {
430        Some(0x66)
431    }
432}
433impl EncodeR<Reg32> for Asm {}
434impl EncodeR<Reg64> for Asm {}
435
436/// Encode helper for memory-register instructions.
437pub(crate) trait EncodeMR<M: Mem> {
438    fn legacy_prefix() -> Option<u8> {
439        None
440    }
441
442    fn rex<T: Reg>(op1: &M, op2: T) -> Option<u8> {
443        if M::is_64() || T::is_64() || op2.is_ext() || op1.base().is_ext() || op1.index().is_ext() {
444            Some(rex(
445                M::is_64() || T::is_64(),
446                op2.idx(),
447                op1.index().idx(),
448                op1.base().idx(),
449            ))
450        } else {
451            None
452        }
453    }
454}
455
456impl EncodeMR<Mem8> for Asm {}
457impl EncodeMR<Mem16> for Asm {
458    fn legacy_prefix() -> Option<u8> {
459        Some(0x66)
460    }
461}
462impl EncodeMR<Mem32> for Asm {}
463impl EncodeMR<Mem64> for Asm {}
464
465/// Encode helper for memory operand instructions.
466pub(crate) trait EncodeM<M: Mem> {
467    fn legacy_prefix() -> Option<u8> {
468        None
469    }
470
471    fn rex(op1: &M) -> Option<u8> {
472        if M::is_64() || op1.base().is_ext() || op1.index().is_ext() {
473            Some(rex(M::is_64(), 0, op1.index().idx(), op1.base().idx()))
474        } else {
475            None
476        }
477    }
478}
479
480impl EncodeM<Mem8> for Asm {}
481impl EncodeM<Mem16> for Asm {
482    fn legacy_prefix() -> Option<u8> {
483        Some(0x66)
484    }
485}
486impl EncodeM<Mem32> for Asm {}
487impl EncodeM<Mem64> for Asm {}
juicebox_asm/asm.rs

juicebox_asm/
asm.rs
