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            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
315        }
316        let checkpoint = self.buffer.checkpoint();
317        let offset = self.buffer.cur_offset();
318        self.b(label);
319        let _ = self
320            .buffer
321            .record_label_patch_site(offset, label, LabelUse::A64Branch26);
322        if self.buffer.error().is_some() {
323            self.buffer.rollback(checkpoint);
324            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
325        }
326        // SAFETY: `b` emits a 26-bit branch at `offset`.
327        unsafe { PatchableSite::new(offset, LabelUse::A64Branch26, 0) }
328    }
329
330    pub fn patchable_bl(&mut self, label: Label) -> PatchableSite {
331        if self.buffer.error().is_some() {
332            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
333        }
334        let checkpoint = self.buffer.checkpoint();
335        let offset = self.buffer.cur_offset();
336        self.bl(label);
337        let _ = self
338            .buffer
339            .record_label_patch_site(offset, label, LabelUse::A64Branch26);
340        if self.buffer.error().is_some() {
341            self.buffer.rollback(checkpoint);
342            return unsafe { PatchableSite::new(u32::MAX, LabelUse::A64Branch26, 0) };
343        }
344        // SAFETY: `bl` emits a 26-bit branch-and-link at `offset`.
345        unsafe { PatchableSite::new(offset, LabelUse::A64Branch26, 0) }
346    }
347
348    /// Patchable immediate materialization via a fixed `movz`/`movk` sequence.
349    ///
350    /// 64-bit destinations use 16 bytes (4 insns); 32-bit destinations use 8 bytes (2 insns).
351    /// Rewrite with [`encode_patchable_mov_imm`] + [`PatchableBlock::rewrite`], or
352    /// [`PatchableBlock::repatch_u64`] is not used here (the block covers whole instructions).
353    pub fn patchable_mov(&mut self, rd: Gp, imm: impl Into<u64>) -> PatchableBlock {
354        let arch = Arch::AArch64;
355        let value = imm.into();
356        if self.buffer.error().is_some() {
357            return unsafe { PatchableBlock::new(u32::MAX, 4, arch) };
358        }
359        let checkpoint = self.buffer.checkpoint();
360        let is_64 = rd.is_gp64();
361        let encoded = encode_patchable_mov_imm(rd.id(), is_64, value);
362        let offset = self.buffer.cur_offset();
363        for chunk in encoded.chunks_exact(4) {
364            let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
365            self.buffer.write_u32(word);
366        }
367        let size = encoded.len() as CodeOffset;
368        let _ = self.buffer.record_patch_block(offset, size, 4);
369        if self.buffer.error().is_some() {
370            self.buffer.rollback(checkpoint);
371            return unsafe { PatchableBlock::new(u32::MAX, size.max(4), arch) };
372        }
373        // SAFETY: fixed movz/movk sequence recorded as a patch block.
374        unsafe { PatchableBlock::new(offset, size, arch) }
375    }
376}
377
378/// Encode a fixed-width patchable mov-immediate sequence for AArch64.
379///
380/// Always emits 2 instructions (8 bytes) for W registers and 4 (16 bytes) for X registers so
381/// rewrites never change layout.
382pub fn encode_patchable_mov_imm(rd: u32, is_64bit: bool, value: u64) -> smallvec::SmallVec<[u8; 16]> {
383    let rd = rd & 0x1f;
384    let mut out = smallvec::SmallVec::new();
385    if is_64bit {
386        const MOVZ: u32 = 0b11010010100000000000000000000000;
387        const MOVK: u32 = 0b11110010100000000000000000000000;
388        let words = [
389            MOVZ | (0 << 21) | (((value as u32) & 0xFFFF) << 5) | rd,
390            MOVK | (1 << 21) | ((((value >> 16) as u32) & 0xFFFF) << 5) | rd,
391            MOVK | (2 << 21) | ((((value >> 32) as u32) & 0xFFFF) << 5) | rd,
392            MOVK | (3 << 21) | ((((value >> 48) as u32) & 0xFFFF) << 5) | rd,
393        ];
394        for w in words {
395            out.extend_from_slice(&w.to_le_bytes());
396        }
397    } else {
398        const MOVZ: u32 = 0b01010010100000000000000000000000;
399        const MOVK: u32 = 0b01110010100000000000000000000000;
400        let value = value as u32;
401        let words = [
402            MOVZ | (0 << 21) | ((value & 0xFFFF) << 5) | rd,
403            MOVK | (1 << 21) | (((value >> 16) & 0xFFFF) << 5) | rd,
404        ];
405        for w in words {
406            out.extend_from_slice(&w.to_le_bytes());
407        }
408    }
409    out
410}
411
412impl InstId {
413    pub const ARM_COND: u32 = 0x78000000;
414    pub const REAL_ID: u32 = 65535;
415    pub const fn with_cc(self, cond: CondCode) -> u32 {
416        let x = self as u32;
417        x | (cond as u32) << Self::ARM_COND.trailing_zeros()
418    }
419
420    pub const fn extract_cc(inst: u32) -> CondCode {
421        unsafe {
422            core::mem::transmute(((inst & Self::ARM_COND) >> Self::ARM_COND.trailing_zeros()) as u8)
423        }
424    }
425
426    pub const fn extract_real_id(inst: u32) -> u32 {
427        inst & Self::REAL_ID
428    }
429}
430
431pub const BLE: u32 = InstId::B.with_cc(CondCode::GE);
432
433impl From<InstId> for u32 {
434    fn from(inst: InstId) -> Self {
435        inst as u32
436    }
437}
438
439// The generated instdb references these through `use super::assembler::*`.
440pub(super) use crate::aarch64::encoder::OffsetType;
441pub use crate::aarch64::encoder::{
442    LogicalImm, count_zero_half_words_64, encode_fp64_to_imm8, encode_logical_imm, is_fp16_imm8,
443    is_fp32_imm8, is_fp64_imm8,
444};
445
446#[cfg(test)]
447mod tests {
448    use super::*;
449    use crate::aarch64::operands::regs;
450    use crate::core::buffer::RelocDistance;
451
452    #[test]
453    fn pic_symbol_load_uses_only_the_got_sequence() {
454        for distance in [RelocDistance::Near, RelocDistance::Far] {
455            let mut environment = Environment::new(Arch::AArch64);
456            environment.set_pic(true);
457            let mut buffer = CodeBuffer::new(environment);
458            let symbol = buffer.extern_sym("external", distance);
459
460            {
461                let mut asm = Assembler::new(&mut buffer);
462                asm.load_constant(regs::x(0), symbol);
463                assert_eq!(
464                    asm.buffer.data(),
465                    &[0x00, 0x00, 0x00, 0x90, 0x00, 0x00, 0x40, 0xF9]
466                );
467                assert_eq!(asm.buffer.relocs().len(), 2);
468                assert_eq!(asm.buffer.relocs()[0].offset, 0);
469                assert_eq!(asm.buffer.relocs()[0].kind, Reloc::Aarch64AdrGotPage21);
470                assert_eq!(asm.buffer.relocs()[0].target, RelocTarget::Sym(symbol));
471                assert_eq!(asm.buffer.relocs()[0].addend, 0);
472                assert_eq!(asm.buffer.relocs()[1].offset, 4);
473                assert_eq!(asm.buffer.relocs()[1].kind, Reloc::Aarch64Ld64GotLo12Nc);
474                assert_eq!(asm.buffer.relocs()[1].target, RelocTarget::Sym(symbol));
475                assert_eq!(asm.buffer.relocs()[1].addend, 0);
476            }
477        }
478    }
479
480    #[test]
481    fn invalid_raw_instruction_ids_are_rejected() {
482        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
483        let mut asm = Assembler::new(&mut buffer);
484
485        asm.emit_n(0u32, &[]);
486        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
487        assert!(asm.buffer.data().is_empty());
488
489        asm.buffer.clear();
490        asm.emit_n(u32::MAX, &[]);
491        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
492        assert!(asm.buffer.data().is_empty());
493
494        asm.buffer.clear();
495        asm.emit_n(InstId::Add as u32 | (1 << 16), &[]);
496        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
497        assert!(asm.buffer.data().is_empty());
498    }
499
500    #[test]
501    fn baseline_rejects_optional_features_before_writing() {
502        let mut buffer = CodeBuffer::new(Environment::baseline(Arch::AArch64));
503        let mut asm = Assembler::new(&mut buffer);
504
505        let error = asm
506            .try_emit_n(
507                InstId::Crc32b,
508                &[
509                    regs::w(0).as_operand(),
510                    regs::w(1).as_operand(),
511                    regs::w(2).as_operand(),
512                ],
513            )
514            .unwrap_err();
515
516        assert!(
517            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("crc32b") && feature.contains("CRC32"))
518        );
519        assert!(asm.buffer.data().is_empty());
520    }
521
522    #[test]
523    fn optional_features_are_enabled_by_default() {
524        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
525        let mut asm = Assembler::new(&mut buffer);
526
527        asm.try_emit_n(
528            InstId::Crc32b,
529            &[
530                regs::w(0).as_operand(),
531                regs::w(1).as_operand(),
532                regs::w(2).as_operand(),
533            ],
534        )
535        .unwrap();
536
537        assert_eq!(asm.buffer.data().len(), 4);
538    }
539
540    #[test]
541    fn mixed_instruction_checks_the_selected_form() {
542        let mut environment = Environment::baseline(Arch::AArch64);
543        environment.set_aarch64_feature(crate::aarch64::CpuFeature::Asimd, true);
544        let mut buffer = CodeBuffer::new(environment);
545        let mut asm = Assembler::new(&mut buffer);
546
547        asm.try_emit_n(
548            InstId::Fadd_v,
549            &[
550                regs::s(0).as_operand(),
551                regs::s(1).as_operand(),
552                regs::s(2).as_operand(),
553            ],
554        )
555        .unwrap();
556        let accepted_len = asm.buffer.data().len();
557
558        let error = asm
559            .try_emit_n(
560                InstId::Fadd_v,
561                &[
562                    regs::h(0).as_operand(),
563                    regs::h(1).as_operand(),
564                    regs::h(2).as_operand(),
565                ],
566            )
567            .unwrap_err();
568
569        assert!(
570            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("fadd Hd") && feature.contains("FP16"))
571        );
572        assert_eq!(asm.buffer.data().len(), accepted_len);
573    }
574
575    #[test]
576    fn raw_emission_rejects_malformed_or_extra_operands_without_mutation() {
577        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
578        let mut asm = Assembler::new(&mut buffer);
579        let malformed = Operand {
580            signature: OperandSignature::from(7),
581            base_id: 0,
582            data: [0; 2],
583        };
584        let invalid_register = Operand {
585            signature: OperandSignature::from(
586                OperandType::Reg as u32 | (31 << OperandSignature::REG_TYPE_SHIFT),
587            ),
588            base_id: 0,
589            data: [0; 2],
590        };
591        let none = Operand::new();
592        let mut invalid_memory = ptr(regs::x(0), 0);
593        invalid_memory.set_base_id(64);
594        let mut invalid_mode = ptr(regs::x(0), 0);
595        invalid_mode
596            .signature
597            .set_field::<{ Mem::SIGNATURE_MEM_OFFSET_MODE_MASK }>(3);
598
599        assert_eq!(
600            asm.try_emit_n(InstId::Add as u32, &[&malformed]),
601            Err(AsmError::InvalidOperand)
602        );
603        assert_eq!(
604            asm.try_emit_n(InstId::Add as u32, &[&invalid_register]),
605            Err(AsmError::InvalidOperand)
606        );
607        assert_eq!(
608            asm.try_emit_n(
609                InstId::Ldr as u32,
610                &[regs::x(1).as_operand(), invalid_memory.as_operand()],
611            ),
612            Err(AsmError::InvalidOperand)
613        );
614        assert_eq!(
615            asm.try_emit_n(
616                InstId::Ldr as u32,
617                &[regs::x(1).as_operand(), invalid_mode.as_operand()],
618            ),
619            Err(AsmError::InvalidOperand)
620        );
621        assert_eq!(
622            asm.try_emit_n(
623                InstId::Add as u32,
624                &[&none, &none, &none, &none, &none, &none, &none],
625            ),
626            Err(AsmError::InvalidOperand)
627        );
628        assert!(asm.buffer.error().is_none());
629        assert!(asm.buffer.data().is_empty());
630    }
631
632    #[test]
633    fn raw_label_registration_error_rolls_back_emission() {
634        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
635        let mut asm = Assembler::new(&mut buffer);
636        let invalid_label = Label::from_id(0);
637
638        assert_eq!(
639            asm.try_emit_n(InstId::B as u32, &[invalid_label.as_operand()]),
640            Err(AsmError::InvalidArgument)
641        );
642        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
643        assert!(asm.buffer.data().is_empty());
644    }
645
646    #[test]
647    fn failed_constant_load_rolls_back_the_whole_sequence() {
648        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
649        let mut asm = Assembler::new(&mut buffer);
650
651        asm.load_constant(regs::x(0), Label::from_id(0));
652
653        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
654        assert!(asm.buffer.data().is_empty());
655    }
656
657    #[test]
658    fn invalid_symbol_load_sets_error_without_mutation() {
659        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
660        let mut asm = Assembler::new(&mut buffer);
661
662        asm.load_constant(regs::x(0), Sym::from_id(u32::MAX));
663
664        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
665        assert!(asm.buffer.data().is_empty());
666    }
667
668    #[test]
669    fn patchable_b_and_mov_handles_work_offline() {
670        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
671        let (site, mov, alt) = {
672            let mut asm = Assembler::new(&mut buffer);
673            let target = asm.get_label();
674            let alt = asm.get_label();
675            let mov = asm.patchable_mov(regs::x(0), 0x1111_2222_3333_4444u64);
676            let site = asm.patchable_b(target);
677            asm.bind_label(target);
678            asm.ret(regs::x(30));
679            asm.bind_label(alt);
680            asm.ret(regs::x(30));
681            let alt_off = asm.buffer.label_offset(alt);
682            (site, mov, alt_off)
683        };
684
685        let code = buffer.finish_patched().unwrap();
686        assert_eq!(mov.size(), 16);
687        let mut bytes = code.data().to_vec();
688        let rewritten = encode_patchable_mov_imm(0, true, 0xAAAA_BBBB_CCCC_DDDD);
689        unsafe {
690            mov.rewrite(&mut bytes, &rewritten).unwrap();
691            site.retarget(&mut bytes, alt).unwrap();
692        }
693        assert_eq!(&bytes[mov.offset() as usize..][..16], rewritten.as_slice());
694    }
695}
asmkit/aarch64/assembler.rs

asmkit/aarch64/
assembler.rs
