asmkit/aarch64/
assembler.rs

1use crate::AsmError;
2use crate::aarch64::operands::*;
3use crate::aarch64::{Gp, Reg, instdb::*};
4use crate::core::arch_traits::Arch;
5use crate::core::buffer::CodeBuffer;
6use crate::core::buffer::{CodeOffset, Constant, LabelUse, Reloc, RelocDistance, RelocTarget};
7use crate::core::globals::CondCode;
8use crate::core::operand::*;
9use crate::core::patch::{PatchableBlock, PatchableSite};
10use crate::core::target::Environment;
11
12pub struct Assembler<'a> {
13    pub(crate) buffer: &'a mut CodeBuffer,
14    /// Scratch error set by the generated emitter during one checked attempt.
15    pub(crate) last_error: Option<AsmError>,
16}
17
18fn validate_raw_operand(buffer: &CodeBuffer, op: &Operand) -> bool {
19    let Some(op_type) = op.signature.try_op_type() else {
20        return false;
21    };
22
23    match op_type {
24        OperandType::None => op.signature.bits() == 0,
25        OperandType::Reg => {
26            let Some(reg_type) = op.signature.try_reg_type() else {
27                return false;
28            };
29            let Some(_) = op.signature.try_reg_group() else {
30                return false;
31            };
32            let expected = Reg::signature_of(reg_type);
33            let mask = OperandSignature::OP_TYPE_MASK
34                | OperandSignature::REG_TYPE_MASK
35                | OperandSignature::REG_GROUP_MASK
36                | OperandSignature::SIZE_MASK;
37            expected.bits() != 0
38                && op.signature.subset(mask) == expected.subset(mask)
39                && match reg_type {
40                    RegType::Gp32 | RegType::Gp64 => op.id() <= 31 || op.id() == Gp::ID_ZR,
41                    RegType::Vec8
42                    | RegType::Vec16
43                    | RegType::Vec32
44                    | RegType::Vec64
45                    | RegType::Vec128 => op.id() <= 31,
46                    _ => false,
47                }
48        }
49        OperandType::Mem => {
50            let Some(base_type) = op.signature.try_mem_base_type() else {
51                return false;
52            };
53            let Some(index_type) = op.signature.try_mem_index_type() else {
54                return false;
55            };
56            let mem = op.as_::<BaseMem>();
57            let offset_mode = op
58                .signature
59                .get_field::<{ Mem::SIGNATURE_MEM_OFFSET_MODE_MASK }>();
60            let shift_op = op
61                .signature
62                .get_field::<{ Mem::SIGNATURE_MEM_SHIFT_OP_MASK }>();
63            let valid_base = match base_type {
64                RegType::None | RegType::LabelTag | RegType::SymTag => true,
65                RegType::Gp32 | RegType::Gp64 => mem.base_id() <= 31,
66                _ => false,
67            };
68            let valid_index = match index_type {
69                RegType::None => mem.index_id() == 0,
70                RegType::Gp32 | RegType::Gp64 => mem.index_id() <= 31,
71                _ => false,
72            };
73            valid_base
74                && valid_index
75                && offset_mode <= OffsetMode::PostIndex as u32
76                && shift_op <= 13
77                && (!mem.has_base_sym()
78                    || buffer.symbol_name(Sym::from_id(mem.base_id())).is_some())
79        }
80        OperandType::Sym => buffer.symbol_name(Sym::from_id(op.id())).is_some(),
81        OperandType::Imm | OperandType::Label => true,
82        OperandType::RegList => false,
83    }
84}
85
86impl crate::core::builder::InstSink for Assembler<'_> {
87    fn arch(&self) -> Arch {
88        self.environment().arch()
89    }
90
91    fn emit_inst(&mut self, inst: &crate::core::inst::Inst) -> Result<(), AsmError> {
92        let ops = inst.operands();
93        let mut refs: smallvec::SmallVec<[&Operand; 6]> = smallvec::SmallVec::new();
94        refs.extend(ops.iter());
95        self.try_emit_n(inst.id(), &refs)
96    }
97
98    fn bind_label(&mut self, label: Label) -> Result<(), AsmError> {
99        self.try_bind_label(label)
100    }
101}
102
103pub trait LoadConstantEmitter<DST, SRC> {
104    fn load_constant(&mut self, dst: DST, src: SRC);
105}
106
107impl LoadConstantEmitter<Gp, Constant> for Assembler<'_> {
108    fn load_constant(&mut self, dst: Gp, src: Constant) {
109        let label = self.buffer.get_label_for_constant(src);
110        self.load_constant(dst, label);
111    }
112}
113
114impl LoadConstantEmitter<Gp, Label> for Assembler<'_> {
115    fn load_constant(&mut self, dst: Gp, src: Label) {
116        self.adrp(dst, src);
117        self.buffer
118            .use_label_at_offset(self.buffer.cur_offset(), src, LabelUse::A64AddAbsLo12);
119        self.add(dst, dst, imm(0));
120    }
121}
122
123impl LoadConstantEmitter<Gp, Sym> for Assembler<'_> {
124    fn load_constant(&mut self, dst: Gp, src: Sym) {
125        let Some(distance) = self.buffer.symbol_distance(src) else {
126            self.buffer.record_error(AsmError::InvalidArgument);
127            return;
128        };
129
130        if self.buffer.env().pic() {
131            // When PIC is enabled, all syms are referenced through the GOT.
132            self.buffer
133                .add_reloc(Reloc::Aarch64AdrGotPage21, RelocTarget::Sym(src), 0);
134            self.adrp(dst, imm(0));
135            self.buffer
136                .add_reloc(Reloc::Aarch64Ld64GotLo12Nc, RelocTarget::Sym(src), 0);
137            self.ldr(dst, ptr(dst, 0));
138            return;
139        }
140
141        match distance {
142            RelocDistance::Near => {
143                self.buffer
144                    .add_reloc(Reloc::Aarch64AdrPrelPgHi21, RelocTarget::Sym(src), 0);
145                self.adrp(dst, imm(0));
146                self.buffer
147                    .add_reloc(Reloc::Aarch64AddAbsLo12Nc, RelocTarget::Sym(src), 0);
148                self.add(dst, dst, imm(0));
149            }
150
151            RelocDistance::Far => {
152                // With absolute offsets we set up a load from a preallocated space, and then jump
153                // over it.
154                //
155                // Emit the following code:
156                //   ldr     rd, #8
157                //   b       #0x10
158                //   <8 byte space>
159                let constant_start = self.buffer.get_label();
160                let constant_end = self.buffer.get_label();
161                self.ldr(dst, label_ptr(constant_start, 0));
162                self.b(constant_end);
163                self.buffer.bind_label(constant_start);
164                self.buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(src), 0);
165                self.buffer.write_u64(0);
166                self.buffer.bind_label(constant_end);
167            }
168        }
169    }
170}
171
172impl<'a> Assembler<'a> {
173    pub fn new(buffer: &'a mut CodeBuffer) -> Self {
174        if buffer.env().arch() != Arch::AArch64 {
175            return Self::poisoned(buffer, AsmError::InvalidArch);
176        }
177        Self::unchecked(buffer)
178    }
179
180    pub fn try_new(buffer: &'a mut CodeBuffer) -> Result<Self, AsmError> {
181        if buffer.env().arch() != Arch::AArch64 {
182            return Err(AsmError::InvalidArch);
183        }
184        Ok(Self::unchecked(buffer))
185    }
186
187    fn unchecked(buffer: &'a mut CodeBuffer) -> Self {
188        Self {
189            buffer,
190            last_error: None,
191        }
192    }
193
194    fn poisoned(buffer: &'a mut CodeBuffer, error: AsmError) -> Self {
195        buffer.record_error(error);
196        Self::unchecked(buffer)
197    }
198
199    /// Returns the environment (arch/mode) this assembler targets.
200    pub fn environment(&self) -> &Environment {
201        self.buffer.env()
202    }
203
204    /// Tests whether the assembler targets a 32-bit mode (always false for A64).
205    pub fn is_32bit(&self) -> bool {
206        self.buffer.env().is_32bit()
207    }
208
209    /// Tests whether the assembler targets a 64-bit mode (always true for A64).
210    pub fn is_64bit(&self) -> bool {
211        self.buffer.env().is_64bit()
212    }
213
214    pub fn get_label(&mut self) -> Label {
215        self.buffer.get_label()
216    }
217
218    pub fn bind_label(&mut self, label: Label) {
219        if let Err(error) = self.try_bind_label(label) {
220            self.buffer.record_error(error);
221        }
222    }
223
224    pub fn try_bind_label(&mut self, label: Label) -> Result<(), AsmError> {
225        self.buffer.try_bind_label(label)
226    }
227
228    /// A helper to load a constant address into a register.
229    ///
230    /// Supported variants are:
231    /// ```text
232    /// +------------------+
233    /// |  DST  |  SRC     |
234    /// +------------------+
235    /// |  Gp   | Label    |
236    /// |  Gp   | Sym      |
237    /// |  Gp   | Constant |
238    /// +------------------+
239    /// ```
240    ///
241    /// Note that `Sym` is loaded based on `self.buffer.pic()` and its distance. If PIC is enabled
242    /// then GOT is always used. Otherwise, if symbol is near it uses `adrp` + `add` combination, and
243    /// for far symbols Abs8 reloc is used and data is embedded right into code.
244    pub fn load_constant<DST, SRC>(&mut self, dst: DST, src: SRC)
245    where
246        Self: LoadConstantEmitter<DST, SRC>,
247    {
248        if self.buffer.error().is_some() {
249            return;
250        }
251        let checkpoint = self.buffer.checkpoint();
252        <Self as LoadConstantEmitter<DST, SRC>>::load_constant(self, dst, src);
253        if self.buffer.error().is_some() {
254            self.buffer.rollback(checkpoint);
255        }
256    }
257
258    #[cfg(test)]
259    fn last_error(&self) -> Option<AsmError> {
260        self.buffer.error().cloned()
261    }
262
263    pub fn emit_n(&mut self, id: impl Into<u32>, ops: &[&Operand]) {
264        if let Err(error) = self.try_emit_n(id, ops) {
265            self.buffer.record_error(error);
266        }
267    }
268
269    pub fn try_emit_n(&mut self, id: impl Into<u32>, ops: &[&Operand]) -> Result<(), AsmError> {
270        if let Some(error) = self.buffer.error().cloned() {
271            return Err(error);
272        }
273        if ops.len() > 6 || ops.iter().any(|op| !validate_raw_operand(self.buffer, op)) {
274            return Err(AsmError::InvalidOperand);
275        }
276        let id = id.into();
277        let checkpoint = self.buffer.checkpoint();
278        self.last_error = None;
279        self._emit(id, ops);
280        if let Some(error) = self.last_error.take() {
281            self.buffer.rollback(checkpoint);
282            return Err(error);
283        }
284        if let Some(error) = self.buffer.error().cloned() {
285            self.buffer.rollback(checkpoint);
286            return Err(error);
287        }
288        Ok(())
289    }
290
291    pub fn data(&self) -> &[u8] {
292        self.buffer.data()
293    }
294
295    pub fn relocs(&self) -> &[crate::core::buffer::AsmReloc] {
296        self.buffer.relocs()
297    }
298
299    pub fn error(&self) -> Option<&AsmError> {
300        self.buffer.error()
301    }
302
303    /// Reserve a nop-filled island for later custom rewriting.
304    pub fn reserve_patch_block(
305        &mut self,
306        size: CodeOffset,
307        align: CodeOffset,
308    ) -> Result<PatchableBlock, AsmError> {
309        self.buffer.reserve_patch_block(size, align)
310    }
311
312    pub fn patchable_b(&mut self, label: Label) -> PatchableSite {
313        if self.buffer.error().is_some() {
314            // SAFETY: poisoned handle for an error path; `u32::MAX` is not a
315            // valid offset into any real image, so applying it fails bounds checks.
316            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
317        }
318        let checkpoint = self.buffer.checkpoint();
319        let offset = self.buffer.cur_offset();
320        self.b(label);
321        let _ = self
322            .buffer
323            .record_label_patch_site(offset, label, LabelUse::A64Branch26);
324        if self.buffer.error().is_some() {
325            self.buffer.rollback(checkpoint);
326            // SAFETY: poisoned handle for an error path; `u32::MAX` is not a
327            // valid offset into any real image, so applying it fails bounds checks.
328            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
329        }
330        // SAFETY: `b` emits a 26-bit branch at `offset`.
331        unsafe { PatchableSite::new(offset, LabelUse::A64Branch26, 0) }
332    }
333
334    pub fn patchable_bl(&mut self, label: Label) -> PatchableSite {
335        if self.buffer.error().is_some() {
336            // SAFETY: poisoned handle for an error path; `u32::MAX` is not a
337            // valid offset into any real image, so applying it fails bounds checks.
338            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
339        }
340        let checkpoint = self.buffer.checkpoint();
341        let offset = self.buffer.cur_offset();
342        self.bl(label);
343        let _ = self
344            .buffer
345            .record_label_patch_site(offset, label, LabelUse::A64Branch26);
346        if self.buffer.error().is_some() {
347            self.buffer.rollback(checkpoint);
348            // SAFETY: poisoned handle for an error path; `u32::MAX` is not a
349            // valid offset into any real image, so applying it fails bounds checks.
350            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
351        }
352        // SAFETY: `bl` emits a 26-bit branch-and-link at `offset`.
353        unsafe { PatchableSite::new(offset, LabelUse::A64Branch26, 0) }
354    }
355
356    /// Patchable immediate materialization via a fixed `movz`/`movk` sequence.
357    ///
358    /// 64-bit destinations use 16 bytes (4 insns); 32-bit destinations use 8 bytes (2 insns).
359    /// Rewrite with [`encode_patchable_mov_imm`] + [`PatchableBlock::rewrite`], or
360    /// [`PatchableBlock::repatch_u64`] is not used here (the block covers whole instructions).
361    pub fn patchable_mov(&mut self, rd: Gp, imm: impl Into<u64>) -> PatchableBlock {
362        let arch = Arch::AArch64;
363        let value = imm.into();
364        if self.buffer.error().is_some() {
365            // SAFETY: poisoned handle for an error path; the sentinel offset is
366            // rejected by the bounds checks when applied.
367            return unsafe { PatchableBlock::new(u32::MAX, 4, arch) };
368        }
369        let checkpoint = self.buffer.checkpoint();
370        let is_64 = rd.is_gp64();
371        let encoded = encode_patchable_mov_imm(rd.id(), is_64, value);
372        let offset = self.buffer.cur_offset();
373        for chunk in encoded.chunks_exact(4) {
374            let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
375            self.buffer.write_u32(word);
376        }
377        let size = encoded.len() as CodeOffset;
378        let _ = self.buffer.record_patch_block(offset, size, 4);
379        if self.buffer.error().is_some() {
380            self.buffer.rollback(checkpoint);
381            // SAFETY: poisoned handle for the rollback path; the sentinel offset
382            // is rejected by the bounds checks when applied.
383            return unsafe { PatchableBlock::new(u32::MAX, size.max(4), arch) };
384        }
385        // SAFETY: fixed movz/movk sequence recorded as a patch block.
386        unsafe { PatchableBlock::new(offset, size, arch) }
387    }
388}
389
390/// Encode a fixed-width patchable mov-immediate sequence for AArch64.
391///
392/// Always emits 2 instructions (8 bytes) for W registers and 4 (16 bytes) for X registers so
393/// rewrites never change layout.
394pub fn encode_patchable_mov_imm(
395    rd: u32,
396    is_64bit: bool,
397    value: u64,
398) -> smallvec::SmallVec<[u8; 16]> {
399    let rd = rd & 0x1f;
400    let mut out = smallvec::SmallVec::new();
401    if is_64bit {
402        const MOVZ: u32 = 0b11010010100000000000000000000000;
403        const MOVK: u32 = 0b11110010100000000000000000000000;
404        let words = [
405            MOVZ | (((value as u32) & 0xFFFF) << 5) | rd,
406            MOVK | (1 << 21) | ((((value >> 16) as u32) & 0xFFFF) << 5) | rd,
407            MOVK | (2 << 21) | ((((value >> 32) as u32) & 0xFFFF) << 5) | rd,
408            MOVK | (3 << 21) | ((((value >> 48) as u32) & 0xFFFF) << 5) | rd,
409        ];
410        for w in words {
411            out.extend_from_slice(&w.to_le_bytes());
412        }
413    } else {
414        const MOVZ: u32 = 0b01010010100000000000000000000000;
415        const MOVK: u32 = 0b01110010100000000000000000000000;
416        let value = value as u32;
417        let words = [
418            MOVZ | ((value & 0xFFFF) << 5) | rd,
419            MOVK | (1 << 21) | (((value >> 16) & 0xFFFF) << 5) | rd,
420        ];
421        for w in words {
422            out.extend_from_slice(&w.to_le_bytes());
423        }
424    }
425    out
426}
427
428impl InstId {
429    pub const ARM_COND: u32 = 0x78000000;
430    pub const REAL_ID: u32 = 65535;
431    pub const fn with_cc(self, cond: CondCode) -> u32 {
432        let x = self as u32;
433        x | (cond as u32) << Self::ARM_COND.trailing_zeros()
434    }
435
436    pub const fn extract_cc(inst: u32) -> CondCode {
437        // SAFETY: `CondCode` is `#[repr(u8)]` with contiguous discriminants
438        // 0..=15, and the 4-bit ARM condition field masked here is always in
439        // that range.
440        unsafe {
441            core::mem::transmute(((inst & Self::ARM_COND) >> Self::ARM_COND.trailing_zeros()) as u8)
442        }
443    }
444
445    pub const fn extract_real_id(inst: u32) -> u32 {
446        inst & Self::REAL_ID
447    }
448}
449
450pub const BLE: u32 = InstId::B.with_cc(CondCode::GE);
451
452impl From<InstId> for u32 {
453    fn from(inst: InstId) -> Self {
454        inst as u32
455    }
456}
457
458// The generated instdb references these through `use super::assembler::*`.
459pub(super) use crate::aarch64::encoder::OffsetType;
460pub use crate::aarch64::encoder::{
461    LogicalImm, count_zero_half_words_64, encode_fp64_to_imm8, encode_logical_imm, is_fp16_imm8,
462    is_fp32_imm8, is_fp64_imm8,
463};
464
465#[cfg(test)]
466mod tests {
467    use super::*;
468    use crate::aarch64::operands::regs;
469    use crate::core::buffer::RelocDistance;
470
471    #[test]
472    fn pic_symbol_load_uses_only_the_got_sequence() {
473        for distance in [RelocDistance::Near, RelocDistance::Far] {
474            let mut environment = Environment::new(Arch::AArch64);
475            environment.set_pic(true);
476            let mut buffer = CodeBuffer::new(environment);
477            let symbol = buffer.extern_sym("external", distance);
478
479            {
480                let mut asm = Assembler::new(&mut buffer);
481                asm.load_constant(regs::x(0), symbol);
482                assert_eq!(
483                    asm.buffer.data(),
484                    &[0x00, 0x00, 0x00, 0x90, 0x00, 0x00, 0x40, 0xF9]
485                );
486                assert_eq!(asm.buffer.relocs().len(), 2);
487                assert_eq!(asm.buffer.relocs()[0].offset, 0);
488                assert_eq!(asm.buffer.relocs()[0].kind, Reloc::Aarch64AdrGotPage21);
489                assert_eq!(asm.buffer.relocs()[0].target, RelocTarget::Sym(symbol));
490                assert_eq!(asm.buffer.relocs()[0].addend, 0);
491                assert_eq!(asm.buffer.relocs()[1].offset, 4);
492                assert_eq!(asm.buffer.relocs()[1].kind, Reloc::Aarch64Ld64GotLo12Nc);
493                assert_eq!(asm.buffer.relocs()[1].target, RelocTarget::Sym(symbol));
494                assert_eq!(asm.buffer.relocs()[1].addend, 0);
495            }
496        }
497    }
498
499    #[test]
500    fn invalid_raw_instruction_ids_are_rejected() {
501        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
502        let mut asm = Assembler::new(&mut buffer);
503
504        asm.emit_n(0u32, &[]);
505        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
506        assert!(asm.buffer.data().is_empty());
507
508        asm.buffer.clear();
509        asm.emit_n(u32::MAX, &[]);
510        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
511        assert!(asm.buffer.data().is_empty());
512
513        asm.buffer.clear();
514        asm.emit_n(InstId::Add as u32 | (1 << 16), &[]);
515        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
516        assert!(asm.buffer.data().is_empty());
517    }
518
519    #[test]
520    fn baseline_rejects_optional_features_before_writing() {
521        let mut buffer = CodeBuffer::new(Environment::baseline(Arch::AArch64));
522        let mut asm = Assembler::new(&mut buffer);
523
524        let error = asm
525            .try_emit_n(
526                InstId::Crc32b,
527                &[
528                    regs::w(0).as_operand(),
529                    regs::w(1).as_operand(),
530                    regs::w(2).as_operand(),
531                ],
532            )
533            .unwrap_err();
534
535        assert!(
536            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("crc32b") && feature.contains("CRC32"))
537        );
538        assert!(asm.buffer.data().is_empty());
539    }
540
541    #[test]
542    fn optional_features_are_enabled_by_default() {
543        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
544        let mut asm = Assembler::new(&mut buffer);
545
546        asm.try_emit_n(
547            InstId::Crc32b,
548            &[
549                regs::w(0).as_operand(),
550                regs::w(1).as_operand(),
551                regs::w(2).as_operand(),
552            ],
553        )
554        .unwrap();
555
556        assert_eq!(asm.buffer.data().len(), 4);
557    }
558
559    #[test]
560    fn mixed_instruction_checks_the_selected_form() {
561        let mut environment = Environment::baseline(Arch::AArch64);
562        environment.set_aarch64_feature(crate::aarch64::CpuFeature::Asimd, true);
563        let mut buffer = CodeBuffer::new(environment);
564        let mut asm = Assembler::new(&mut buffer);
565
566        asm.try_emit_n(
567            InstId::Fadd_v,
568            &[
569                regs::s(0).as_operand(),
570                regs::s(1).as_operand(),
571                regs::s(2).as_operand(),
572            ],
573        )
574        .unwrap();
575        let accepted_len = asm.buffer.data().len();
576
577        let error = asm
578            .try_emit_n(
579                InstId::Fadd_v,
580                &[
581                    regs::h(0).as_operand(),
582                    regs::h(1).as_operand(),
583                    regs::h(2).as_operand(),
584                ],
585            )
586            .unwrap_err();
587
588        assert!(
589            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("fadd Hd") && feature.contains("FP16"))
590        );
591        assert_eq!(asm.buffer.data().len(), accepted_len);
592    }
593
594    #[test]
595    fn raw_emission_rejects_malformed_or_extra_operands_without_mutation() {
596        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
597        let mut asm = Assembler::new(&mut buffer);
598        let malformed = Operand {
599            signature: OperandSignature::from(7),
600            base_id: 0,
601            data: [0; 2],
602        };
603        let invalid_register = Operand {
604            signature: OperandSignature::from(
605                OperandType::Reg as u32 | (31 << OperandSignature::REG_TYPE_SHIFT),
606            ),
607            base_id: 0,
608            data: [0; 2],
609        };
610        let none = Operand::new();
611        let mut invalid_memory = ptr(regs::x(0), 0);
612        invalid_memory.set_base_id(64);
613        let mut invalid_mode = ptr(regs::x(0), 0);
614        invalid_mode
615            .signature
616            .set_field::<{ Mem::SIGNATURE_MEM_OFFSET_MODE_MASK }>(3);
617
618        assert_eq!(
619            asm.try_emit_n(InstId::Add as u32, &[&malformed]),
620            Err(AsmError::InvalidOperand)
621        );
622        assert_eq!(
623            asm.try_emit_n(InstId::Add as u32, &[&invalid_register]),
624            Err(AsmError::InvalidOperand)
625        );
626        assert_eq!(
627            asm.try_emit_n(
628                InstId::Ldr as u32,
629                &[regs::x(1).as_operand(), invalid_memory.as_operand()],
630            ),
631            Err(AsmError::InvalidOperand)
632        );
633        assert_eq!(
634            asm.try_emit_n(
635                InstId::Ldr as u32,
636                &[regs::x(1).as_operand(), invalid_mode.as_operand()],
637            ),
638            Err(AsmError::InvalidOperand)
639        );
640        assert_eq!(
641            asm.try_emit_n(
642                InstId::Add as u32,
643                &[&none, &none, &none, &none, &none, &none, &none],
644            ),
645            Err(AsmError::InvalidOperand)
646        );
647        assert!(asm.buffer.error().is_none());
648        assert!(asm.buffer.data().is_empty());
649    }
650
651    #[test]
652    fn raw_label_registration_error_rolls_back_emission() {
653        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
654        let mut asm = Assembler::new(&mut buffer);
655        let invalid_label = Label::from_id(0);
656
657        assert_eq!(
658            asm.try_emit_n(InstId::B as u32, &[invalid_label.as_operand()]),
659            Err(AsmError::InvalidArgument)
660        );
661        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
662        assert!(asm.buffer.data().is_empty());
663    }
664
665    #[test]
666    fn failed_constant_load_rolls_back_the_whole_sequence() {
667        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
668        let mut asm = Assembler::new(&mut buffer);
669
670        asm.load_constant(regs::x(0), Label::from_id(0));
671
672        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
673        assert!(asm.buffer.data().is_empty());
674    }
675
676    #[test]
677    fn invalid_symbol_load_sets_error_without_mutation() {
678        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
679        let mut asm = Assembler::new(&mut buffer);
680
681        asm.load_constant(regs::x(0), Sym::from_id(u32::MAX));
682
683        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
684        assert!(asm.buffer.data().is_empty());
685    }
686
687    #[test]
688    fn patchable_b_and_mov_handles_work_offline() {
689        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
690        let (site, mov, alt) = {
691            let mut asm = Assembler::new(&mut buffer);
692            let target = asm.get_label();
693            let alt = asm.get_label();
694            let mov = asm.patchable_mov(regs::x(0), 0x1111_2222_3333_4444u64);
695            let site = asm.patchable_b(target);
696            asm.bind_label(target);
697            asm.ret(regs::x(30));
698            asm.bind_label(alt);
699            asm.ret(regs::x(30));
700            let alt_off = asm.buffer.label_offset(alt);
701            (site, mov, alt_off)
702        };
703
704        let code = buffer.finish_patched().unwrap();
705        assert_eq!(mov.size(), 16);
706        let mut bytes = code.data().to_vec();
707        let rewritten = encode_patchable_mov_imm(0, true, 0xAAAA_BBBB_CCCC_DDDD);
708        unsafe {
709            mov.rewrite(&mut bytes, &rewritten).unwrap();
710            site.retarget(&mut bytes, alt).unwrap();
711        }
712        assert_eq!(&bytes[mov.offset() as usize..][..16], rewritten.as_slice());
713    }
714}
asmkit/aarch64/assembler.rs

asmkit/aarch64/
assembler.rs
