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::{Constant, LabelUse, Reloc, RelocDistance, RelocTarget};
7use crate::core::globals::CondCode;
8use crate::core::operand::*;
9use crate::core::target::Environment;
10
11pub struct Assembler<'a> {
12    pub(crate) buffer: &'a mut CodeBuffer,
13    /// Scratch error set by the generated emitter during one checked attempt.
14    pub(crate) last_error: Option<AsmError>,
15}
16
17fn validate_raw_operand(buffer: &CodeBuffer, op: &Operand) -> bool {
18    let Some(op_type) = op.signature.try_op_type() else {
19        return false;
20    };
21
22    match op_type {
23        OperandType::None => op.signature.bits() == 0,
24        OperandType::Reg => {
25            let Some(reg_type) = op.signature.try_reg_type() else {
26                return false;
27            };
28            let Some(_) = op.signature.try_reg_group() else {
29                return false;
30            };
31            let expected = Reg::signature_of(reg_type);
32            let mask = OperandSignature::OP_TYPE_MASK
33                | OperandSignature::REG_TYPE_MASK
34                | OperandSignature::REG_GROUP_MASK
35                | OperandSignature::SIZE_MASK;
36            expected.bits() != 0
37                && op.signature.subset(mask) == expected.subset(mask)
38                && match reg_type {
39                    RegType::Gp32 | RegType::Gp64 => op.id() <= 31 || op.id() == Gp::ID_ZR,
40                    RegType::Vec8
41                    | RegType::Vec16
42                    | RegType::Vec32
43                    | RegType::Vec64
44                    | RegType::Vec128 => op.id() <= 31,
45                    _ => false,
46                }
47        }
48        OperandType::Mem => {
49            let Some(base_type) = op.signature.try_mem_base_type() else {
50                return false;
51            };
52            let Some(index_type) = op.signature.try_mem_index_type() else {
53                return false;
54            };
55            let mem = op.as_::<BaseMem>();
56            let offset_mode = op
57                .signature
58                .get_field::<{ Mem::SIGNATURE_MEM_OFFSET_MODE_MASK }>();
59            let shift_op = op
60                .signature
61                .get_field::<{ Mem::SIGNATURE_MEM_SHIFT_OP_MASK }>();
62            let valid_base = match base_type {
63                RegType::None | RegType::LabelTag | RegType::SymTag => true,
64                RegType::Gp32 | RegType::Gp64 => mem.base_id() <= 31,
65                _ => false,
66            };
67            let valid_index = match index_type {
68                RegType::None => mem.index_id() == 0,
69                RegType::Gp32 | RegType::Gp64 => mem.index_id() <= 31,
70                _ => false,
71            };
72            valid_base
73                && valid_index
74                && offset_mode <= OffsetMode::PostIndex as u32
75                && shift_op <= 13
76                && (!mem.has_base_sym()
77                    || buffer.symbol_name(Sym::from_id(mem.base_id())).is_some())
78        }
79        OperandType::Sym => buffer.symbol_name(Sym::from_id(op.id())).is_some(),
80        OperandType::Imm | OperandType::Label => true,
81        OperandType::RegList => false,
82    }
83}
84
85impl crate::core::builder::InstSink for Assembler<'_> {
86    fn arch(&self) -> Arch {
87        self.environment().arch()
88    }
89
90    fn emit_inst(&mut self, inst: &crate::core::inst::Inst) -> Result<(), AsmError> {
91        let ops = inst.operands();
92        let mut refs: smallvec::SmallVec<[&Operand; 6]> = smallvec::SmallVec::new();
93        refs.extend(ops.iter());
94        self.try_emit_n(inst.id(), &refs)
95    }
96
97    fn bind_label(&mut self, label: Label) -> Result<(), AsmError> {
98        self.try_bind_label(label)
99    }
100}
101
102pub trait LoadConstantEmitter<DST, SRC> {
103    fn load_constant(&mut self, dst: DST, src: SRC);
104}
105
106impl LoadConstantEmitter<Gp, Constant> for Assembler<'_> {
107    fn load_constant(&mut self, dst: Gp, src: Constant) {
108        let label = self.buffer.get_label_for_constant(src);
109        self.load_constant(dst, label);
110    }
111}
112
113impl LoadConstantEmitter<Gp, Label> for Assembler<'_> {
114    fn load_constant(&mut self, dst: Gp, src: Label) {
115        self.adrp(dst, src);
116        self.buffer
117            .use_label_at_offset(self.buffer.cur_offset(), src, LabelUse::A64AddAbsLo12);
118        self.add(dst, dst, imm(0));
119    }
120}
121
122impl LoadConstantEmitter<Gp, Sym> for Assembler<'_> {
123    fn load_constant(&mut self, dst: Gp, src: Sym) {
124        let Some(distance) = self.buffer.symbol_distance(src) else {
125            self.buffer.record_error(AsmError::InvalidArgument);
126            return;
127        };
128
129        if self.buffer.env().pic() {
130            // When PIC is enabled, all syms are referenced through the GOT.
131            self.buffer
132                .add_reloc(Reloc::Aarch64AdrGotPage21, RelocTarget::Sym(src), 0);
133            self.adrp(dst, imm(0));
134            self.buffer
135                .add_reloc(Reloc::Aarch64Ld64GotLo12Nc, RelocTarget::Sym(src), 0);
136            self.ldr(dst, ptr(dst, 0));
137            return;
138        }
139
140        match distance {
141            RelocDistance::Near => {
142                self.buffer
143                    .add_reloc(Reloc::Aarch64AdrPrelPgHi21, RelocTarget::Sym(src), 0);
144                self.adrp(dst, imm(0));
145                self.buffer
146                    .add_reloc(Reloc::Aarch64AddAbsLo12Nc, RelocTarget::Sym(src), 0);
147                self.add(dst, dst, imm(0));
148            }
149
150            RelocDistance::Far => {
151                // With absolute offsets we set up a load from a preallocated space, and then jump
152                // over it.
153                //
154                // Emit the following code:
155                //   ldr     rd, #8
156                //   b       #0x10
157                //   <8 byte space>
158                let constant_start = self.buffer.get_label();
159                let constant_end = self.buffer.get_label();
160                self.ldr(dst, label_ptr(constant_start, 0));
161                self.b(constant_end);
162                self.buffer.bind_label(constant_start);
163                self.buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(src), 0);
164                self.buffer.write_u64(0);
165                self.buffer.bind_label(constant_end);
166            }
167        }
168    }
169}
170
171impl<'a> Assembler<'a> {
172    pub fn new(buffer: &'a mut CodeBuffer) -> Self {
173        if buffer.env().arch() != Arch::AArch64 {
174            return Self::poisoned(buffer, AsmError::InvalidArch);
175        }
176        Self::unchecked(buffer)
177    }
178
179    pub fn try_new(buffer: &'a mut CodeBuffer) -> Result<Self, AsmError> {
180        if buffer.env().arch() != Arch::AArch64 {
181            return Err(AsmError::InvalidArch);
182        }
183        Ok(Self::unchecked(buffer))
184    }
185
186    fn unchecked(buffer: &'a mut CodeBuffer) -> Self {
187        Self {
188            buffer,
189            last_error: None,
190        }
191    }
192
193    fn poisoned(buffer: &'a mut CodeBuffer, error: AsmError) -> Self {
194        buffer.record_error(error);
195        Self::unchecked(buffer)
196    }
197
198    /// Returns the environment (arch/mode) this assembler targets.
199    pub fn environment(&self) -> &Environment {
200        self.buffer.env()
201    }
202
203    /// Tests whether the assembler targets a 32-bit mode (always false for A64).
204    pub fn is_32bit(&self) -> bool {
205        self.buffer.env().is_32bit()
206    }
207
208    /// Tests whether the assembler targets a 64-bit mode (always true for A64).
209    pub fn is_64bit(&self) -> bool {
210        self.buffer.env().is_64bit()
211    }
212
213    pub fn get_label(&mut self) -> Label {
214        self.buffer.get_label()
215    }
216
217    pub fn bind_label(&mut self, label: Label) {
218        if let Err(error) = self.try_bind_label(label) {
219            self.buffer.record_error(error);
220        }
221    }
222
223    pub fn try_bind_label(&mut self, label: Label) -> Result<(), AsmError> {
224        self.buffer.try_bind_label(label)
225    }
226
227    /// A helper to load a constant address into a register.
228    ///
229    /// Supported variants are:
230    /// ```text
231    /// +------------------+
232    /// |  DST  |  SRC     |
233    /// +------------------+
234    /// |  Gp   | Label    |
235    /// |  Gp   | Sym      |
236    /// |  Gp   | Constant |
237    /// +------------------+
238    /// ```
239    ///
240    /// Note that `Sym` is loaded based on `self.buffer.pic()` and its distance. If PIC is enabled
241    /// then GOT is always used. Otherwise, if symbol is near it uses `adrp` + `add` combination, and
242    /// for far symbols Abs8 reloc is used and data is embedded right into code.
243    pub fn load_constant<DST, SRC>(&mut self, dst: DST, src: SRC)
244    where
245        Self: LoadConstantEmitter<DST, SRC>,
246    {
247        if self.buffer.error().is_some() {
248            return;
249        }
250        let checkpoint = self.buffer.checkpoint();
251        <Self as LoadConstantEmitter<DST, SRC>>::load_constant(self, dst, src);
252        if self.buffer.error().is_some() {
253            self.buffer.rollback(checkpoint);
254        }
255    }
256
257    #[cfg(test)]
258    fn last_error(&self) -> Option<AsmError> {
259        self.buffer.error().cloned()
260    }
261
262    pub fn emit_n(&mut self, id: impl Into<u32>, ops: &[&Operand]) {
263        if let Err(error) = self.try_emit_n(id, ops) {
264            self.buffer.record_error(error);
265        }
266    }
267
268    pub fn try_emit_n(&mut self, id: impl Into<u32>, ops: &[&Operand]) -> Result<(), AsmError> {
269        if let Some(error) = self.buffer.error().cloned() {
270            return Err(error);
271        }
272        if ops.len() > 6 || ops.iter().any(|op| !validate_raw_operand(self.buffer, op)) {
273            return Err(AsmError::InvalidOperand);
274        }
275        let id = id.into();
276        let checkpoint = self.buffer.checkpoint();
277        self.last_error = None;
278        self._emit(id, ops);
279        if let Some(error) = self.last_error.take() {
280            self.buffer.rollback(checkpoint);
281            return Err(error);
282        }
283        if let Some(error) = self.buffer.error().cloned() {
284            self.buffer.rollback(checkpoint);
285            return Err(error);
286        }
287        Ok(())
288    }
289
290    pub fn data(&self) -> &[u8] {
291        self.buffer.data()
292    }
293
294    pub fn relocs(&self) -> &[crate::core::buffer::AsmReloc] {
295        self.buffer.relocs()
296    }
297
298    pub fn error(&self) -> Option<&AsmError> {
299        self.buffer.error()
300    }
301}
302
303impl InstId {
304    pub const ARM_COND: u32 = 0x78000000;
305    pub const REAL_ID: u32 = 65535;
306    pub const fn with_cc(self, cond: CondCode) -> u32 {
307        let x = self as u32;
308        x | (cond as u32) << Self::ARM_COND.trailing_zeros()
309    }
310
311    pub const fn extract_cc(inst: u32) -> CondCode {
312        unsafe {
313            core::mem::transmute(((inst & Self::ARM_COND) >> Self::ARM_COND.trailing_zeros()) as u8)
314        }
315    }
316
317    pub const fn extract_real_id(inst: u32) -> u32 {
318        inst & Self::REAL_ID
319    }
320}
321
322pub const BLE: u32 = InstId::B.with_cc(CondCode::GE);
323
324impl From<InstId> for u32 {
325    fn from(inst: InstId) -> Self {
326        inst as u32
327    }
328}
329
330// The generated instdb references these through `use super::assembler::*`.
331pub(super) use crate::aarch64::encoder::OffsetType;
332pub use crate::aarch64::encoder::{
333    LogicalImm, count_zero_half_words_64, encode_fp64_to_imm8, encode_logical_imm, is_fp16_imm8,
334    is_fp32_imm8, is_fp64_imm8,
335};
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340    use crate::aarch64::operands::regs;
341    use crate::core::buffer::RelocDistance;
342
343    #[test]
344    fn pic_symbol_load_uses_only_the_got_sequence() {
345        for distance in [RelocDistance::Near, RelocDistance::Far] {
346            let mut environment = Environment::new(Arch::AArch64);
347            environment.set_pic(true);
348            let mut buffer = CodeBuffer::new(environment);
349            let symbol = buffer.extern_sym("external", distance);
350
351            {
352                let mut asm = Assembler::new(&mut buffer);
353                asm.load_constant(regs::x(0), symbol);
354                assert_eq!(
355                    asm.buffer.data(),
356                    &[0x00, 0x00, 0x00, 0x90, 0x00, 0x00, 0x40, 0xF9]
357                );
358                assert_eq!(asm.buffer.relocs().len(), 2);
359                assert_eq!(asm.buffer.relocs()[0].offset, 0);
360                assert_eq!(asm.buffer.relocs()[0].kind, Reloc::Aarch64AdrGotPage21);
361                assert_eq!(asm.buffer.relocs()[0].target, RelocTarget::Sym(symbol));
362                assert_eq!(asm.buffer.relocs()[0].addend, 0);
363                assert_eq!(asm.buffer.relocs()[1].offset, 4);
364                assert_eq!(asm.buffer.relocs()[1].kind, Reloc::Aarch64Ld64GotLo12Nc);
365                assert_eq!(asm.buffer.relocs()[1].target, RelocTarget::Sym(symbol));
366                assert_eq!(asm.buffer.relocs()[1].addend, 0);
367            }
368        }
369    }
370
371    #[test]
372    fn invalid_raw_instruction_ids_are_rejected() {
373        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
374        let mut asm = Assembler::new(&mut buffer);
375
376        asm.emit_n(0u32, &[]);
377        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
378        assert!(asm.buffer.data().is_empty());
379
380        asm.buffer.clear();
381        asm.emit_n(u32::MAX, &[]);
382        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
383        assert!(asm.buffer.data().is_empty());
384
385        asm.buffer.clear();
386        asm.emit_n(InstId::Add as u32 | (1 << 16), &[]);
387        assert_eq!(asm.last_error(), Some(AsmError::InvalidInstruction));
388        assert!(asm.buffer.data().is_empty());
389    }
390
391    #[test]
392    fn baseline_rejects_optional_features_before_writing() {
393        let mut buffer = CodeBuffer::new(Environment::baseline(Arch::AArch64));
394        let mut asm = Assembler::new(&mut buffer);
395
396        let error = asm
397            .try_emit_n(
398                InstId::Crc32b,
399                &[
400                    regs::w(0).as_operand(),
401                    regs::w(1).as_operand(),
402                    regs::w(2).as_operand(),
403                ],
404            )
405            .unwrap_err();
406
407        assert!(
408            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("crc32b") && feature.contains("CRC32"))
409        );
410        assert!(asm.buffer.data().is_empty());
411    }
412
413    #[test]
414    fn optional_features_are_enabled_by_default() {
415        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
416        let mut asm = Assembler::new(&mut buffer);
417
418        asm.try_emit_n(
419            InstId::Crc32b,
420            &[
421                regs::w(0).as_operand(),
422                regs::w(1).as_operand(),
423                regs::w(2).as_operand(),
424            ],
425        )
426        .unwrap();
427
428        assert_eq!(asm.buffer.data().len(), 4);
429    }
430
431    #[test]
432    fn mixed_instruction_checks_the_selected_form() {
433        let mut environment = Environment::baseline(Arch::AArch64);
434        environment.set_aarch64_feature(crate::aarch64::CpuFeature::Asimd, true);
435        let mut buffer = CodeBuffer::new(environment);
436        let mut asm = Assembler::new(&mut buffer);
437
438        asm.try_emit_n(
439            InstId::Fadd_v,
440            &[
441                regs::s(0).as_operand(),
442                regs::s(1).as_operand(),
443                regs::s(2).as_operand(),
444            ],
445        )
446        .unwrap();
447        let accepted_len = asm.buffer.data().len();
448
449        let error = asm
450            .try_emit_n(
451                InstId::Fadd_v,
452                &[
453                    regs::h(0).as_operand(),
454                    regs::h(1).as_operand(),
455                    regs::h(2).as_operand(),
456                ],
457            )
458            .unwrap_err();
459
460        assert!(
461            matches!(error, AsmError::MissingCpuFeature { feature } if feature.contains("fadd Hd") && feature.contains("FP16"))
462        );
463        assert_eq!(asm.buffer.data().len(), accepted_len);
464    }
465
466    #[test]
467    fn raw_emission_rejects_malformed_or_extra_operands_without_mutation() {
468        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
469        let mut asm = Assembler::new(&mut buffer);
470        let malformed = Operand {
471            signature: OperandSignature::from(7),
472            base_id: 0,
473            data: [0; 2],
474        };
475        let invalid_register = Operand {
476            signature: OperandSignature::from(
477                OperandType::Reg as u32 | (31 << OperandSignature::REG_TYPE_SHIFT),
478            ),
479            base_id: 0,
480            data: [0; 2],
481        };
482        let none = Operand::new();
483        let mut invalid_memory = ptr(regs::x(0), 0);
484        invalid_memory.set_base_id(64);
485        let mut invalid_mode = ptr(regs::x(0), 0);
486        invalid_mode
487            .signature
488            .set_field::<{ Mem::SIGNATURE_MEM_OFFSET_MODE_MASK }>(3);
489
490        assert_eq!(
491            asm.try_emit_n(InstId::Add as u32, &[&malformed]),
492            Err(AsmError::InvalidOperand)
493        );
494        assert_eq!(
495            asm.try_emit_n(InstId::Add as u32, &[&invalid_register]),
496            Err(AsmError::InvalidOperand)
497        );
498        assert_eq!(
499            asm.try_emit_n(
500                InstId::Ldr as u32,
501                &[regs::x(1).as_operand(), invalid_memory.as_operand()],
502            ),
503            Err(AsmError::InvalidOperand)
504        );
505        assert_eq!(
506            asm.try_emit_n(
507                InstId::Ldr as u32,
508                &[regs::x(1).as_operand(), invalid_mode.as_operand()],
509            ),
510            Err(AsmError::InvalidOperand)
511        );
512        assert_eq!(
513            asm.try_emit_n(
514                InstId::Add as u32,
515                &[&none, &none, &none, &none, &none, &none, &none],
516            ),
517            Err(AsmError::InvalidOperand)
518        );
519        assert!(asm.buffer.error().is_none());
520        assert!(asm.buffer.data().is_empty());
521    }
522
523    #[test]
524    fn raw_label_registration_error_rolls_back_emission() {
525        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
526        let mut asm = Assembler::new(&mut buffer);
527        let invalid_label = Label::from_id(0);
528
529        assert_eq!(
530            asm.try_emit_n(InstId::B as u32, &[invalid_label.as_operand()]),
531            Err(AsmError::InvalidArgument)
532        );
533        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
534        assert!(asm.buffer.data().is_empty());
535    }
536
537    #[test]
538    fn failed_constant_load_rolls_back_the_whole_sequence() {
539        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
540        let mut asm = Assembler::new(&mut buffer);
541
542        asm.load_constant(regs::x(0), Label::from_id(0));
543
544        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
545        assert!(asm.buffer.data().is_empty());
546    }
547
548    #[test]
549    fn invalid_symbol_load_sets_error_without_mutation() {
550        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
551        let mut asm = Assembler::new(&mut buffer);
552
553        asm.load_constant(regs::x(0), Sym::from_id(u32::MAX));
554
555        assert_eq!(asm.buffer.error(), Some(&AsmError::InvalidArgument));
556        assert!(asm.buffer.data().is_empty());
557    }
558}
asmkit/aarch64/assembler.rs

asmkit/aarch64/
assembler.rs
