Skip to main content

asmkit/core/
buffer.rs

1use alloc::{borrow::Cow, collections::BinaryHeap, vec::Vec};
2use core::fmt;
3
4use smallvec::SmallVec;
5
6use crate::AsmError;
7use crate::core::arch_traits::Arch;
8use crate::core::patch::{
9    PatchBlock, PatchBlockId, PatchCatalog, PatchSite, PatchSiteId, PatchableBlock,
10    fill_with_nops, minimum_patch_alignment,
11};
12#[cfg(feature = "riscv")]
13use crate::riscv;
14
15#[cfg(feature = "jit")]
16use crate::core::jit_allocator::{JitAllocator, Span};
17
18use super::{
19    operand::{Label, Sym},
20    target::Environment,
21};
22
23/// A buffer of output to be produced, fixed up, and then emitted to a CodeSink
24/// in bulk.
25///
26/// This struct uses `SmallVec`s to support small-ish function bodies without
27/// any heap allocation. As such, it will be several kilobytes large. This is
28/// likely fine as long as it is stack-allocated for function emission then
29/// thrown away; but beware if many buffer objects are retained persistently.
30pub struct CodeBuffer {
31    env: Environment,
32    data: SmallVec<[u8; 1024]>,
33    relocs: SmallVec<[AsmReloc; 16]>,
34    symbols: SmallVec<[SymData; 16]>,
35    defined_symbols: SmallVec<[(ExternalName, Label); 4]>,
36    label_offsets: SmallVec<[CodeOffset; 16]>,
37    pending_fixup_records: SmallVec<[AsmFixup; 16]>,
38    pending_fixup_deadline: u32,
39    pending_constants: SmallVec<[Constant; 16]>,
40    pending_constants_size: CodeOffset,
41    used_constants: SmallVec<[(Constant, CodeOffset); 4]>,
42    constants: SmallVec<[(ConstantData, AsmConstant); 4]>,
43    fixup_records: BinaryHeap<AsmFixup>,
44    patch_blocks: SmallVec<[PendingPatchBlock; 4]>,
45    patch_sites: SmallVec<[PendingPatchSite; 8]>,
46    #[cfg(feature = "x86")]
47    x86_branch_relaxations: SmallVec<[X86BranchRelaxation; 8]>,
48    #[cfg(feature = "x86")]
49    alignment_constraints: SmallVec<[(CodeOffset, CodeOffset); 4]>,
50    error: Option<AsmError>,
51}
52
53/// Private rollback point for one raw instruction emission.
54///
55/// Raw backend emission may append bytes and metadata, and x86 may rewrite
56/// bytes appended by the same attempt. It must not mutate metadata or bytes
57/// that predate this checkpoint; island emission is a finalization operation.
58#[derive(Clone, Copy)]
59#[cfg(any(feature = "x86", feature = "aarch64", feature = "riscv"))]
60pub(crate) struct EmissionCheckpoint {
61    data_len: usize,
62    relocs_len: usize,
63    symbols_len: usize,
64    defined_symbols_len: usize,
65    label_offsets_len: usize,
66    pending_fixup_records_len: usize,
67    pending_fixup_deadline: u32,
68    pending_constants_len: usize,
69    pending_constants_size: CodeOffset,
70    used_constants_len: usize,
71    constants_len: usize,
72    patch_blocks_len: usize,
73    patch_sites_len: usize,
74    #[cfg(feature = "x86")]
75    x86_branch_relaxations_len: usize,
76    #[cfg(feature = "x86")]
77    alignment_constraints_len: usize,
78}
79
80#[cfg(feature = "x86")]
81#[derive(Clone, Copy)]
82struct X86BranchRelaxation {
83    opcode_offset: CodeOffset,
84    label: Label,
85    short_opcode: u8,
86    near_size: u8,
87}
88
89#[derive(Clone, Copy)]
90struct PendingPatchBlock {
91    offset: CodeOffset,
92    size: CodeOffset,
93    align: CodeOffset,
94}
95
96#[derive(Clone, Copy)]
97enum PendingPatchTarget {
98    Offset(CodeOffset),
99    Label(Label),
100}
101
102#[derive(Clone, Copy)]
103struct PendingPatchSite {
104    offset: CodeOffset,
105    kind: LabelUse,
106    target: PendingPatchTarget,
107    addend: i64,
108}
109
110/// An external name in a user-defined symbol table.
111///
112/// Cranelift-style opaque key: asmkit does not interpret `namespace` or `index`.
113/// Hosts commonly use separate namespaces for functions vs data, but that
114/// convention is not enforced here.
115#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
116pub struct UserExternalName {
117    pub namespace: u32,
118    pub index: u32,
119}
120
121impl UserExternalName {
122    /// Creates a new user external name.
123    pub const fn new(namespace: u32, index: u32) -> Self {
124        Self { namespace, index }
125    }
126}
127
128impl fmt::Display for UserExternalName {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        write!(f, "u{}:{}", self.namespace, self.index)
131    }
132}
133
134/// Name of an external (or exported) symbol.
135///
136/// Either a string [`Symbol`](ExternalName::Symbol) for human-readable /
137/// ELF-like names, or a [`User`](ExternalName::User) namespace+index key for
138/// backends that do not use string symbols.
139#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
140pub enum ExternalName {
141    Symbol(Cow<'static, str>),
142    User(UserExternalName),
143}
144
145impl ExternalName {
146    /// Creates a user-defined external name from `namespace` and `index`.
147    pub const fn user(namespace: u32, index: u32) -> Self {
148        Self::User(UserExternalName::new(namespace, index))
149    }
150}
151
152impl fmt::Display for ExternalName {
153    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
154        match self {
155            Self::Symbol(name) => f.write_str(name),
156            Self::User(name) => fmt::Display::fmt(name, f),
157        }
158    }
159}
160
161impl From<UserExternalName> for ExternalName {
162    fn from(name: UserExternalName) -> Self {
163        Self::User(name)
164    }
165}
166
167impl From<&'static str> for ExternalName {
168    fn from(name: &'static str) -> Self {
169        Self::Symbol(Cow::Borrowed(name))
170    }
171}
172
173impl From<alloc::string::String> for ExternalName {
174    fn from(name: alloc::string::String) -> Self {
175        Self::Symbol(Cow::Owned(name))
176    }
177}
178
179impl From<Cow<'static, str>> for ExternalName {
180    fn from(name: Cow<'static, str>) -> Self {
181        Self::Symbol(name)
182    }
183}
184#[derive(Clone, PartialEq, Eq, Hash, Debug)]
185pub enum RelocTarget {
186    Sym(Sym),
187    Label(Label),
188}
189
190#[derive(Copy, Clone, PartialEq, Eq, Debug)]
191pub enum RelocDistance {
192    Near,
193    Far,
194}
195
196#[derive(Clone, PartialEq, Eq)]
197pub(crate) struct SymData {
198    pub(crate) name: ExternalName,
199    pub(crate) distance: RelocDistance,
200}
201
202/// A relocation resulting from emitting assembly.
203#[derive(Clone, PartialEq, Eq, Hash, Debug)]
204pub struct AsmReloc {
205    pub offset: CodeOffset,
206    pub kind: Reloc,
207    pub addend: i64,
208    pub target: RelocTarget,
209}
210/// A fixup to perform on the buffer once code is emitted.
211/// Fixups always refer to labels and patch the code based on label offsets.
212/// Hence, they are like relocations, but internal to one buffer.
213#[derive(Clone, Copy, PartialEq, PartialOrd, Ord, Eq)]
214pub(crate) struct AsmFixup {
215    pub label: Label,
216    pub offset: CodeOffset,
217    pub kind: LabelUse,
218}
219
220impl AsmFixup {
221    fn deadline(&self) -> CodeOffset {
222        self.offset.saturating_sub(self.kind.max_pos_range())
223    }
224}
225
226/// Metadata about a constant.
227#[derive(Clone, Copy)]
228struct AsmConstant {
229    /// A label which has not yet been bound which can be used for this
230    /// constant.
231    ///
232    /// This is lazily created when a label is requested for a constant and is
233    /// cleared when a constant is emitted.
234    upcoming_label: Option<Label>,
235    /// Required alignment.
236    align: CodeOffset,
237    /// The byte size of this constant.
238    size: usize,
239}
240
241/// A `CodeBuffer` once emission is completed: holds generated code and records,
242/// without fixups. This allows the type to be independent of the backend.
243pub struct CodeBufferFinalized {
244    pub(crate) data: SmallVec<[u8; 1024]>,
245    pub(crate) relocs: SmallVec<[AsmReloc; 16]>,
246    pub(crate) symbols: SmallVec<[SymData; 16]>,
247    pub(crate) label_offsets: SmallVec<[CodeOffset; 16]>,
248    pub(crate) defined_symbols: SmallVec<[(ExternalName, CodeOffset); 4]>,
249    pub(crate) alignment: u32,
250    pub(crate) patch_catalog: PatchCatalog,
251}
252
253/// Executable memory loaded from a finalized code buffer with relocations applied.
254#[cfg(feature = "jit")]
255pub struct LoadedRelocatedCode {
256    span: Span,
257    code_size: usize,
258    got_targets: Vec<RelocTarget>,
259}
260
261#[cfg(feature = "jit")]
262impl LoadedRelocatedCode {
263    pub const fn rx(&self) -> *const u8 {
264        self.span.rx()
265    }
266
267    pub const fn rw(&self) -> *mut u8 {
268        self.span.rw()
269    }
270
271    pub const fn span(&self) -> &Span {
272        &self.span
273    }
274
275    pub const fn code_size(&self) -> usize {
276        self.code_size
277    }
278
279    pub fn got_targets(&self) -> &[RelocTarget] {
280        &self.got_targets
281    }
282
283    pub fn got_size(&self) -> usize {
284        self.got_targets.len() * core::mem::size_of::<usize>()
285    }
286
287    pub fn got_rx(&self) -> *const u8 {
288        self.rx().wrapping_add(self.code_size)
289    }
290
291    pub fn got_rw(&self) -> *mut u8 {
292        self.rw().wrapping_add(self.code_size)
293    }
294}
295
296pub fn reloc_uses_got(kind: Reloc) -> bool {
297    matches!(
298        kind,
299        Reloc::X86GOTPCRel4
300            | Reloc::RiscvGotHi20
301            | Reloc::RiscvPCRelLo12I
302            | Reloc::Aarch64AdrGotPage21
303            | Reloc::Aarch64Ld64GotLo12Nc
304    )
305}
306
307#[cfg(feature = "jit")]
308pub(crate) fn got_slot_index(got_targets: &[RelocTarget], target: &RelocTarget) -> Option<usize> {
309    got_targets.iter().position(|item| item == target)
310}
311
312impl CodeBufferFinalized {
313    pub fn total_size(&self) -> usize {
314        self.data.len()
315    }
316
317    pub fn data(&self) -> &[u8] {
318        &self.data[..]
319    }
320
321    pub fn symbol_name(&self, sym: Sym) -> Option<&ExternalName> {
322        self.symbols
323            .get(sym.id() as usize)
324            .map(|symbol| &symbol.name)
325    }
326
327    pub fn symbol_distance(&self, sym: Sym) -> Option<RelocDistance> {
328        self.symbols
329            .get(sym.id() as usize)
330            .map(|symbol| symbol.distance)
331    }
332
333    /// Offset of a symbol exported with [`CodeBuffer::bind_symbol`], if present.
334    ///
335    /// After loading, the runtime address of the symbol is `rx + offset`.
336    pub fn defined_symbol_offset(&self, name: &ExternalName) -> Option<CodeOffset> {
337        self.defined_symbols
338            .iter()
339            .find(|(defined, _)| defined == name)
340            .map(|(_, offset)| *offset)
341    }
342
343    /// Like [`Self::defined_symbol_offset`], for string [`ExternalName::Symbol`] exports.
344    pub fn defined_symbol_str(&self, name: &str) -> Option<CodeOffset> {
345        self.defined_symbols
346            .iter()
347            .find(|(defined, _)| matches!(defined, ExternalName::Symbol(s) if s.as_ref() == name))
348            .map(|(_, offset)| *offset)
349    }
350
351    pub fn relocs(&self) -> &[AsmReloc] {
352        &self.relocs[..]
353    }
354
355    pub fn alignment(&self) -> u32 {
356        self.alignment
357    }
358
359    /// Allocate this code buffer in executable memory and return a `Span` referring to it.
360    /// This will also write the code into the allocated memory. To execute
361    /// code you can simply use [`span.rx()`](Span::rx) to get a pointer to read+exec memory
362    /// and transmute that to a function pointer of the appropriate type.
363    #[cfg(feature = "jit")]
364    pub fn allocate(&self, jit_allocator: &mut JitAllocator) -> Result<Span, AsmError> {
365        let mut span = jit_allocator.alloc(self.data().len())?;
366
367        unsafe {
368            jit_allocator.write(&mut span, |span| {
369                span.rw()
370                    .copy_from_nonoverlapping(self.data().as_ptr(), self.data().len());
371            })?;
372        }
373
374        Ok(span)
375    }
376
377    /// Allocate executable memory and apply relocations, including GOT setup in JIT mode.
378    ///
379    /// GOT entries are created automatically for relocations that require them and populated
380    /// with values returned by `get_address`.
381    ///
382    /// Relocations targeting a [`Label`](RelocTarget::Label) are resolved internally against
383    /// the label offsets recorded at finalize time; the callbacks only see external symbols.
384    #[cfg(feature = "jit")]
385    pub fn allocate_relocated(
386        &self,
387        jit_allocator: &mut JitAllocator,
388        get_address: impl Fn(&RelocTarget) -> *const u8,
389        get_plt_entry: impl Fn(&RelocTarget) -> *const u8,
390    ) -> Result<LoadedRelocatedCode, AsmError> {
391        let mut got_targets = Vec::new();
392
393        for reloc in &self.relocs {
394            if reloc_uses_got(reloc.kind) && !got_targets.iter().any(|item| item == &reloc.target) {
395                got_targets.push(reloc.target.clone());
396            }
397        }
398
399        let got_size = got_targets
400            .len()
401            .checked_mul(core::mem::size_of::<usize>())
402            .ok_or(AsmError::TooLarge)?;
403        let total_size = self
404            .data()
405            .len()
406            .checked_add(got_size)
407            .ok_or(AsmError::TooLarge)?;
408        let mut span = jit_allocator.alloc(total_size)?;
409
410        let mut relocation_result = Ok(());
411        unsafe {
412            jit_allocator.write(&mut span, |span| {
413                relocation_result = (|| {
414                    span.rw()
415                        .copy_from_nonoverlapping(self.data().as_ptr(), self.data().len());
416
417                    let rx = span.rx();
418                    let resolve = |target: &RelocTarget,
419                                   fallback: &dyn Fn(&RelocTarget) -> *const u8|
420                     -> Result<*const u8, AsmError> {
421                        if let RelocTarget::Label(label) = target {
422                            let offset = self
423                                .label_offsets
424                                .get(label.id() as usize)
425                                .copied()
426                                .ok_or(AsmError::InvalidArgument)?;
427                            if offset == u32::MAX || offset as usize > self.data().len() {
428                                return Err(AsmError::UnboundLabel);
429                            }
430                            return Ok(rx.add(offset as usize));
431                        }
432                        let address = fallback(target);
433                        if address.is_null() {
434                            return Err(AsmError::InvalidArgument);
435                        }
436                        Ok(address)
437                    };
438
439                    let got_rw = span.rw().add(self.data().len()) as *mut usize;
440                    for (index, target) in got_targets.iter().enumerate() {
441                        let addr = resolve(target, &get_address)?;
442                        got_rw.add(index).write_unaligned(addr.addr());
443                    }
444
445                    let got_rx = rx.add(self.data().len());
446                    perform_relocations(
447                        span.rw(),
448                        rx,
449                        self.data().len(),
450                        &self.relocs,
451                        |target| resolve(target, &get_address),
452                        |target| {
453                            let index = got_slot_index(&got_targets, target)
454                                .ok_or(AsmError::InvalidState)?;
455                            let offset = index
456                                .checked_mul(core::mem::size_of::<usize>())
457                                .ok_or(AsmError::TooLarge)?;
458                            Ok(got_rx.add(offset))
459                        },
460                        |target| resolve(target, &get_plt_entry),
461                    )
462                })();
463            })?;
464        }
465        relocation_result?;
466
467        Ok(LoadedRelocatedCode {
468            span,
469            code_size: self.data().len(),
470            got_targets,
471        })
472    }
473
474    /// Allocate executable memory and apply relocations, resolving external
475    /// symbols through `resolve`.
476    ///
477    /// This is the ergonomic counterpart of [`Self::allocate_relocated`]: every
478    /// undefined external symbol (declared with [`CodeBuffer::extern_sym`],
479    /// [`CodeBuffer::extern_user`], or surviving a
480    /// [`Linker`](crate::core::linker::Linker) link) is passed to `resolve`,
481    /// which must return its address. Symbols defined inside the image (label
482    /// targets and linked-in definitions) are resolved internally.
483    #[cfg(feature = "jit")]
484    pub fn allocate_resolved(
485        &self,
486        jit_allocator: &mut JitAllocator,
487        resolve: impl Fn(&ExternalName) -> *const u8,
488    ) -> Result<LoadedRelocatedCode, AsmError> {
489        let by_name = |target: &RelocTarget| match target {
490            RelocTarget::Sym(sym) => match self.symbol_name(*sym) {
491                Some(name) => resolve(name),
492                None => core::ptr::null(),
493            },
494            // Label targets are resolved internally by `allocate_relocated`.
495            RelocTarget::Label(_) => core::ptr::null(),
496        };
497
498        self.allocate_relocated(jit_allocator, by_name, by_name)
499    }
500}
501
502impl CodeBuffer {
503    /// Creates a buffer for `env`.
504    pub fn new(env: Environment) -> Self {
505        Self {
506            env,
507            data: SmallVec::new(),
508            relocs: SmallVec::new(),
509            symbols: SmallVec::new(),
510            defined_symbols: SmallVec::new(),
511            label_offsets: SmallVec::new(),
512            pending_fixup_records: SmallVec::new(),
513            pending_fixup_deadline: 0,
514            pending_constants: SmallVec::new(),
515            pending_constants_size: 0,
516            used_constants: SmallVec::new(),
517            constants: SmallVec::new(),
518            fixup_records: BinaryHeap::new(),
519            patch_blocks: SmallVec::new(),
520            patch_sites: SmallVec::new(),
521            #[cfg(feature = "x86")]
522            x86_branch_relaxations: SmallVec::new(),
523            #[cfg(feature = "x86")]
524            alignment_constraints: SmallVec::new(),
525            error: None,
526        }
527    }
528
529    /// Creates a buffer for the host target.
530    pub fn host() -> Self {
531        Self::new(Environment::host())
532    }
533
534    pub fn clear(&mut self) {
535        self.data.clear();
536        self.relocs.clear();
537        self.label_offsets.clear();
538        self.pending_fixup_records.clear();
539        self.constants.clear();
540        self.fixup_records.clear();
541        self.symbols.clear();
542        self.defined_symbols.clear();
543        self.used_constants.clear();
544        self.pending_fixup_deadline = 0;
545        self.pending_constants_size = 0;
546        self.pending_constants.clear();
547        self.patch_blocks.clear();
548        self.patch_sites.clear();
549        #[cfg(feature = "x86")]
550        self.x86_branch_relaxations.clear();
551        #[cfg(feature = "x86")]
552        self.alignment_constraints.clear();
553        self.error = None;
554    }
555
556    /// Returns the first error recorded by a void assembler operation.
557    pub fn error(&self) -> Option<&AsmError> {
558        self.error.as_ref()
559    }
560
561    pub(crate) fn record_error(&mut self, error: AsmError) {
562        if self.error.is_none() {
563            self.error = Some(error);
564        }
565    }
566
567    #[cfg(any(feature = "x86", feature = "aarch64", feature = "riscv"))]
568    pub(crate) fn checkpoint(&self) -> EmissionCheckpoint {
569        EmissionCheckpoint {
570            data_len: self.data.len(),
571            relocs_len: self.relocs.len(),
572            symbols_len: self.symbols.len(),
573            defined_symbols_len: self.defined_symbols.len(),
574            label_offsets_len: self.label_offsets.len(),
575            pending_fixup_records_len: self.pending_fixup_records.len(),
576            pending_fixup_deadline: self.pending_fixup_deadline,
577            pending_constants_len: self.pending_constants.len(),
578            pending_constants_size: self.pending_constants_size,
579            used_constants_len: self.used_constants.len(),
580            constants_len: self.constants.len(),
581            patch_blocks_len: self.patch_blocks.len(),
582            patch_sites_len: self.patch_sites.len(),
583            #[cfg(feature = "x86")]
584            x86_branch_relaxations_len: self.x86_branch_relaxations.len(),
585            #[cfg(feature = "x86")]
586            alignment_constraints_len: self.alignment_constraints.len(),
587        }
588    }
589
590    #[cfg(any(feature = "x86", feature = "aarch64", feature = "riscv"))]
591    pub(crate) fn rollback(&mut self, checkpoint: EmissionCheckpoint) {
592        self.data.truncate(checkpoint.data_len);
593        self.relocs.truncate(checkpoint.relocs_len);
594        self.symbols.truncate(checkpoint.symbols_len);
595        self.defined_symbols
596            .truncate(checkpoint.defined_symbols_len);
597        self.label_offsets.truncate(checkpoint.label_offsets_len);
598        self.pending_fixup_records
599            .truncate(checkpoint.pending_fixup_records_len);
600        self.pending_fixup_deadline = checkpoint.pending_fixup_deadline;
601        self.pending_constants
602            .truncate(checkpoint.pending_constants_len);
603        self.pending_constants_size = checkpoint.pending_constants_size;
604        self.used_constants.truncate(checkpoint.used_constants_len);
605        self.constants.truncate(checkpoint.constants_len);
606        self.patch_blocks.truncate(checkpoint.patch_blocks_len);
607        self.patch_sites.truncate(checkpoint.patch_sites_len);
608        #[cfg(feature = "x86")]
609        self.x86_branch_relaxations
610            .truncate(checkpoint.x86_branch_relaxations_len);
611        #[cfg(feature = "x86")]
612        self.alignment_constraints
613            .truncate(checkpoint.alignment_constraints_len);
614    }
615    pub fn env(&self) -> &Environment {
616        &self.env
617    }
618
619    pub fn data(&self) -> &[u8] {
620        &self.data
621    }
622
623    /// Returns the byte at `offset`.
624    ///
625    /// Used by encoder patch-up paths (e.g. X86 LEA/abs32 fixups).
626    #[cfg(feature = "x86")]
627    pub(crate) fn byte_at(&self, offset: CodeOffset) -> u8 {
628        self.data[offset as usize]
629    }
630
631    /// Overwrites the byte at `offset`.
632    #[cfg(feature = "x86")]
633    pub(crate) fn set_byte_at(&mut self, offset: CodeOffset, value: u8) {
634        self.data[offset as usize] = value;
635    }
636
637    /// Inserts a byte at `offset`, shifting all subsequent bytes.
638    ///
639    /// This is a rare encoder patch-up operation; prefer appending.
640    #[cfg(feature = "x86")]
641    pub(crate) fn insert_at(&mut self, offset: CodeOffset, value: u8) {
642        self.data.insert(offset as usize, value);
643    }
644
645    /// Removes the byte at `offset`, shifting all subsequent bytes.
646    ///
647    /// This is a rare encoder patch-up operation; prefer appending.
648    #[cfg(feature = "x86")]
649    pub(crate) fn remove_at(&mut self, offset: CodeOffset) {
650        self.data.remove(offset as usize);
651    }
652
653    #[cfg(feature = "x86")]
654    pub(crate) fn record_x86_branch_relaxation(
655        &mut self,
656        opcode_offset: CodeOffset,
657        label: Label,
658        short_opcode: u8,
659        near_size: u8,
660    ) {
661        debug_assert!(matches!(near_size, 5 | 6));
662        debug_assert!((opcode_offset as usize) + near_size as usize <= self.data.len());
663        debug_assert!(self.label_offsets.get(label.id() as usize).is_some());
664        self.x86_branch_relaxations.push(X86BranchRelaxation {
665            opcode_offset,
666            label,
667            short_opcode,
668            near_size,
669        });
670    }
671
672    pub fn relocs(&self) -> &[AsmReloc] {
673        &self.relocs
674    }
675
676    pub fn put1(&mut self, value: u8) {
677        if self.error.is_some() {
678            return;
679        }
680        self.data.push(value);
681    }
682
683    pub fn put2(&mut self, value: u16) {
684        if self.error.is_some() {
685            return;
686        }
687        self.data.extend_from_slice(&value.to_ne_bytes());
688    }
689
690    pub fn put4(&mut self, value: u32) {
691        if self.error.is_some() {
692            return;
693        }
694        self.data.extend_from_slice(&value.to_ne_bytes());
695    }
696
697    pub fn put8(&mut self, value: u64) {
698        if self.error.is_some() {
699            return;
700        }
701        self.data.extend_from_slice(&value.to_ne_bytes());
702    }
703
704    pub fn write_u8(&mut self, value: u8) {
705        if self.error.is_some() {
706            return;
707        }
708        self.data.push(value);
709    }
710
711    pub fn write_u16(&mut self, value: u16) {
712        if self.error.is_some() {
713            return;
714        }
715        self.data.extend_from_slice(&value.to_ne_bytes());
716    }
717
718    pub fn write_u32(&mut self, value: u32) {
719        if self.error.is_some() {
720            return;
721        }
722        self.data.extend_from_slice(&value.to_ne_bytes());
723    }
724
725    pub fn write_u64(&mut self, value: u64) {
726        if self.error.is_some() {
727            return;
728        }
729        self.data.extend_from_slice(&value.to_ne_bytes());
730    }
731
732    pub fn add_symbol(&mut self, name: impl Into<ExternalName>, distance: RelocDistance) -> Sym {
733        if self.error.is_some() {
734            return Sym::new();
735        }
736        let ix = self.symbols.len();
737        self.symbols.push(SymData {
738            distance,
739            name: name.into(),
740        });
741
742        Sym::from_id(ix as u32)
743    }
744
745    /// Declares an external symbol by name, deduplicating: repeated calls with
746    /// the same name return the same `Sym` (the distance of the first
747    /// declaration wins).
748    ///
749    /// The returned symbol can be passed to the architecture assemblers (e.g.
750    /// `ptr64_sym` on x86) which then record the appropriate relocation based
751    /// on the symbol's distance.
752    pub fn extern_sym(
753        &mut self,
754        name: impl Into<Cow<'static, str>>,
755        distance: RelocDistance,
756    ) -> Sym {
757        if self.error.is_some() {
758            return Sym::new();
759        }
760        let name = ExternalName::Symbol(name.into());
761        if let Some(ix) = self.symbols.iter().position(|sym| sym.name == name) {
762            return Sym::from_id(ix as u32);
763        }
764
765        self.add_symbol(name, distance)
766    }
767
768    /// Declares an external symbol by user namespace+index, deduplicating:
769    /// repeated calls with the same key return the same `Sym` (the distance of
770    /// the first declaration wins).
771    pub fn extern_user(&mut self, namespace: u32, index: u32, distance: RelocDistance) -> Sym {
772        if self.error.is_some() {
773            return Sym::new();
774        }
775        let name = ExternalName::user(namespace, index);
776        if let Some(ix) = self.symbols.iter().position(|sym| sym.name == name) {
777            return Sym::from_id(ix as u32);
778        }
779
780        self.add_symbol(name, distance)
781    }
782
783    /// Exports `label` under `name`, making it a defined symbol that other
784    /// modules can resolve at link time (see [`crate::core::linker::Linker`]).
785    ///
786    /// `name` may be a string ([`ExternalName::Symbol`]) or a user key
787    /// ([`ExternalName::User`]).
788    pub fn bind_symbol(&mut self, name: impl Into<ExternalName>, label: Label) {
789        if self.error.is_some() {
790            return;
791        }
792        if self.label_offsets.get(label.id() as usize).is_none() {
793            self.record_error(AsmError::InvalidArgument);
794            return;
795        }
796        self.defined_symbols.push((name.into(), label));
797    }
798
799    pub fn symbol_distance(&self, sym: Sym) -> Option<RelocDistance> {
800        self.symbols
801            .get(sym.id() as usize)
802            .map(|symbol| symbol.distance)
803    }
804
805    pub fn symbol_name(&self, sym: Sym) -> Option<&ExternalName> {
806        self.symbols
807            .get(sym.id() as usize)
808            .map(|symbol| &symbol.name)
809    }
810
811    pub fn get_label(&mut self) -> Label {
812        if self.error.is_some() {
813            return Label::new();
814        }
815        let l = self.label_offsets.len();
816        self.label_offsets.push(u32::MAX);
817        Label::from_id(l as _)
818    }
819
820    pub fn is_bound(&self, label: Label) -> bool {
821        self.label_offsets
822            .get(label.id() as usize)
823            .is_some_and(|offset| *offset != u32::MAX)
824    }
825
826    /// Number of labels created so far; label ids below this count are valid.
827    pub fn label_count(&self) -> u32 {
828        self.label_offsets.len() as u32
829    }
830
831    pub fn get_label_for_constant(&mut self, constant: Constant) -> Label {
832        if self.error.is_some() {
833            return Label::new();
834        }
835        let Some((_, metadata)) = self.constants.get(constant.0 as usize) else {
836            self.record_error(AsmError::InvalidArgument);
837            return Label::new();
838        };
839        let metadata = *metadata;
840        let AsmConstant {
841            upcoming_label,
842            size,
843            ..
844        } = metadata;
845        if let Some(label) = upcoming_label {
846            return label;
847        }
848
849        let label = self.get_label();
850        self.pending_constants.push(constant);
851        self.pending_constants_size += size as u32;
852        self.constants[constant.0 as usize].1.upcoming_label = Some(label);
853        label
854    }
855
856    pub fn add_constant(&mut self, constant: impl Into<ConstantData>) -> Constant {
857        if self.error.is_some() {
858            return Constant(u32::MAX);
859        }
860        let c = self.constants.len() as u32;
861        let data = constant.into();
862        let x = AsmConstant {
863            upcoming_label: None,
864            align: data.alignment() as _,
865            size: data.as_slice().len(),
866        };
867        self.constants.push((data, x));
868        Constant(c)
869    }
870
871    pub fn use_label_at_offset(&mut self, offset: CodeOffset, label: Label, kind: LabelUse) {
872        if self.error.is_some() {
873            return;
874        }
875        if let Err(error) = kind.validate_for_arch(self.env.arch()) {
876            self.record_error(error);
877            return;
878        }
879        if (offset as usize).checked_add(kind.patch_size()).is_none() {
880            self.record_error(AsmError::TooLarge);
881            return;
882        }
883        if self.label_offsets.get(label.id() as usize).is_none() {
884            self.record_error(AsmError::InvalidArgument);
885            return;
886        }
887        let fixup = AsmFixup {
888            kind,
889            label,
890            offset,
891        };
892
893        self.pending_fixup_records.push(fixup);
894        self.pending_fixup_deadline = self.pending_fixup_deadline.min(fixup.deadline());
895    }
896
897    /// Align up to the given alignment.
898    pub fn try_align_to(&mut self, align_to: CodeOffset) -> Result<(), AsmError> {
899        if let Some(error) = self.error.clone() {
900            return Err(error);
901        }
902        if !align_to.is_power_of_two() {
903            return Err(AsmError::InvalidArgument);
904        }
905        while self.cur_offset() & (align_to - 1) != 0 {
906            self.write_u8(0);
907        }
908        #[cfg(feature = "x86")]
909        if align_to > 1 {
910            self.alignment_constraints
911                .push((self.cur_offset(), align_to));
912        }
913        Ok(())
914    }
915
916    pub fn align_to(&mut self, align_to: CodeOffset) {
917        if let Err(error) = self.try_align_to(align_to) {
918            self.record_error(error);
919        }
920    }
921
922    pub fn cur_offset(&self) -> CodeOffset {
923        self.data.len() as _
924    }
925
926    pub fn try_bind_label(&mut self, label: Label) -> Result<(), AsmError> {
927        if let Some(error) = self.error.clone() {
928            return Err(error);
929        }
930        let current_offset = self.cur_offset();
931        let Some(offset) = self.label_offsets.get_mut(label.id() as usize) else {
932            return Err(AsmError::InvalidArgument);
933        };
934        if *offset != u32::MAX {
935            return Err(AsmError::InvalidState);
936        }
937        *offset = current_offset;
938        Ok(())
939    }
940
941    pub fn bind_label(&mut self, label: Label) {
942        if self.error.is_some() {
943            return;
944        }
945        if let Err(error) = self.try_bind_label(label) {
946            self.record_error(error);
947        }
948    }
949
950    pub fn label_offset(&self, label: Label) -> u32 {
951        self.label_offsets
952            .get(label.id() as usize)
953            .copied()
954            .unwrap_or(u32::MAX)
955    }
956
957    pub fn add_reloc(&mut self, kind: Reloc, target: RelocTarget, addend: i64) {
958        let offset = self.cur_offset();
959        self.add_reloc_at_offset(offset, kind, target, addend);
960    }
961
962    pub fn add_reloc_at_offset(
963        &mut self,
964        offset: CodeOffset,
965        kind: Reloc,
966        target: RelocTarget,
967        addend: i64,
968    ) {
969        if self.error.is_some() {
970            return;
971        }
972        if !kind.supports_arch(self.env.arch()) {
973            self.record_error(AsmError::InvalidArch);
974            return;
975        }
976        let valid_target = match target {
977            RelocTarget::Sym(sym) => self.symbols.get(sym.id() as usize).is_some(),
978            RelocTarget::Label(label) => self.label_offsets.get(label.id() as usize).is_some(),
979        };
980        if !valid_target {
981            self.record_error(AsmError::InvalidArgument);
982            return;
983        }
984        self.relocs.push(AsmReloc {
985            addend,
986            kind,
987            offset,
988            target,
989        })
990    }
991
992    /// Reserve a nop-filled island for later custom rewriting (JSC `padBeforePatch`).
993    ///
994    /// Returns a [`PatchableBlock`] handle; the block is also recorded in the patch catalog
995    /// for `finish_patched` / linker rebasing.
996    pub fn reserve_patch_block(
997        &mut self,
998        size: CodeOffset,
999        align: CodeOffset,
1000    ) -> Result<PatchableBlock, AsmError> {
1001        if let Some(error) = self.error.clone() {
1002            return Err(error);
1003        }
1004        let min_align = minimum_patch_alignment(self.env.arch());
1005        let align = align.max(min_align);
1006        if size == 0 || !align.is_power_of_two() {
1007            return Err(AsmError::InvalidArgument);
1008        }
1009        let nop_size = match self.env.arch() {
1010            Arch::X86 | Arch::X64 => 1,
1011            Arch::AArch64 | Arch::RISCV32 | Arch::RISCV64 => 4,
1012            _ => return Err(AsmError::InvalidArch),
1013        };
1014        if size as usize % nop_size != 0 {
1015            return Err(AsmError::InvalidArgument);
1016        }
1017
1018        self.try_align_to(align)?;
1019        let arch = self.env.arch();
1020        let offset = self.cur_offset();
1021        let block = self.get_appended_space(size as usize);
1022        fill_with_nops(arch, block)?;
1023
1024        self.patch_blocks.push(PendingPatchBlock {
1025            offset,
1026            size,
1027            align,
1028        });
1029        // SAFETY: we just reserved and recorded this nop island.
1030        Ok(unsafe { PatchableBlock::new(offset, size, arch) })
1031    }
1032
1033    pub fn record_patch_block(
1034        &mut self,
1035        offset: CodeOffset,
1036        size: CodeOffset,
1037        align: CodeOffset,
1038    ) -> PatchBlockId {
1039        match self.try_record_patch_block(offset, size, align) {
1040            Ok(id) => id,
1041            Err(error) => {
1042                self.record_error(error);
1043                PatchBlockId::from_index(usize::MAX)
1044            }
1045        }
1046    }
1047
1048    pub fn try_record_patch_block(
1049        &mut self,
1050        offset: CodeOffset,
1051        size: CodeOffset,
1052        align: CodeOffset,
1053    ) -> Result<PatchBlockId, AsmError> {
1054        if let Some(error) = self.error.clone() {
1055            return Err(error);
1056        }
1057        if size == 0 || align == 0 || !align.is_power_of_two() || offset & (align - 1) != 0 {
1058            return Err(AsmError::InvalidArgument);
1059        }
1060        let end = (offset as usize)
1061            .checked_add(size as usize)
1062            .ok_or(AsmError::TooLarge)?;
1063        if end > self.data.len() {
1064            return Err(AsmError::InvalidArgument);
1065        }
1066        let id = PatchBlockId::from_index(self.patch_blocks.len());
1067        self.patch_blocks.push(PendingPatchBlock {
1068            offset,
1069            size,
1070            align,
1071        });
1072        Ok(id)
1073    }
1074
1075    pub fn record_patch_site(
1076        &mut self,
1077        offset: CodeOffset,
1078        kind: LabelUse,
1079        target_offset: CodeOffset,
1080    ) -> PatchSiteId {
1081        match self.try_record_patch_site(offset, kind, target_offset) {
1082            Ok(id) => id,
1083            Err(error) => {
1084                self.record_error(error);
1085                PatchSiteId::from_index(usize::MAX)
1086            }
1087        }
1088    }
1089
1090    pub fn try_record_patch_site(
1091        &mut self,
1092        offset: CodeOffset,
1093        kind: LabelUse,
1094        target_offset: CodeOffset,
1095    ) -> Result<PatchSiteId, AsmError> {
1096        if let Some(error) = self.error.clone() {
1097            return Err(error);
1098        }
1099        self.validate_patch_site_offset(offset, kind)?;
1100        let id = PatchSiteId::from_index(self.patch_sites.len());
1101        self.patch_sites.push(PendingPatchSite {
1102            offset,
1103            kind,
1104            target: PendingPatchTarget::Offset(target_offset),
1105            addend: 0,
1106        });
1107        Ok(id)
1108    }
1109
1110    pub fn record_label_patch_site(
1111        &mut self,
1112        offset: CodeOffset,
1113        label: Label,
1114        kind: LabelUse,
1115    ) -> PatchSiteId {
1116        match self.try_record_label_patch_site(offset, label, kind) {
1117            Ok(id) => id,
1118            Err(error) => {
1119                self.record_error(error);
1120                PatchSiteId::from_index(usize::MAX)
1121            }
1122        }
1123    }
1124
1125    pub fn try_record_label_patch_site(
1126        &mut self,
1127        offset: CodeOffset,
1128        label: Label,
1129        kind: LabelUse,
1130    ) -> Result<PatchSiteId, AsmError> {
1131        if let Some(error) = self.error.clone() {
1132            return Err(error);
1133        }
1134        if self.label_offsets.get(label.id() as usize).is_none() {
1135            return Err(AsmError::InvalidArgument);
1136        }
1137        self.validate_patch_site_offset(offset, kind)?;
1138        let id = PatchSiteId::from_index(self.patch_sites.len());
1139        self.patch_sites.push(PendingPatchSite {
1140            offset,
1141            kind,
1142            target: PendingPatchTarget::Label(label),
1143            addend: 0,
1144        });
1145        Ok(id)
1146    }
1147
1148    fn validate_patch_site_offset(
1149        &self,
1150        offset: CodeOffset,
1151        kind: LabelUse,
1152    ) -> Result<(), AsmError> {
1153        kind.validate_for_arch(self.env.arch())?;
1154        let end = (offset as usize)
1155            .checked_add(kind.patch_size())
1156            .ok_or(AsmError::TooLarge)?;
1157        if end > self.data.len() {
1158            return Err(AsmError::InvalidArgument);
1159        }
1160        Ok(())
1161    }
1162
1163    fn handle_fixup(&mut self, fixup: AsmFixup) -> Result<(), AsmError> {
1164        let AsmFixup {
1165            kind,
1166            label,
1167            offset,
1168        } = fixup;
1169        let start = offset;
1170        let end = (offset as usize)
1171            .checked_add(kind.patch_size())
1172            .ok_or(AsmError::TooLarge)?;
1173        if end > self.data.len() {
1174            return Err(AsmError::InvalidArgument);
1175        }
1176
1177        let label_offset = self.label_offset(label);
1178        if label_offset != u32::MAX {
1179            if !kind.can_reach(offset, label_offset) {
1180                if label_offset < offset {
1181                    self.emit_veneer(label, offset, kind)?;
1182                } else {
1183                    return Err(AsmError::TooLarge);
1184                }
1185            } else {
1186                let slice = &mut self.data[start as usize..end];
1187
1188                kind.patch(slice, start, label_offset);
1189            }
1190        } else {
1191            // If the offset of this label is not known at this time then
1192            // that means that a veneer is required because after this
1193            // island the target can't be in range of the original target.
1194            self.emit_veneer(label, offset, kind)?;
1195        }
1196        Ok(())
1197    }
1198
1199    /// Emits a "veneer" the `kind` code at `offset` to jump to `label`.
1200    ///
1201    /// This will generate extra machine code, using `kind`, to get a
1202    /// larger-jump-kind than `kind` allows. The code at `offset` is then
1203    /// patched to jump to our new code, and then the new code is enqueued for
1204    /// a fixup to get processed at some later time.
1205    pub fn emit_veneer(
1206        &mut self,
1207        label: Label,
1208        offset: CodeOffset,
1209        kind: LabelUse,
1210    ) -> Result<(), AsmError> {
1211        if let Some(error) = self.error.clone() {
1212            return Err(error);
1213        }
1214        kind.validate_for_arch(self.env.arch())?;
1215        if !kind.supports_veneer() {
1216            return Err(AsmError::UnsupportedInstruction {
1217                reason: "branch range requires an unsupported veneer",
1218            });
1219        }
1220        if self.label_offsets.get(label.id() as usize).is_none() {
1221            return Err(AsmError::InvalidArgument);
1222        }
1223        let start = offset as usize;
1224        let end = start
1225            .checked_add(kind.patch_size())
1226            .ok_or(AsmError::TooLarge)?;
1227        if end > self.data.len() {
1228            return Err(AsmError::InvalidArgument);
1229        }
1230        #[cfg(not(feature = "riscv"))]
1231        {
1232            let _ = (label, offset, kind);
1233            Err(AsmError::UnsupportedInstruction {
1234                reason: "RISC-V veneer support is disabled",
1235            })
1236        }
1237
1238        #[cfg(feature = "riscv")]
1239        {
1240            self.try_align_to(kind.align() as _)?;
1241            let veneer_offset = self.cur_offset();
1242            let slice = &mut self.data[start..end];
1243
1244            kind.patch(slice, offset, veneer_offset);
1245            let veneer_slice = self.get_appended_space(kind.veneer_size());
1246            let (veneer_fixup_off, veneer_label_use) =
1247                kind.generate_veneer(veneer_slice, veneer_offset);
1248
1249            // Register a new use of `label` with our new veneer fixup and
1250            // offset. This'll recalculate deadlines accordingly and
1251            // enqueue this fixup to get processed at some later
1252            // time.
1253            self.use_label_at_offset(veneer_fixup_off, label, veneer_label_use);
1254            if let Some(error) = self.error.clone() {
1255                return Err(error);
1256            }
1257            Ok(())
1258        }
1259    }
1260
1261    /// Reserve appended space and return a mutable slice referring to it.
1262    pub(crate) fn get_appended_space(&mut self, len: usize) -> &mut [u8] {
1263        let off = self.data.len();
1264        let new_len = self.data.len() + len;
1265        self.data.resize(new_len, 0);
1266        &mut self.data[off..]
1267
1268        // Post-invariant: as for `put1()`.
1269    }
1270
1271    /// Returns the maximal offset that islands can reach if `distance` more
1272    /// bytes are appended.
1273    ///
1274    /// This is used to determine if veneers need insertions since jumps that
1275    /// can't reach past this point must get a veneer of some form.
1276    fn worst_case_end_of_island(&self, distance: CodeOffset) -> CodeOffset {
1277        // Assume that all fixups will require veneers and that the veneers are
1278        // the worst-case size for each platform. This is an over-generalization
1279        // to avoid iterating over the `fixup_records` list or maintaining
1280        // information about it as we go along.
1281        let island_worst_case_size =
1282            ((self.fixup_records.len() + self.pending_fixup_records.len()) as u32) * 20
1283                + self.pending_constants_size;
1284        self.cur_offset()
1285            .saturating_add(distance)
1286            .saturating_add(island_worst_case_size)
1287    }
1288
1289    fn should_apply_fixup(&self, fixup: &AsmFixup, forced_threshold: CodeOffset) -> bool {
1290        let label_offset = self.label_offset(fixup.label);
1291        label_offset != u32::MAX
1292            || fixup.offset.saturating_add(fixup.kind.max_pos_range()) < forced_threshold
1293    }
1294    /// Is an island needed within the next N bytes?
1295    pub fn island_needed(&mut self, distance: CodeOffset) -> bool {
1296        let deadline = match self.fixup_records.peek() {
1297            Some(fixup) => fixup
1298                .offset
1299                .saturating_add(fixup.kind.max_pos_range())
1300                .min(self.pending_fixup_deadline),
1301            None => self.pending_fixup_deadline,
1302        };
1303
1304        deadline < u32::MAX && self.worst_case_end_of_island(distance) > deadline
1305    }
1306
1307    /// Emit all pending constants and required pending veneers.
1308    pub fn emit_island(&mut self, distance: CodeOffset) -> Result<(), AsmError> {
1309        if let Some(error) = self.error.clone() {
1310            return Err(error);
1311        }
1312        let forced_threshold = self.worst_case_end_of_island(distance);
1313
1314        for constant in core::mem::take(&mut self.pending_constants) {
1315            let (_, AsmConstant { align, size, .. }) = self.constants[constant.0 as usize];
1316            let label = self.constants[constant.0 as usize]
1317                .1
1318                .upcoming_label
1319                .take()
1320                .unwrap();
1321            self.try_align_to(align as _)?;
1322            self.try_bind_label(label)?;
1323            self.used_constants.push((constant, self.cur_offset()));
1324            self.get_appended_space(size);
1325        }
1326        // Either handle all pending fixups because they're ready or move them
1327        // onto the `BinaryHeap` tracking all pending fixups if they aren't
1328        // ready.
1329        for fixup in core::mem::take(&mut self.pending_fixup_records) {
1330            if self.should_apply_fixup(&fixup, forced_threshold) {
1331                self.handle_fixup(fixup)?;
1332            } else {
1333                self.fixup_records.push(fixup);
1334            }
1335        }
1336
1337        self.pending_fixup_deadline = u32::MAX;
1338
1339        while let Some(fixup) = self.fixup_records.peek() {
1340            // If this fixup shouldn't be applied, that means its label isn't
1341            // defined yet and there'll be remaining space to apply a veneer if
1342            // necessary in the future after this island. In that situation
1343            // because `fixup_records` is sorted by deadline this loop can
1344            // exit.
1345            if !self.should_apply_fixup(fixup, forced_threshold) {
1346                break;
1347            }
1348            let fixup = self.fixup_records.pop().unwrap();
1349            self.handle_fixup(fixup)?;
1350        }
1351        Ok(())
1352    }
1353
1354    fn finish_emission_maybe_forcing_veneers(&mut self) -> Result<(), AsmError> {
1355        while !self.pending_constants.is_empty()
1356            || !self.pending_fixup_records.is_empty()
1357            || !self.fixup_records.is_empty()
1358        {
1359            // `emit_island()` will emit any pending veneers and constants, and
1360            // as a side-effect, will also take care of any fixups with resolved
1361            // labels eagerly.
1362            self.emit_island(u32::MAX)?;
1363        }
1364        Ok(())
1365    }
1366
1367    #[cfg(feature = "x86")]
1368    fn relax_x86_branches(&mut self) -> Result<(), AsmError> {
1369        loop {
1370            let mut selected = None;
1371
1372            for (index, &candidate) in self.x86_branch_relaxations.iter().enumerate() {
1373                let start = candidate.opcode_offset as usize;
1374                let near_size = candidate.near_size as usize;
1375                let end = start.checked_add(near_size).ok_or(AsmError::TooLarge)?;
1376                if end > self.data.len() {
1377                    return Err(AsmError::InvalidState);
1378                }
1379                let valid_encoding = match candidate.near_size {
1380                    5 => self.data[start] == 0xE9 && candidate.short_opcode == 0xEB,
1381                    6 => {
1382                        self.data[start] == 0x0F
1383                            && self.data[start + 1] == candidate.short_opcode.wrapping_add(0x10)
1384                    }
1385                    _ => false,
1386                };
1387                if !valid_encoding {
1388                    return Err(AsmError::InvalidState);
1389                }
1390
1391                let target = self.label_offset(candidate.label);
1392                if target == u32::MAX {
1393                    return Err(AsmError::UnboundLabel);
1394                }
1395                let old_end = candidate
1396                    .opcode_offset
1397                    .checked_add(candidate.near_size as u32)
1398                    .ok_or(AsmError::TooLarge)?;
1399                if target > candidate.opcode_offset && target < old_end {
1400                    return Err(AsmError::InvalidState);
1401                }
1402
1403                let removed = candidate.near_size as u32 - 2;
1404                let target_after = if target >= old_end {
1405                    target - removed
1406                } else {
1407                    target
1408                };
1409                let short_end = candidate
1410                    .opcode_offset
1411                    .checked_add(2)
1412                    .ok_or(AsmError::TooLarge)?;
1413                let displacement = i64::from(target_after) - i64::from(short_end);
1414                let Ok(displacement) = i8::try_from(displacement) else {
1415                    continue;
1416                };
1417
1418                let cut_start = short_end;
1419                let cut_end = old_end;
1420                let preserves_alignment =
1421                    self.alignment_constraints.iter().all(|&(offset, align)| {
1422                        offset < cut_end || (offset - removed) & (align - 1) == 0
1423                    }) && self.patch_blocks.iter().all(|block| {
1424                        block.offset < cut_end || (block.offset - removed) & (block.align - 1) == 0
1425                    });
1426                if !preserves_alignment {
1427                    continue;
1428                }
1429
1430                selected = Some((index, candidate, displacement, cut_start, cut_end));
1431                break;
1432            }
1433
1434            let Some((index, candidate, displacement, cut_start, cut_end)) = selected else {
1435                return Ok(());
1436            };
1437            let removed = cut_end - cut_start;
1438            let fixup_offset = cut_end - LabelUse::X86JmpRel32.patch_size() as u32;
1439            let is_candidate_fixup = |fixup: &AsmFixup| {
1440                fixup.offset == fixup_offset
1441                    && fixup.label == candidate.label
1442                    && fixup.kind == LabelUse::X86JmpRel32
1443            };
1444
1445            let overlaps_cut = |offset: CodeOffset, size: usize| -> Result<bool, AsmError> {
1446                let end = offset
1447                    .checked_add(u32::try_from(size).map_err(|_| AsmError::TooLarge)?)
1448                    .ok_or(AsmError::TooLarge)?;
1449                Ok(offset < cut_end && end > cut_start)
1450            };
1451            for reloc in &self.relocs {
1452                if overlaps_cut(reloc.offset, relocation_patch_size(reloc.kind)?)? {
1453                    return Err(AsmError::InvalidState);
1454                }
1455            }
1456            for fixup in self
1457                .pending_fixup_records
1458                .iter()
1459                .chain(self.fixup_records.iter())
1460            {
1461                if !is_candidate_fixup(fixup)
1462                    && overlaps_cut(fixup.offset, fixup.kind.patch_size())?
1463                {
1464                    return Err(AsmError::InvalidState);
1465                }
1466            }
1467            for block in &self.patch_blocks {
1468                if overlaps_cut(block.offset, block.size as usize)? {
1469                    return Err(AsmError::InvalidState);
1470                }
1471            }
1472            for site in &self.patch_sites {
1473                if overlaps_cut(site.offset, site.kind.patch_size())?
1474                    || matches!(site.target, PendingPatchTarget::Offset(offset) if (cut_start..cut_end).contains(&offset))
1475                {
1476                    return Err(AsmError::InvalidState);
1477                }
1478            }
1479            if self
1480                .label_offsets
1481                .iter()
1482                .any(|&offset| offset != u32::MAX && (cut_start..cut_end).contains(&offset))
1483                || self
1484                    .used_constants
1485                    .iter()
1486                    .any(|&(_, offset)| (cut_start..cut_end).contains(&offset))
1487                || self
1488                    .alignment_constraints
1489                    .iter()
1490                    .any(|&(offset, _)| (cut_start..cut_end).contains(&offset))
1491            {
1492                return Err(AsmError::InvalidState);
1493            }
1494            for (other_index, other) in self.x86_branch_relaxations.iter().enumerate() {
1495                if other_index != index
1496                    && overlaps_cut(other.opcode_offset, other.near_size as usize)?
1497                {
1498                    return Err(AsmError::InvalidState);
1499                }
1500            }
1501
1502            let start = candidate.opcode_offset as usize;
1503            self.data[start] = candidate.short_opcode;
1504            self.data[start + 1] = displacement as u8;
1505            self.data.drain(cut_start as usize..cut_end as usize);
1506            self.x86_branch_relaxations.remove(index);
1507
1508            self.pending_fixup_records
1509                .retain(|fixup| !is_candidate_fixup(fixup));
1510            let mut fixups = core::mem::take(&mut self.fixup_records).into_vec();
1511            fixups.retain(|fixup| !is_candidate_fixup(fixup));
1512
1513            let rebase = |offset: &mut CodeOffset| -> Result<(), AsmError> {
1514                if *offset >= cut_end {
1515                    *offset -= removed;
1516                } else if *offset >= cut_start {
1517                    return Err(AsmError::InvalidState);
1518                }
1519                Ok(())
1520            };
1521
1522            for reloc in &mut self.relocs {
1523                rebase(&mut reloc.offset)?;
1524            }
1525            for offset in &mut self.label_offsets {
1526                if *offset != u32::MAX {
1527                    rebase(offset)?;
1528                }
1529            }
1530            for fixup in &mut self.pending_fixup_records {
1531                rebase(&mut fixup.offset)?;
1532            }
1533            for fixup in &mut fixups {
1534                rebase(&mut fixup.offset)?;
1535            }
1536            self.fixup_records = BinaryHeap::from(fixups);
1537            self.pending_fixup_deadline = self
1538                .pending_fixup_records
1539                .iter()
1540                .map(AsmFixup::deadline)
1541                .min()
1542                .unwrap_or(u32::MAX);
1543            for (_, offset) in &mut self.used_constants {
1544                rebase(offset)?;
1545            }
1546            for block in &mut self.patch_blocks {
1547                rebase(&mut block.offset)?;
1548            }
1549            for site in &mut self.patch_sites {
1550                rebase(&mut site.offset)?;
1551                if let PendingPatchTarget::Offset(offset) = &mut site.target {
1552                    rebase(offset)?;
1553                }
1554            }
1555            for relaxation in &mut self.x86_branch_relaxations {
1556                rebase(&mut relaxation.opcode_offset)?;
1557            }
1558            for (offset, _) in &mut self.alignment_constraints {
1559                rebase(offset)?;
1560            }
1561        }
1562    }
1563
1564    /// Reject finalization failures that can be determined before island
1565    /// emission mutates the buffer.
1566    fn preflight_finalization(&self) -> Result<(), AsmError> {
1567        for fixup in self
1568            .pending_fixup_records
1569            .iter()
1570            .chain(self.fixup_records.iter())
1571        {
1572            let end = (fixup.offset as usize)
1573                .checked_add(fixup.kind.patch_size())
1574                .ok_or(AsmError::TooLarge)?;
1575            if end > self.data.len() {
1576                return Err(AsmError::InvalidArgument);
1577            }
1578            let label_offset = self.label_offset(fixup.label);
1579            if label_offset == u32::MAX {
1580                return Err(AsmError::UnboundLabel);
1581            }
1582            if fixup.kind.can_reach(fixup.offset, label_offset) {
1583                continue;
1584            }
1585            if matches!(
1586                fixup.kind,
1587                LabelUse::A64Branch14 | LabelUse::A64Branch19 | LabelUse::A64Branch26
1588            ) && !fixup.kind.supports_veneer()
1589            {
1590                return Err(AsmError::UnsupportedInstruction {
1591                    reason: "AArch64 branch veneers are not implemented",
1592                });
1593            }
1594            if label_offset >= fixup.offset {
1595                return Err(AsmError::TooLarge);
1596            }
1597            if !fixup.kind.supports_veneer() {
1598                return Err(AsmError::UnsupportedInstruction {
1599                    reason: "branch range requires an unsupported veneer",
1600                });
1601            }
1602            #[cfg(not(feature = "riscv"))]
1603            return Err(AsmError::UnsupportedInstruction {
1604                reason: "RISC-V veneer support is disabled",
1605            });
1606        }
1607        Ok(())
1608    }
1609
1610    fn finish_constants(&mut self) -> u32 {
1611        let mut alignment = 32;
1612
1613        for (constant, offset) in core::mem::take(&mut self.used_constants) {
1614            let constant = &self.constants[constant.0 as usize].0;
1615            let data = constant.as_slice();
1616            self.data[offset as usize..][..data.len()].copy_from_slice(data);
1617            alignment = constant.alignment().max(alignment);
1618        }
1619
1620        alignment as _
1621    }
1622
1623    fn resolve_patch_catalog(&self, validate_ranges: bool) -> Result<PatchCatalog, AsmError> {
1624        let mut blocks = SmallVec::new();
1625        let mut sites = SmallVec::new();
1626
1627        for block in &self.patch_blocks {
1628            blocks.push(PatchBlock {
1629                offset: block.offset,
1630                size: block.size,
1631                align: block.align,
1632            });
1633        }
1634
1635        for site in &self.patch_sites {
1636            let target_offset = match site.target {
1637                PendingPatchTarget::Offset(offset) => offset,
1638                PendingPatchTarget::Label(label) => self.label_offset(label),
1639            };
1640
1641            if target_offset == u32::MAX {
1642                return Err(AsmError::InvalidState);
1643            }
1644
1645            if validate_ranges && !site.kind.can_reach(site.offset, target_offset) {
1646                return Err(AsmError::TooLarge);
1647            }
1648
1649            sites.push(PatchSite {
1650                offset: site.offset,
1651                kind: site.kind,
1652                current_target: target_offset,
1653                addend: site.addend,
1654            });
1655        }
1656
1657        Ok(PatchCatalog::with_parts(self.env.arch(), blocks, sites))
1658    }
1659
1660    fn resolved_defined_symbols(&self) -> SmallVec<[(ExternalName, CodeOffset); 4]> {
1661        // Unbound labels keep the sentinel offset; the linker reports them as
1662        // an error.
1663        self.defined_symbols
1664            .iter()
1665            .map(|(name, label)| (name.clone(), self.label_offset(*label)))
1666            .collect()
1667    }
1668
1669    pub fn finish_patched(mut self) -> Result<CodeBufferFinalized, AsmError> {
1670        if let Some(error) = self.error.take() {
1671            return Err(error);
1672        }
1673        if self.has_unbound_labels() {
1674            return Err(AsmError::UnboundLabel);
1675        }
1676        self.preflight_finalization()?;
1677        #[cfg(feature = "x86")]
1678        self.relax_x86_branches()?;
1679        #[cfg(feature = "x86")]
1680        self.preflight_finalization()?;
1681        self.finish_emission_maybe_forcing_veneers()?;
1682        let patch_catalog = self.resolve_patch_catalog(true)?;
1683        let alignment = self.finish_constants();
1684        let defined_symbols = self.resolved_defined_symbols();
1685        Ok(CodeBufferFinalized {
1686            data: self.data,
1687            relocs: self.relocs,
1688            symbols: self.symbols,
1689            label_offsets: self.label_offsets,
1690            defined_symbols,
1691            alignment,
1692            patch_catalog,
1693        })
1694    }
1695
1696    pub fn finish(&mut self) -> Result<CodeBufferFinalized, AsmError> {
1697        if let Some(error) = self.error.clone() {
1698            return Err(error);
1699        }
1700        if self.has_unbound_labels() {
1701            return Err(AsmError::UnboundLabel);
1702        }
1703        self.preflight_finalization()?;
1704        #[cfg(feature = "x86")]
1705        self.relax_x86_branches()?;
1706        #[cfg(feature = "x86")]
1707        self.preflight_finalization()?;
1708        self.finish_emission_maybe_forcing_veneers()?;
1709        let patch_catalog = self.resolve_patch_catalog(false)?;
1710        let alignment = self.finish_constants();
1711        Ok(CodeBufferFinalized {
1712            data: self.data.clone(),
1713            relocs: self.relocs.clone(),
1714            symbols: self.symbols.clone(),
1715            label_offsets: self.label_offsets.clone(),
1716            defined_symbols: self.resolved_defined_symbols(),
1717            alignment,
1718            patch_catalog,
1719        })
1720    }
1721
1722    fn has_unbound_labels(&self) -> bool {
1723        let is_unbound = |label: Label| {
1724            self.label_offsets
1725                .get(label.id() as usize)
1726                .is_none_or(|offset| *offset == u32::MAX)
1727        };
1728        self.pending_fixup_records
1729            .iter()
1730            .any(|fixup| is_unbound(fixup.label))
1731            || self
1732                .fixup_records
1733                .iter()
1734                .any(|fixup| is_unbound(fixup.label))
1735            || self
1736                .relocs
1737                .iter()
1738                .any(|reloc| matches!(reloc.target, RelocTarget::Label(label) if is_unbound(label)))
1739            || self
1740                .defined_symbols
1741                .iter()
1742                .any(|(_, label)| is_unbound(*label))
1743            || self.patch_sites.iter().any(
1744                |site| matches!(site.target, PendingPatchTarget::Label(label) if is_unbound(label)),
1745            )
1746    }
1747}
1748
1749#[derive(Clone, PartialEq, Eq, Debug, Hash)]
1750pub enum ConstantData {
1751    WellKnown(&'static [u8]),
1752    U64([u8; 8]),
1753    Bytes(Vec<u8>),
1754}
1755
1756impl ConstantData {
1757    pub fn as_slice(&self) -> &[u8] {
1758        match self {
1759            ConstantData::WellKnown(data) => data,
1760            ConstantData::U64(data) => data.as_ref(),
1761            ConstantData::Bytes(data) => data,
1762        }
1763    }
1764
1765    pub fn alignment(&self) -> usize {
1766        if self.as_slice().len() <= 8 { 8 } else { 16 }
1767    }
1768}
1769
1770/// A use of a constant by one or mroe assembly instructions.
1771#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1772pub struct Constant(pub(crate) u32);
1773
1774impl From<&'static str> for ConstantData {
1775    fn from(value: &'static str) -> Self {
1776        Self::WellKnown(value.as_bytes())
1777    }
1778}
1779
1780impl From<[u8; 8]> for ConstantData {
1781    fn from(value: [u8; 8]) -> Self {
1782        Self::U64(value)
1783    }
1784}
1785
1786impl From<Vec<u8>> for ConstantData {
1787    fn from(value: Vec<u8>) -> Self {
1788        Self::Bytes(value)
1789    }
1790}
1791
1792impl From<&'static [u8]> for ConstantData {
1793    fn from(value: &'static [u8]) -> Self {
1794        Self::WellKnown(value)
1795    }
1796}
1797
1798impl From<u64> for ConstantData {
1799    fn from(value: u64) -> Self {
1800        Self::U64(value.to_ne_bytes())
1801    }
1802}
1803
1804/// Offset in bytes from the beginning of the function.
1805///
1806/// Cranelift can be used as a cross compiler, so we don't want to use a type like `usize` which
1807/// depends on the *host* platform, not the *target* platform.
1808pub type CodeOffset = u32;
1809
1810/// Addend to add to the symbol value.
1811pub type Addend = i64;
1812
1813/// Relocation kinds for every ISA
1814#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
1815pub enum Reloc {
1816    /// absolute 4-byte
1817    Abs4,
1818    /// absolute 8-byte
1819    Abs8,
1820    /// x86 PC-relative 4-byte
1821    X86PCRel4,
1822    /// x86 call to PC-relative 4-byte
1823    X86CallPCRel4,
1824    /// x86 call to PLT-relative 4-byte
1825    X86CallPLTRel4,
1826    /// x86 GOT PC-relative 4-byte
1827    X86GOTPCRel4,
1828    /// The 32-bit offset of the target from the beginning of its section.
1829    /// Equivalent to `IMAGE_REL_AMD64_SECREL`.
1830    /// See: [PE Format](https://docs.microsoft.com/en-us/windows/win32/debug/pe-format)
1831    X86SecRel,
1832    /// Arm32 call target
1833    Arm32Call,
1834    /// Arm64 call target. Encoded as bottom 26 bits of instruction. This
1835    /// value is sign-extended, multiplied by 4, and added to the PC of
1836    /// the call instruction to form the destination address.
1837    Arm64Call,
1838
1839    /// Elf x86_64 32 bit signed PC relative offset to two GOT entries for GD symbol.
1840    ElfX86_64TlsGd,
1841
1842    /// Mach-O x86_64 32 bit signed PC relative offset to a `__thread_vars` entry.
1843    MachOX86_64Tlv,
1844
1845    /// Mach-O Aarch64 TLS
1846    /// PC-relative distance to the page of the TLVP slot.
1847    MachOAarch64TlsAdrPage21,
1848
1849    /// Mach-O Aarch64 TLS
1850    /// Offset within page of TLVP slot.
1851    MachOAarch64TlsAdrPageOff12,
1852
1853    /// Aarch64 TLSDESC Adr Page21
1854    /// This is equivalent to `R_AARCH64_TLSDESC_ADR_PAGE21` in the [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#57105thread-local-storage-descriptors)
1855    Aarch64TlsDescAdrPage21,
1856
1857    /// Aarch64 TLSDESC Ld64 Lo12
1858    /// This is equivalent to `R_AARCH64_TLSDESC_LD64_LO12` in the [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#57105thread-local-storage-descriptors)
1859    Aarch64TlsDescLd64Lo12,
1860
1861    /// Aarch64 TLSDESC Add Lo12
1862    /// This is equivalent to `R_AARCH64_TLSGD_ADD_LO12` in the [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#57105thread-local-storage-descriptors)
1863    Aarch64TlsDescAddLo12,
1864
1865    /// Aarch64 TLSDESC Call
1866    /// This is equivalent to `R_AARCH64_TLSDESC_CALL` in the [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#57105thread-local-storage-descriptors)
1867    Aarch64TlsDescCall,
1868
1869    /// AArch64 GOT Page
1870    /// Set the immediate value of an ADRP to bits 32:12 of X; check that –2^32 <= X < 2^32
1871    /// This is equivalent to `R_AARCH64_ADR_GOT_PAGE` (311) in the  [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#static-aarch64-relocations)
1872    Aarch64AdrGotPage21,
1873
1874    /// AArch64 GOT Low bits
1875
1876    /// Set the LD/ST immediate field to bits 11:3 of X. No overflow check; check that X&7 = 0
1877    /// This is equivalent to `R_AARCH64_LD64_GOT_LO12_NC` (312) in the  [aaelf64](https://github.com/ARM-software/abi-aa/blob/2bcab1e3b22d55170c563c3c7940134089176746/aaelf64/aaelf64.rst#static-aarch64-relocations)
1878    Aarch64Ld64GotLo12Nc,
1879
1880    /// Equivalent of `R_AARCH64_ADR_PREL_PG_HI21`.
1881    Aarch64AdrPrelPgHi21,
1882    /// Equivalent of `R_AARCH64_ADD_ABS_LO12_NC`.
1883    Aarch64AddAbsLo12Nc,
1884
1885    /// RISC-V Absolute address: 64-bit address.
1886    RiscvAbs8,
1887
1888    /// RISC-V Call PLT: 32-bit PC-relative function call, macros call, tail (PIC)
1889    ///
1890    /// Despite having PLT in the name, this relocation is also used for normal calls.
1891    /// The non-PLT version of this relocation has been deprecated.
1892    ///
1893    /// This is the `R_RISCV_CALL_PLT` relocation from the RISC-V ELF psABI document.
1894    /// <https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#procedure-calls>
1895    RiscvCallPlt,
1896
1897    /// RISC-V TLS GD: High 20 bits of 32-bit PC-relative TLS GD GOT reference,
1898    ///
1899    /// This is the `R_RISCV_TLS_GD_HI20` relocation from the RISC-V ELF psABI document.
1900    /// <https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#global-dynamic>
1901    RiscvTlsGdHi20,
1902
1903    /// Low 12 bits of a 32-bit PC-relative relocation (I-Type instruction)
1904    ///
1905    /// This is the `R_RISCV_PCREL_LO12_I` relocation from the RISC-V ELF psABI document.
1906    /// <https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#pc-relative-symbol-addresses>
1907    RiscvPCRelLo12I,
1908
1909    /// High 20 bits of a 32-bit PC-relative GOT offset relocation
1910    ///
1911    /// This is the `R_RISCV_GOT_HI20` relocation from the RISC-V ELF psABI document.
1912    /// <https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#pc-relative-symbol-addresses>
1913    RiscvGotHi20,
1914}
1915
1916impl Reloc {
1917    const fn supports_arch(self, arch: Arch) -> bool {
1918        match self {
1919            Self::Abs4 | Self::Abs8 => true,
1920            Self::X86PCRel4
1921            | Self::X86CallPCRel4
1922            | Self::X86CallPLTRel4
1923            | Self::X86GOTPCRel4
1924            | Self::X86SecRel
1925            | Self::ElfX86_64TlsGd
1926            | Self::MachOX86_64Tlv => {
1927                cfg!(feature = "x86") && matches!(arch, Arch::X86 | Arch::X64)
1928            }
1929            Self::Arm32Call => false,
1930            Self::Arm64Call
1931            | Self::MachOAarch64TlsAdrPage21
1932            | Self::MachOAarch64TlsAdrPageOff12
1933            | Self::Aarch64TlsDescAdrPage21
1934            | Self::Aarch64TlsDescLd64Lo12
1935            | Self::Aarch64TlsDescAddLo12
1936            | Self::Aarch64TlsDescCall
1937            | Self::Aarch64AdrGotPage21
1938            | Self::Aarch64Ld64GotLo12Nc
1939            | Self::Aarch64AdrPrelPgHi21
1940            | Self::Aarch64AddAbsLo12Nc => {
1941                cfg!(feature = "aarch64") && matches!(arch, Arch::AArch64)
1942            }
1943            Self::RiscvAbs8
1944            | Self::RiscvCallPlt
1945            | Self::RiscvTlsGdHi20
1946            | Self::RiscvPCRelLo12I
1947            | Self::RiscvGotHi20 => {
1948                cfg!(feature = "riscv") && matches!(arch, Arch::RISCV32 | Arch::RISCV64)
1949            }
1950        }
1951    }
1952}
1953
1954#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug, Hash)]
1955pub enum LabelUse {
1956    X86JmpRel32,
1957    /// 20-bit branch offset (unconditional branches). PC-rel, offset is
1958    /// imm << 1. Immediate is 20 signed bits. Use in Jal instructions.
1959    RVJal20,
1960    /// The unconditional jump instructions all use PC-relative
1961    /// addressing to help support position independent code. The JALR
1962    /// instruction was defined to enable a two-instruction sequence to
1963    /// jump anywhere in a 32-bit absolute address range. A LUI
1964    /// instruction can first load rs1 with the upper 20 bits of a
1965    /// target address, then JALR can add in the lower bits. Similarly,
1966    /// AUIPC then JALR can jump anywhere in a 32-bit pc-relative
1967    /// address range.
1968    RVPCRel32,
1969
1970    /// All branch instructions use the B-type instruction format. The
1971    /// 12-bit B-immediate encodes signed offsets in multiples of 2, and
1972    /// is added to the current pc to give the target address. The
1973    /// conditional branch range is ±4 KiB.
1974    RVB12,
1975
1976    /// Equivalent to the `R_RISCV_PCREL_HI20` relocation, Allows setting
1977    /// the immediate field of an `auipc` instruction.
1978    RVPCRelHi20,
1979
1980    /// Similar to the `R_RISCV_PCREL_LO12_I` relocation but pointing to
1981    /// the final address, instead of the `PCREL_HI20` label. Allows setting
1982    /// the immediate field of I Type instructions such as `addi` or `lw`.
1983    ///
1984    /// Since we currently don't support offsets in labels, this relocation has
1985    /// an implicit offset of 4.
1986    RVPCRelLo12I,
1987
1988    /// 11-bit PC-relative jump offset. Equivalent to the `RVC_JUMP` relocation
1989    RVCJump,
1990    /// 9-bit PC-relative branch offset.
1991    RVCB9,
1992    /// 14-bit branch offset (conditional branches). PC-rel, offset is imm <<
1993    /// 2. Immediate is 14 signed bits, in bits 18:5. Used by tbz and tbnz.
1994    A64Branch14,
1995    /// 19-bit branch offset (conditional branches). PC-rel, offset is imm << 2. Immediate is 19
1996    /// signed bits, in bits 23:5. Used by cbz, cbnz, b.cond.
1997    A64Branch19,
1998    /// 26-bit branch offset (unconditional branches). PC-rel, offset is imm << 2. Immediate is 26
1999    /// signed bits, in bits 25:0. Used by b, bl.
2000    A64Branch26,
2001    /// 19-bit offset for LDR (load literal). PC-rel, offset is imm << 2. Immediate is 19 signed bits,
2002    /// in bits 23:5.
2003    A64Ldr19,
2004    /// 21-bit offset for ADR (get address of label). PC-rel, offset is not shifted. Immediate is
2005    /// 21 signed bits, with high 19 bits in bits 23:5 and low 2 bits in bits 30:29.
2006    A64Adr21,
2007    /// 21-bit offset for ADRP (get address of label). PC-rel, offset is shifted. Immediate is
2008    /// 21 signed bits, with high 19 bits in bits 23:5 and low 2 bits in bits 30:29.
2009    A64Adrp21,
2010    A64Ldr12,
2011
2012    A64AddAbsLo12,
2013}
2014
2015impl LabelUse {
2016    fn validate_for_arch(self, arch: Arch) -> Result<(), AsmError> {
2017        if self == Self::A64Ldr12 {
2018            return Err(AsmError::UnsupportedInstruction {
2019                reason: "AArch64 LDR12 label patching is not implemented",
2020            });
2021        }
2022        if !self.supports_arch(arch) {
2023            return Err(AsmError::InvalidArch);
2024        }
2025        Ok(())
2026    }
2027
2028    const fn supports_arch(self, arch: Arch) -> bool {
2029        match self {
2030            Self::X86JmpRel32 => cfg!(feature = "x86") && matches!(arch, Arch::X86 | Arch::X64),
2031            Self::RVJal20
2032            | Self::RVPCRel32
2033            | Self::RVB12
2034            | Self::RVPCRelHi20
2035            | Self::RVPCRelLo12I
2036            | Self::RVCJump
2037            | Self::RVCB9 => {
2038                cfg!(feature = "riscv") && matches!(arch, Arch::RISCV32 | Arch::RISCV64)
2039            }
2040            Self::A64Branch14
2041            | Self::A64Branch19
2042            | Self::A64Branch26
2043            | Self::A64Ldr19
2044            | Self::A64Adr21
2045            | Self::A64Adrp21
2046            | Self::A64Ldr12
2047            | Self::A64AddAbsLo12 => cfg!(feature = "aarch64") && matches!(arch, Arch::AArch64),
2048        }
2049    }
2050
2051    pub fn can_reach(&self, use_offset: CodeOffset, label_offset: CodeOffset) -> bool {
2052        let delta = (label_offset as i64) - (use_offset as i64);
2053
2054        match self {
2055            Self::X86JmpRel32 => {
2056                let disp = delta - 4;
2057                i32::try_from(disp).is_ok()
2058            }
2059            Self::RVJal20 => delta % 2 == 0 && (-(1 << 20)..=((1 << 20) - 2)).contains(&delta),
2060            Self::RVB12 => delta % 2 == 0 && (-(1 << 12)..=((1 << 12) - 2)).contains(&delta),
2061            Self::RVCJump => delta % 2 == 0 && (-(1 << 11)..=((1 << 11) - 2)).contains(&delta),
2062            Self::RVCB9 => delta % 2 == 0 && (-(1 << 8)..=((1 << 8) - 2)).contains(&delta),
2063            Self::RVPCRelHi20 | Self::RVPCRelLo12I | Self::RVPCRel32 => {
2064                i32::try_from(delta).is_ok()
2065            }
2066            Self::A64Branch14 => delta % 4 == 0 && (-(1 << 15)..=((1 << 15) - 4)).contains(&delta),
2067            Self::A64Branch19 | Self::A64Ldr19 => {
2068                delta % 4 == 0 && (-(1 << 20)..=((1 << 20) - 4)).contains(&delta)
2069            }
2070            Self::A64Branch26 => delta % 4 == 0 && (-(1 << 27)..=((1 << 27) - 4)).contains(&delta),
2071            Self::A64Adr21 => (-(1 << 20)..=((1 << 20) - 1)).contains(&delta),
2072            Self::A64Adrp21 => {
2073                let page_delta = ((label_offset & !0xfff) as i64) - ((use_offset & !0xfff) as i64);
2074                page_delta % 4096 == 0 && (-(1 << 32)..=((1 << 32) - 4096)).contains(&page_delta)
2075            }
2076
2077            Self::A64AddAbsLo12 => {
2078                delta % 4096 == delta && (-(1 << 12)..=(1 << 12) - 1).contains(&delta)
2079            }
2080
2081            Self::A64Ldr12 => true,
2082        }
2083    }
2084
2085    /// Maximum PC-relative range (positive), inclusive.
2086    pub const fn max_pos_range(self) -> CodeOffset {
2087        match self {
2088            LabelUse::RVJal20 => ((1 << 19) - 1) * 2,
2089            LabelUse::RVPCRelLo12I | LabelUse::RVPCRelHi20 | LabelUse::RVPCRel32 => {
2090                let imm20_max: i64 = ((1 << 19) - 1) << 12;
2091                let imm12_max = (1 << 11) - 1;
2092                (imm20_max + imm12_max) as _
2093            }
2094            LabelUse::RVB12 => ((1 << 11) - 1) * 2,
2095            LabelUse::RVCB9 => ((1 << 8) - 1) * 2,
2096            LabelUse::RVCJump => ((1 << 10) - 1) * 2,
2097            LabelUse::X86JmpRel32 => i32::MAX as _,
2098            _ => u32::MAX,
2099        }
2100    }
2101
2102    pub const fn max_neg_range(self) -> CodeOffset {
2103        match self {
2104            LabelUse::RVPCRel32 => {
2105                let imm20_max: i64 = (1 << 19) << 12;
2106                let imm12_max = 1 << 11;
2107                (-imm20_max - imm12_max) as CodeOffset
2108            }
2109            _ => self.max_pos_range() + 2,
2110        }
2111    }
2112
2113    pub const fn patch_size(&self) -> usize {
2114        match self {
2115            Self::X86JmpRel32 => 4,
2116            Self::RVCJump | Self::RVCB9 => 2,
2117            Self::RVJal20 | Self::RVB12 | Self::RVPCRelHi20 | Self::RVPCRelLo12I => 4,
2118            Self::RVPCRel32 => 8,
2119            _ => 4,
2120        }
2121    }
2122
2123    pub const fn align(&self) -> usize {
2124        match self {
2125            Self::X86JmpRel32 => 1,
2126            Self::RVCJump => 4,
2127            Self::RVJal20 | Self::RVB12 | Self::RVCB9 | Self::RVPCRelHi20 | Self::RVPCRelLo12I => 4,
2128            Self::RVPCRel32 => 4,
2129            _ => 4,
2130        }
2131    }
2132
2133    pub const fn supports_veneer(&self) -> bool {
2134        matches!(self, Self::RVB12 | Self::RVJal20 | Self::RVCJump)
2135    }
2136
2137    #[cfg(feature = "riscv")]
2138    pub(crate) fn veneer_size(&self) -> usize {
2139        debug_assert!(self.supports_veneer());
2140        8
2141    }
2142
2143    #[cfg(feature = "riscv")]
2144    pub(crate) fn generate_veneer(
2145        &self,
2146        buffer: &mut [u8],
2147        veneer_offset: CodeOffset,
2148    ) -> (CodeOffset, Self) {
2149        debug_assert!(self.supports_veneer());
2150        let base = riscv::X31;
2151
2152        {
2153            let x = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2154                .encode()
2155                .set_rd(base.id())
2156                .value
2157                .to_le_bytes();
2158            buffer[0] = x[0];
2159            buffer[1] = x[1];
2160            buffer[2] = x[2];
2161            buffer[3] = x[3];
2162        }
2163
2164        {
2165            let x = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2166                .encode()
2167                .set_rd(riscv::ZERO.id())
2168                .set_rs1(base.id())
2169                .value
2170                .to_le_bytes();
2171            buffer[4] = x[0];
2172            buffer[5] = x[1];
2173            buffer[6] = x[2];
2174            buffer[7] = x[3];
2175        }
2176
2177        (veneer_offset, LabelUse::RVPCRel32)
2178    }
2179    pub(crate) fn patch(
2180        &self,
2181        buffer: &mut [u8],
2182        use_offset: CodeOffset,
2183        label_offset: CodeOffset,
2184    ) {
2185        let addend = match self {
2186            Self::X86JmpRel32 => i64::from(u32::from_le_bytes([
2187                buffer[0], buffer[1], buffer[2], buffer[3],
2188            ])),
2189            _ => 0,
2190        };
2191
2192        self.patch_with_addend(buffer, use_offset, label_offset, addend);
2193    }
2194
2195    pub(crate) fn patch_with_addend(
2196        &self,
2197        buffer: &mut [u8],
2198        use_offset: CodeOffset,
2199        label_offset: CodeOffset,
2200        addend: i64,
2201    ) {
2202        let pc_reli = (label_offset as i64) - (use_offset as i64);
2203
2204        let pc_rel = pc_reli as u32;
2205
2206        match self {
2207            Self::X86JmpRel32 => {
2208                let value = pc_rel.wrapping_add(addend as u32).wrapping_sub(4);
2209
2210                buffer.copy_from_slice(&value.to_le_bytes());
2211            }
2212
2213            Self::RVJal20 => {
2214                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2215                let offset = pc_rel;
2216                let v = ((offset >> 12 & 0b1111_1111) << 12)
2217                    | ((offset >> 11 & 0b1) << 20)
2218                    | ((offset >> 1 & 0b11_1111_1111) << 21)
2219                    | ((offset >> 20 & 0b1) << 31);
2220                buffer[0..4].clone_from_slice(&u32::to_le_bytes(insn | v));
2221            }
2222
2223            Self::RVPCRel32 => {
2224                #[cfg(feature = "riscv")]
2225                {
2226                    let (imm20, imm12) = generate_imm(pc_rel as u64);
2227                    let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2228                    let insn2 = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]);
2229
2230                    let auipc = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2231                        .encode()
2232                        .set_imm20(0);
2233                    let jalr = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2234                        .encode()
2235                        .set_rd(0)
2236                        .set_rs1(0)
2237                        .set_imm12(0);
2238
2239                    buffer[0..4].copy_from_slice(&(insn | auipc.value | imm20).to_le_bytes());
2240                    buffer[4..8].copy_from_slice(&(insn2 | jalr.value | imm12).to_le_bytes());
2241                }
2242                #[cfg(not(feature = "riscv"))]
2243                {
2244                    panic!("RISC-V veneers aren't supported without the `riscv` feature");
2245                }
2246            }
2247
2248            Self::RVB12 => {
2249                #[cfg(feature = "riscv")]
2250                {
2251                    let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2252                    let offset = pc_rel;
2253                    let v = ((offset >> 11 & 0b1) << 7)
2254                        | ((offset >> 1 & 0b1111) << 8)
2255                        | ((offset >> 5 & 0b11_1111) << 25)
2256                        | ((offset >> 12 & 0b1) << 31);
2257                    buffer[0..4].clone_from_slice(&u32::to_le_bytes(insn | v));
2258                }
2259                #[cfg(not(feature = "riscv"))]
2260                {
2261                    panic!("RISC-V veneers aren't supported without the `riscv` feature");
2262                }
2263            }
2264
2265            Self::RVPCRelHi20 => {
2266                // See https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#pc-relative-symbol-addresses
2267                //
2268                // We need to add 0x800 to ensure that we land at the next page as soon as it goes out of range for the
2269                // Lo12 relocation. That relocation is signed and has a maximum range of -2048..2047. So when we get an
2270                // offset of 2048, we need to land at the next page and subtract instead.
2271                let offset = pc_reli as u32;
2272                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2273                let hi20 = offset.wrapping_add(0x800) >> 12;
2274                let insn = (insn & 0xfff) | (hi20 << 12);
2275                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2276            }
2277
2278            Self::RVPCRelLo12I => {
2279                // `offset` is the offset from the current instruction to the target address.
2280                //
2281                // However we are trying to compute the offset to the target address from the previous instruction.
2282                // The previous instruction should be the one that contains the PCRelHi20 relocation and
2283                // stores/references the program counter (`auipc` usually).
2284                //
2285                // Since we are trying to compute the offset from the previous instruction, we can
2286                // represent it as offset = target_address - (current_instruction_address - 4)
2287                // which is equivalent to offset = target_address - current_instruction_address + 4.
2288                //
2289                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2290
2291                let lo12 = (pc_reli + 4) as u32 & 0xfff;
2292                let insn = (insn & 0xFFFFF) | (lo12 << 20);
2293                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2294            }
2295
2296            Self::RVCJump => {
2297                debug_assert!(pc_rel & 1 == 0);
2298
2299                #[cfg(feature = "riscv")]
2300                {
2301                    let insn = riscv::opcodes::Inst::new(riscv::Opcode::CJ)
2302                        .encode()
2303                        .set_c_imm12(pc_rel as _);
2304                    buffer[0..2].clone_from_slice(&(insn.value as u16).to_le_bytes());
2305                }
2306                #[cfg(not(feature = "riscv"))]
2307                {
2308                    panic!("RISC-V jumps aren't supported without the `riscv` feature");
2309                }
2310            }
2311
2312            Self::RVCB9 => {
2313                debug_assert!(pc_rel & 1 == 0);
2314
2315                #[cfg(feature = "riscv")]
2316                {
2317                    let insn = riscv::opcodes::Inst::new(riscv::Opcode::BEQZ)
2318                        .encode()
2319                        .set_c_bimm9lohi(pc_rel as _);
2320                    buffer[0..2].clone_from_slice(&(insn.value as u16).to_le_bytes());
2321                }
2322                #[cfg(not(feature = "riscv"))]
2323                {
2324                    panic!("RISC-V veneers aren't supported without the `riscv` feature");
2325                }
2326            }
2327
2328            Self::A64Branch14 => {
2329                debug_assert!(pc_reli & 0b11 == 0);
2330
2331                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2332                let imm14 = ((pc_reli >> 2) as i32 as u32) & 0x3fff;
2333                let insn = (insn & !0x0007ffe0) | (imm14 << 5);
2334                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2335            }
2336
2337            Self::A64Branch19 | Self::A64Ldr19 => {
2338                debug_assert!(pc_reli & 0b11 == 0);
2339
2340                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2341                let imm19 = ((pc_reli >> 2) as i32 as u32) & 0x7ffff;
2342                let insn = (insn & !0x00ffffe0) | (imm19 << 5);
2343                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2344            }
2345
2346            Self::A64Branch26 => {
2347                debug_assert!(pc_reli & 0b11 == 0);
2348
2349                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2350                let imm26 = ((pc_reli >> 2) as i32 as u32) & 0x03ff_ffff;
2351                let insn = (insn & !0x03ff_ffff) | imm26;
2352                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2353            }
2354
2355            Self::A64Adr21 => {
2356                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2357                let imm21 = (pc_reli as i32 as u32) & 0x1f_ffff;
2358                let immlo = imm21 & 0x3;
2359                let immhi = (imm21 >> 2) & 0x7ffff;
2360                let insn = (insn & !0x60ff_ffe0) | (immlo << 29) | (immhi << 5);
2361                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2362            }
2363
2364            Self::A64Adrp21 => {
2365                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2366
2367                // 1. Calculate the page-aligned PC and Target
2368                let pc_page = (use_offset as i64) & !0xFFF;
2369                let target_page = ((label_offset as i64) + addend) & !0xFFF;
2370
2371                // 2. Calculate the offset in pages
2372                let page_offset = (target_page - pc_page) >> 12;
2373
2374                // 3. Encode the 21-bit signed immediate
2375                let imm21 = (page_offset as u32) & 0x1F_FFFF;
2376                let immlo = imm21 & 0x3; // Lowest 2 bits
2377                let immhi = (imm21 >> 2) & 0x7FFFF; // Upper 19 bits
2378
2379                // 4. Clear existing immediate bits and insert new ones
2380                // Bits 29..31 (immlo) and Bits 5..24 (immhi)
2381                let insn = (insn & !0x60FF_FFE0) | (immlo << 29) | (immhi << 5);
2382
2383                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2384            }
2385
2386            Self::A64AddAbsLo12 => {
2387                let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2388
2389                let imm12 = ((pc_reli as i32 as u32) & 0xfff) << 10;
2390                let insn = insn | imm12;
2391                buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2392            }
2393
2394            _ => todo!(),
2395        }
2396    }
2397}
2398
2399pub const fn is_imm12(val: i64) -> bool {
2400    val >= -2048 && val <= 2047
2401}
2402
2403#[allow(dead_code)]
2404pub(crate) fn generate_imm(value: u64) -> (u32, u32) {
2405    #[cfg(not(feature = "riscv"))]
2406    {
2407        let _ = value;
2408        panic!("Can't generate RISC-V immediates without the `riscv` feature");
2409    }
2410    #[cfg(feature = "riscv")]
2411    {
2412        if is_imm12(value as _) {
2413            return (
2414                0,
2415                riscv::opcodes::InstructionValue::new(0)
2416                    .set_imm12(value as i64 as i32)
2417                    .value,
2418            );
2419        }
2420
2421        let value = value as i64;
2422
2423        let mod_num = 4096i64;
2424        let (imm20, imm12) = if value > 0 {
2425            let mut imm20 = value / mod_num;
2426            let mut imm12 = value % mod_num;
2427
2428            if imm12 >= 2048 {
2429                imm12 -= mod_num;
2430                imm20 += 1;
2431            }
2432
2433            (imm20, imm12)
2434        } else {
2435            let value_abs = value.abs();
2436            let imm20 = value_abs / mod_num;
2437            let imm12 = value_abs % mod_num;
2438            let mut imm20 = -imm20;
2439            let mut imm12 = -imm12;
2440            if imm12 < -2048 {
2441                imm12 += mod_num;
2442                imm20 -= 1;
2443            }
2444            (imm20, imm12)
2445        };
2446        (
2447            riscv::opcodes::InstructionValue::new(0)
2448                .set_imm20(imm20 as _)
2449                .value,
2450            riscv::opcodes::InstructionValue::new(0)
2451                .set_imm12(imm12 as _)
2452                .value,
2453        )
2454    }
2455}
2456
2457pub(crate) fn relocation_patch_size(kind: Reloc) -> Result<usize, AsmError> {
2458    match kind {
2459        Reloc::Abs4
2460        | Reloc::X86PCRel4
2461        | Reloc::X86CallPCRel4
2462        | Reloc::X86CallPLTRel4
2463        | Reloc::X86GOTPCRel4
2464        | Reloc::RiscvGotHi20
2465        | Reloc::RiscvPCRelLo12I
2466        | Reloc::Aarch64AdrPrelPgHi21
2467        | Reloc::Aarch64AddAbsLo12Nc
2468        | Reloc::Aarch64AdrGotPage21
2469        | Reloc::Aarch64Ld64GotLo12Nc => Ok(4),
2470        Reloc::Abs8 | Reloc::RiscvAbs8 | Reloc::RiscvCallPlt => Ok(8),
2471        _ => Err(AsmError::InvalidArgument),
2472    }
2473}
2474
2475fn checked_address(base: usize, addend: i64) -> Result<usize, AsmError> {
2476    if addend >= 0 {
2477        base.checked_add(usize::try_from(addend).map_err(|_| AsmError::TooLarge)?)
2478    } else {
2479        base.checked_sub(usize::try_from(addend.unsigned_abs()).map_err(|_| AsmError::TooLarge)?)
2480    }
2481    .ok_or(AsmError::TooLarge)
2482}
2483
2484fn checked_pcrel(target: usize, instruction: usize) -> Result<i64, AsmError> {
2485    i64::try_from(target as i128 - instruction as i128).map_err(|_| AsmError::TooLarge)
2486}
2487
2488/// Applies relocations to one writable code span.
2489///
2490/// Unsupported relocation kinds, invalid resolver results, out-of-range patch
2491/// spans, address overflow, and displacement overflow are returned as errors.
2492///
2493/// # Safety
2494///
2495/// `code` and `code_rx` must refer to writable and executable views of the same
2496/// `code_size`-byte allocation. Resolver pointers are treated as addresses and
2497/// are never dereferenced.
2498pub unsafe fn perform_relocations(
2499    code: *mut u8,
2500    code_rx: *const u8,
2501    code_size: usize,
2502    relocs: &[AsmReloc],
2503    get_address: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2504    get_got_entry: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2505    get_plt_entry: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2506) -> Result<(), AsmError> {
2507    use core::ptr::write_unaligned;
2508
2509    if code.is_null() || code_rx.is_null() {
2510        return Err(AsmError::InvalidArgument);
2511    }
2512
2513    for &AsmReloc {
2514        addend,
2515        kind,
2516        offset,
2517        ref target,
2518    } in relocs
2519    {
2520        let patch_size = relocation_patch_size(kind)?;
2521        let patch_start = offset as usize;
2522        let patch_end = patch_start
2523            .checked_add(patch_size)
2524            .ok_or(AsmError::TooLarge)?;
2525        if patch_end > code_size {
2526            return Err(AsmError::InvalidArgument);
2527        }
2528        let at = unsafe { code.add(patch_start) };
2529        let atrx = code_rx
2530            .addr()
2531            .checked_add(patch_start)
2532            .ok_or(AsmError::TooLarge)?;
2533        let resolve = |resolver: &dyn Fn(&RelocTarget) -> Result<*const u8, AsmError>| {
2534            let base = resolver(target)?;
2535            if base.is_null() {
2536                return Err(AsmError::InvalidArgument);
2537            }
2538            checked_address(base.addr(), addend)
2539        };
2540
2541        match kind {
2542            Reloc::Abs4 => {
2543                let what = resolve(&get_address)?;
2544                let what = u32::try_from(what).map_err(|_| AsmError::TooLarge)?;
2545                unsafe {
2546                    write_unaligned(at as *mut u32, what);
2547                }
2548            }
2549
2550            Reloc::Abs8 | Reloc::RiscvAbs8 => {
2551                let what = resolve(&get_address)?;
2552                let what = u64::try_from(what).map_err(|_| AsmError::TooLarge)?;
2553                unsafe {
2554                    write_unaligned(at as *mut u64, what);
2555                }
2556            }
2557
2558            Reloc::X86PCRel4 | Reloc::X86CallPCRel4 => {
2559                let what = resolve(&get_address)?;
2560                let pcrel =
2561                    i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2562
2563                unsafe {
2564                    write_unaligned(at as *mut i32, pcrel);
2565                }
2566            }
2567
2568            Reloc::X86GOTPCRel4 => {
2569                let what = resolve(&get_got_entry)?;
2570                let pcrel =
2571                    i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2572
2573                unsafe {
2574                    write_unaligned(at as *mut i32, pcrel);
2575                }
2576            }
2577
2578            Reloc::X86CallPLTRel4 => {
2579                let what = resolve(&get_plt_entry)?;
2580                let pcrel =
2581                    i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2582                unsafe { write_unaligned(at as *mut i32, pcrel) };
2583            }
2584
2585            Reloc::RiscvGotHi20 => {
2586                let what = resolve(&get_got_entry)?;
2587                let pc_rel =
2588                    i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2589                unsafe {
2590                    let buffer = core::slice::from_raw_parts_mut(at, 4);
2591                    let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2592                    let hi20 = (pc_rel as u32).wrapping_add(0x800) >> 12;
2593                    let insn = (insn & 0xfff) | (hi20 << 12);
2594                    buffer.copy_from_slice(&insn.to_le_bytes());
2595                }
2596            }
2597
2598            Reloc::RiscvPCRelLo12I => {
2599                let what = resolve(&get_got_entry)?;
2600                let pc_rel = checked_pcrel(what, atrx)?;
2601                let pc_rel = i32::try_from(pc_rel).map_err(|_| AsmError::TooLarge)?;
2602
2603                unsafe {
2604                    let buffer = core::slice::from_raw_parts_mut(at, 4);
2605                    let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2606                    let lo12 = ((pc_rel as i64 + 4) as u32) & 0xfff;
2607                    let insn = (insn & 0xFFFFF) | (lo12 << 20);
2608                    buffer.copy_from_slice(&insn.to_le_bytes());
2609                }
2610            }
2611
2612            Reloc::RiscvCallPlt => {
2613                #[cfg(not(feature = "riscv"))]
2614                {
2615                    return Err(AsmError::InvalidArgument);
2616                }
2617                #[cfg(feature = "riscv")]
2618                {
2619                    // A R_RISCV_CALL_PLT relocation expects auipc+jalr instruction pair.
2620                    // It is the equivalent of two relocations:
2621                    // 1. R_RISCV_PCREL_HI20 on the `auipc`
2622                    // 2. R_RISCV_PCREL_LO12_I on the `jalr`
2623
2624                    let what = resolve(&get_address)?;
2625                    let pcrel = i32::try_from(checked_pcrel(what, atrx)?)
2626                        .map_err(|_| AsmError::TooLarge)?;
2627
2628                    // See https://github.com/riscv-non-isa/riscv-elf-psabi-doc/blob/master/riscv-elf.adoc#pc-relative-symbol-addresses
2629                    // for a better explanation of the following code.
2630                    //
2631                    // Unlike the regular symbol relocations, here both "sub-relocations" point to the same address.
2632                    //
2633                    // `pcrel` is a signed value (+/- 2GiB range), when splitting it into two parts, we need to
2634                    // ensure that `hi20` is close enough to `pcrel` to be able to add `lo12` to it and still
2635                    // get a valid address.
2636                    //
2637                    // `lo12` is also a signed offset (+/- 2KiB range) relative to the `hi20` value.
2638                    //
2639                    // `hi20` should also be shifted right to be the "true" value. But we also need it
2640                    // left shifted for the `lo12` calculation and it also matches the instruction encoding.
2641                    let hi20 = pcrel.wrapping_add(0x800) as u32 & 0xFFFFF000u32;
2642                    let lo12 = (pcrel as u32).wrapping_sub(hi20) & 0xFFF;
2643
2644                    unsafe {
2645                        let auipc_addr = at as *mut u32;
2646                        let auipc = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2647                            .encode()
2648                            .set_imm20(hi20 as _)
2649                            .value;
2650                        write_unaligned(auipc_addr, auipc_addr.read_unaligned() | auipc);
2651
2652                        let jalr_addr = at.add(4) as *mut u32;
2653                        let jalr = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2654                            .encode()
2655                            .set_imm12(lo12 as _)
2656                            .value;
2657                        write_unaligned(jalr_addr, jalr_addr.read_unaligned() | jalr);
2658                    }
2659                }
2660            }
2661
2662            Reloc::Aarch64AdrPrelPgHi21 => {
2663                let what = resolve(&get_address)?;
2664                let pages = ((what & !0xfff) as i128 - (atrx & !0xfff) as i128) >> 12;
2665                if !(-(1i128 << 20)..(1i128 << 20)).contains(&pages) {
2666                    return Err(AsmError::TooLarge);
2667                }
2668                let iptr = at as *mut u32;
2669                let imm21 = pages as u32 & 0x1f_ffff;
2670                let lo = (imm21 & 0x3) << 29;
2671                let hi = ((imm21 >> 2) & 0x7ffff) << 5;
2672                unsafe {
2673                    let insn = iptr.read_unaligned();
2674                    write_unaligned(iptr, insn | lo | hi);
2675                }
2676            }
2677
2678            Reloc::Aarch64AddAbsLo12Nc => {
2679                let what = resolve(&get_address)?;
2680                let iptr = at as *mut u32;
2681                let imm12 = (what as u32 & 0xfff) << 10;
2682                unsafe {
2683                    let insn = iptr.read_unaligned();
2684                    write_unaligned(iptr, insn | imm12);
2685                }
2686            }
2687
2688            Reloc::Aarch64AdrGotPage21 => {
2689                let what = resolve(&get_got_entry)?;
2690                let pages = ((what & !0xfff) as i128 - (atrx & !0xfff) as i128) >> 12;
2691                if !(-(1i128 << 20)..(1i128 << 20)).contains(&pages) {
2692                    return Err(AsmError::TooLarge);
2693                }
2694                let iptr = at as *mut u32;
2695                let imm21 = pages as u32 & 0x1f_ffff;
2696                let lo = (imm21 & 0x3) << 29;
2697                let hi = ((imm21 >> 2) & 0x7ffff) << 5;
2698                unsafe {
2699                    let insn = iptr.read_unaligned();
2700                    write_unaligned(iptr, insn | lo | hi);
2701                }
2702            }
2703
2704            Reloc::Aarch64Ld64GotLo12Nc => {
2705                let what = resolve(&get_got_entry)?;
2706                if what & 7 != 0 {
2707                    return Err(AsmError::InvalidArgument);
2708                }
2709                let iptr = at as *mut u32;
2710                let imm12 = ((what as u32 & 0xfff) >> 3) << 10;
2711                unsafe {
2712                    let insn = iptr.read_unaligned();
2713                    write_unaligned(iptr, insn | imm12);
2714                }
2715            }
2716
2717            _ => return Err(AsmError::InvalidArgument),
2718        }
2719    }
2720    Ok(())
2721}
2722
2723#[cfg(test)]
2724mod tests {
2725    use super::*;
2726
2727    fn resolved(address: usize) -> impl Fn(&RelocTarget) -> Result<*const u8, AsmError> {
2728        move |_| Ok(address as *const u8)
2729    }
2730
2731    fn unresolved(_: &RelocTarget) -> Result<*const u8, AsmError> {
2732        Ok(core::ptr::null())
2733    }
2734
2735    #[test]
2736    fn relocation_rejects_patch_outside_code() {
2737        let mut code = [0u8; 4];
2738        let relocs = [AsmReloc {
2739            offset: 1,
2740            kind: Reloc::Abs4,
2741            addend: 0,
2742            target: RelocTarget::Label(Label::from_id(0)),
2743        }];
2744
2745        let result = unsafe {
2746            perform_relocations(
2747                code.as_mut_ptr(),
2748                code.as_ptr(),
2749                code.len(),
2750                &relocs,
2751                resolved(1),
2752                resolved(1),
2753                resolved(1),
2754            )
2755        };
2756
2757        assert_eq!(result, Err(AsmError::InvalidArgument));
2758        assert_eq!(code, [0; 4]);
2759    }
2760
2761    #[test]
2762    fn relocation_rejects_null_and_overflowing_targets() {
2763        let mut code = [0u8; 8];
2764        let reloc = AsmReloc {
2765            offset: 0,
2766            kind: Reloc::Abs8,
2767            addend: 0,
2768            target: RelocTarget::Label(Label::from_id(0)),
2769        };
2770
2771        let null_result = unsafe {
2772            perform_relocations(
2773                code.as_mut_ptr(),
2774                code.as_ptr(),
2775                code.len(),
2776                core::slice::from_ref(&reloc),
2777                unresolved,
2778                unresolved,
2779                unresolved,
2780            )
2781        };
2782        assert_eq!(null_result, Err(AsmError::InvalidArgument));
2783
2784        let overflowing = AsmReloc {
2785            addend: -2,
2786            ..reloc
2787        };
2788        let overflow_result = unsafe {
2789            perform_relocations(
2790                code.as_mut_ptr(),
2791                code.as_ptr(),
2792                code.len(),
2793                core::slice::from_ref(&overflowing),
2794                resolved(1),
2795                resolved(1),
2796                resolved(1),
2797            )
2798        };
2799        assert_eq!(overflow_result, Err(AsmError::TooLarge));
2800    }
2801
2802    #[test]
2803    fn poisoned_buffer_rejects_raw_mutation_and_patch_metadata() {
2804        let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2805        buffer.write_u32(0);
2806        let bytes = buffer.data().to_vec();
2807        buffer.record_error(AsmError::InvalidOperand);
2808
2809        buffer.write_u8(1);
2810        buffer.put8(2);
2811        assert!(
2812            !buffer
2813                .add_symbol(ExternalName::user(0, 1), RelocDistance::Far)
2814                .is_valid()
2815        );
2816        assert_eq!(buffer.add_constant(3u64), Constant(u32::MAX));
2817        assert!(!buffer.get_label().is_valid());
2818        buffer.add_reloc(Reloc::Abs4, RelocTarget::Label(Label::from_id(0)), 0);
2819        let patch = buffer.try_record_patch_site(0, LabelUse::X86JmpRel32, 0);
2820
2821        assert_eq!(patch, Err(AsmError::InvalidOperand));
2822        assert_eq!(buffer.data(), bytes);
2823        assert!(buffer.relocs().is_empty());
2824        assert!(buffer.patch_sites.is_empty());
2825        assert!(matches!(buffer.finish(), Err(AsmError::InvalidOperand)));
2826    }
2827
2828    #[test]
2829    fn target_rejects_foreign_relocations_and_patch_kinds() {
2830        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2831        let label = buffer.get_label();
2832        buffer.add_reloc(Reloc::X86PCRel4, RelocTarget::Label(label), 0);
2833        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2834        assert!(buffer.relocs().is_empty());
2835
2836        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2837        buffer.write_u32(0);
2838        assert_eq!(
2839            buffer.try_record_patch_site(0, LabelUse::X86JmpRel32, 0),
2840            Err(AsmError::InvalidArch)
2841        );
2842        assert!(buffer.patch_sites.is_empty());
2843
2844        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2845        let label = buffer.get_label();
2846        buffer.use_label_at_offset(0, label, LabelUse::RVJal20);
2847        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2848        assert!(buffer.pending_fixup_records.is_empty());
2849
2850        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2851        let label = buffer.get_label();
2852        assert_eq!(
2853            buffer.emit_veneer(label, 0, LabelUse::RVJal20),
2854            Err(AsmError::InvalidArch)
2855        );
2856        assert!(buffer.data().is_empty());
2857
2858        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2859        buffer.write_u32(0);
2860        assert_eq!(
2861            buffer.try_record_patch_site(0, LabelUse::A64Ldr12, 0),
2862            Err(AsmError::UnsupportedInstruction {
2863                reason: "AArch64 LDR12 label patching is not implemented",
2864            })
2865        );
2866    }
2867
2868    #[cfg(not(feature = "x86"))]
2869    #[test]
2870    fn disabled_x86_label_use_is_rejected_without_a_fixup() {
2871        let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2872        let label = buffer.get_label();
2873        buffer.write_u32(0);
2874        let bytes = buffer.data().to_vec();
2875
2876        buffer.use_label_at_offset(0, label, LabelUse::X86JmpRel32);
2877
2878        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2879        assert_eq!(buffer.data(), bytes);
2880        assert!(buffer.pending_fixup_records.is_empty());
2881        assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2882
2883        let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2884        let label = buffer.get_label();
2885        buffer.add_reloc(Reloc::X86PCRel4, RelocTarget::Label(label), 0);
2886        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2887        assert!(buffer.relocs().is_empty());
2888    }
2889
2890    #[cfg(not(feature = "riscv"))]
2891    #[test]
2892    fn disabled_riscv_label_use_is_rejected_without_a_fixup() {
2893        let mut buffer = CodeBuffer::new(Environment::new(Arch::RISCV64));
2894        let label = buffer.get_label();
2895        buffer.write_u32(0);
2896        let bytes = buffer.data().to_vec();
2897
2898        buffer.use_label_at_offset(0, label, LabelUse::RVPCRel32);
2899
2900        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2901        assert_eq!(buffer.data(), bytes);
2902        assert!(buffer.pending_fixup_records.is_empty());
2903        assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2904
2905        let mut buffer = CodeBuffer::new(Environment::new(Arch::RISCV64));
2906        let label = buffer.get_label();
2907        buffer.add_reloc(Reloc::RiscvCallPlt, RelocTarget::Label(label), 0);
2908        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2909        assert!(buffer.relocs().is_empty());
2910    }
2911
2912    #[cfg(not(feature = "aarch64"))]
2913    #[test]
2914    fn disabled_aarch64_label_use_is_rejected_without_a_fixup() {
2915        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2916        let label = buffer.get_label();
2917        buffer.write_u32(0);
2918        let bytes = buffer.data().to_vec();
2919
2920        buffer.use_label_at_offset(0, label, LabelUse::A64Branch26);
2921
2922        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2923        assert_eq!(buffer.data(), bytes);
2924        assert!(buffer.pending_fixup_records.is_empty());
2925        assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2926
2927        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2928        let label = buffer.get_label();
2929        buffer.add_reloc(Reloc::Arm64Call, RelocTarget::Label(label), 0);
2930        assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2931        assert!(buffer.relocs().is_empty());
2932    }
2933
2934    #[test]
2935    fn poisoned_buffer_does_not_consume_pending_island_state() {
2936        let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2937        buffer.add_constant(7u64);
2938        let pending_constants = buffer.pending_constants.len();
2939        buffer.record_error(AsmError::InvalidOperand);
2940
2941        assert_eq!(buffer.emit_island(0), Err(AsmError::InvalidOperand));
2942        assert_eq!(buffer.pending_constants.len(), pending_constants);
2943        assert!(buffer.data().is_empty());
2944    }
2945
2946    #[cfg(feature = "jit")]
2947    #[test]
2948    fn allocate_resolved_rejects_unresolved_symbols() {
2949        let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2950        let symbol = buffer.add_symbol(ExternalName::user(0, 7), RelocDistance::Far);
2951        buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(symbol), 0);
2952        buffer.write_u64(0);
2953        let code = buffer.finish().unwrap();
2954        let mut allocator = JitAllocator::new(Default::default());
2955
2956        let result = code.allocate_resolved(&mut allocator, |_| core::ptr::null());
2957
2958        assert_eq!(result.err(), Some(AsmError::InvalidArgument));
2959    }
2960
2961    #[cfg(feature = "jit")]
2962    #[test]
2963    fn allocate_resolved_resolves_user_external_names() {
2964        let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2965        let symbol = buffer.extern_user(1, 2, RelocDistance::Far);
2966        buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(symbol), 0);
2967        buffer.write_u64(0);
2968        let code = buffer.finish().unwrap();
2969        let mut allocator = JitAllocator::new(Default::default());
2970
2971        // Any non-null address is accepted; Abs8 just patches the pointer.
2972        let target = 0x1000usize as *const u8;
2973        let loaded = code
2974            .allocate_resolved(&mut allocator, |name| match name {
2975                ExternalName::User(u) if u.namespace == 1 && u.index == 2 => target,
2976                _ => core::ptr::null(),
2977            })
2978            .unwrap();
2979
2980        unsafe {
2981            let patched = core::ptr::read_unaligned(loaded.rx() as *const usize);
2982            assert_eq!(patched, target as usize);
2983        }
2984    }
2985
2986    #[test]
2987    fn extern_sym_deduplicates_by_name() {
2988        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
2989        let first = buf.extern_sym("puts", RelocDistance::Far);
2990        let other = buf.extern_sym("printf", RelocDistance::Near);
2991        let again = buf.extern_sym("puts", RelocDistance::Near);
2992
2993        assert_eq!(first, again);
2994        assert_ne!(first, other);
2995        // The first declaration's distance wins.
2996        assert_eq!(buf.symbol_distance(first), Some(RelocDistance::Far));
2997        assert_eq!(buf.symbol_distance(other), Some(RelocDistance::Near));
2998    }
2999
3000    #[test]
3001    fn extern_user_deduplicates_by_namespace_and_index() {
3002        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3003        let first = buf.extern_user(0, 1, RelocDistance::Far);
3004        let other_ns = buf.extern_user(1, 1, RelocDistance::Near);
3005        let other_idx = buf.extern_user(0, 2, RelocDistance::Near);
3006        let again = buf.extern_user(0, 1, RelocDistance::Near);
3007
3008        assert_eq!(first, again);
3009        assert_ne!(first, other_ns);
3010        assert_ne!(first, other_idx);
3011        assert_eq!(buf.symbol_distance(first), Some(RelocDistance::Far));
3012        assert_eq!(
3013            buf.symbol_name(first),
3014            Some(&ExternalName::user(0, 1))
3015        );
3016    }
3017
3018    #[test]
3019    fn unknown_symbols_are_fallible() {
3020        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3021        let missing = Sym::from_id(u32::MAX);
3022
3023        assert_eq!(buf.symbol_name(missing), None);
3024        assert_eq!(buf.symbol_distance(missing), None);
3025
3026        let finalized = buf.finish().unwrap();
3027        assert_eq!(finalized.symbol_name(missing), None);
3028        assert_eq!(finalized.symbol_distance(missing), None);
3029    }
3030
3031    #[test]
3032    fn defined_symbols_are_resolved_at_finish() {
3033        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3034        buf.write_u8(0x90);
3035        let entry = buf.get_label();
3036        buf.bind_label(entry);
3037        buf.bind_symbol("entry", entry);
3038        buf.write_u8(0xC3);
3039
3040        let result = buf.finish().unwrap();
3041        assert_eq!(result.defined_symbol_str("entry"), Some(1));
3042        assert_eq!(result.defined_symbol_str("missing"), None);
3043        assert_eq!(
3044            result.defined_symbol_offset(&ExternalName::from("entry")),
3045            Some(1)
3046        );
3047    }
3048
3049    #[test]
3050    fn invalid_label_binding_is_reported_without_mutating_labels() {
3051        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3052        let invalid = Label::from_id(0);
3053
3054        assert_eq!(buf.try_bind_label(invalid), Err(AsmError::InvalidArgument));
3055        assert_eq!(buf.label_count(), 0);
3056
3057        buf.bind_label(invalid);
3058        assert_eq!(buf.error(), Some(&AsmError::InvalidArgument));
3059        assert!(matches!(buf.finish(), Err(AsmError::InvalidArgument)));
3060    }
3061
3062    #[cfg(feature = "x86")]
3063    #[test]
3064    fn invalid_patch_registration_does_not_create_metadata() {
3065        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3066
3067        assert_eq!(
3068            buf.try_record_patch_site(0, LabelUse::X86JmpRel32, 0),
3069            Err(AsmError::InvalidArgument)
3070        );
3071        assert!(buf.patch_sites.is_empty());
3072
3073        buf.record_patch_site(0, LabelUse::X86JmpRel32, 0);
3074        assert_eq!(buf.error(), Some(&AsmError::InvalidArgument));
3075        assert!(buf.patch_sites.is_empty());
3076    }
3077
3078    #[cfg(feature = "x86")]
3079    #[test]
3080    fn finish_rejects_an_unbound_fixup() {
3081        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3082        let label = buf.get_label();
3083        buf.write_u32(0);
3084        buf.use_label_at_offset(0, label, LabelUse::X86JmpRel32);
3085
3086        assert!(matches!(buf.finish(), Err(AsmError::UnboundLabel)));
3087    }
3088
3089    #[cfg(feature = "x86")]
3090    #[test]
3091    fn unsupported_veneer_leaves_buffer_unchanged() {
3092        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3093        let label = buf.get_label();
3094        buf.write_u32(0);
3095        let original = buf.data().to_vec();
3096
3097        assert_eq!(
3098            buf.emit_veneer(label, 0, LabelUse::X86JmpRel32),
3099            Err(AsmError::UnsupportedInstruction {
3100                reason: "branch range requires an unsupported veneer",
3101            })
3102        );
3103        assert_eq!(buf.data(), original);
3104    }
3105
3106    #[test]
3107    fn branch_ranges_include_the_last_encodable_offset() {
3108        for (kind, range) in [
3109            (LabelUse::A64Branch14, 1 << 15),
3110            (LabelUse::A64Branch19, 1 << 20),
3111            (LabelUse::A64Branch26, 1 << 27),
3112        ] {
3113            assert!(kind.can_reach(0, range - 4));
3114            assert!(!kind.can_reach(0, range));
3115            assert!(kind.can_reach(range, 0));
3116            assert!(!kind.can_reach(range + 4, 0));
3117        }
3118
3119        assert!(LabelUse::RVB12.can_reach(0, (1 << 12) - 2));
3120        assert!(!LabelUse::RVB12.can_reach(0, 1 << 12));
3121        assert!(LabelUse::RVB12.can_reach(1 << 12, 0));
3122        assert!(!LabelUse::RVB12.can_reach((1 << 12) + 2, 0));
3123    }
3124
3125    #[test]
3126    fn supported_veneer_ranges_cover_both_exact_boundaries() {
3127        for (kind, positive, negative, step) in [
3128            (LabelUse::RVJal20, (1 << 20) - 2, 1 << 20, 2),
3129            (LabelUse::RVB12, (1 << 12) - 2, 1 << 12, 2),
3130            (LabelUse::RVCJump, (1 << 11) - 2, 1 << 11, 2),
3131        ] {
3132            assert!(kind.supports_veneer());
3133            assert!(kind.can_reach(0, positive));
3134            assert!(!kind.can_reach(0, positive + step));
3135            assert!(kind.can_reach(negative, 0));
3136            assert!(!kind.can_reach(negative + step, 0));
3137        }
3138    }
3139
3140    #[cfg(feature = "riscv")]
3141    #[test]
3142    fn riscv_veneer_is_finalized_as_a_fixup() {
3143        let mut buf = CodeBuffer::new(Environment::new(Arch::RISCV64));
3144        let label = buf.get_label();
3145        buf.bind_label(label);
3146        buf.write_u32(0);
3147
3148        buf.emit_veneer(label, 0, LabelUse::RVB12).unwrap();
3149        let code = buf.finish().unwrap();
3150        assert_eq!(code.data().len(), 12);
3151        // The source branch reaches the veneer inserted directly after it.
3152        assert_eq!(&code.data()[0..4], &[0x00, 0x02, 0x00, 0x00]);
3153    }
3154
3155    #[cfg(feature = "riscv")]
3156    #[test]
3157    fn riscv_supported_veneers_are_inserted_by_islands() {
3158        for kind in [LabelUse::RVJal20, LabelUse::RVB12, LabelUse::RVCJump] {
3159            let mut buf = CodeBuffer::new(Environment::new(Arch::RISCV64));
3160            let label = buf.get_label();
3161            let patch_size = kind.patch_size();
3162            buf.get_appended_space(patch_size);
3163            buf.use_label_at_offset(0, label, kind);
3164
3165            buf.emit_island(u32::MAX).unwrap();
3166            buf.bind_label(label);
3167            let code = buf.finish().unwrap();
3168
3169            assert_eq!(code.data().len(), 12);
3170            assert!(kind.can_reach(0, 4));
3171        }
3172    }
3173
3174    #[cfg(feature = "riscv")]
3175    #[test]
3176    fn riscv_large_images_hit_exact_branch_boundaries() {
3177        for (kind, positive, step) in [
3178            (LabelUse::RVJal20, (1 << 20) - 2, 2),
3179            (LabelUse::RVB12, (1 << 12) - 2, 2),
3180            (LabelUse::RVCJump, (1 << 11) - 2, 2),
3181        ] {
3182            let mut exact = CodeBuffer::new(Environment::new(Arch::RISCV64));
3183            let label = exact.get_label();
3184            exact.get_appended_space(kind.patch_size());
3185            exact.use_label_at_offset(0, label, kind);
3186            exact.get_appended_space(positive as usize - kind.patch_size());
3187            exact.bind_label(label);
3188            assert!(exact.finish().is_ok());
3189
3190            let mut outside = CodeBuffer::new(Environment::new(Arch::RISCV64));
3191            let label = outside.get_label();
3192            outside.get_appended_space(kind.patch_size());
3193            outside.use_label_at_offset(0, label, kind);
3194            outside.get_appended_space((positive + step) as usize - kind.patch_size());
3195            outside.bind_label(label);
3196            assert_eq!(outside.finish().err(), Some(AsmError::TooLarge));
3197        }
3198    }
3199
3200    #[cfg(feature = "aarch64")]
3201    #[test]
3202    fn out_of_range_aarch64_branch_reports_unsupported_veneer() {
3203        let mut buf = CodeBuffer::new(Environment::new(Arch::AArch64));
3204        let label = buf.get_label();
3205        buf.write_u32(0);
3206        buf.use_label_at_offset(0, label, LabelUse::A64Branch14);
3207        buf.get_appended_space((1 << 15) - 4);
3208        buf.bind_label(label);
3209
3210        assert_eq!(
3211            buf.finish().err(),
3212            Some(AsmError::UnsupportedInstruction {
3213                reason: "AArch64 branch veneers are not implemented",
3214            })
3215        );
3216    }
3217
3218    #[cfg(feature = "x86")]
3219    #[test]
3220    fn x86_branch_relaxation_honors_rel8_boundaries() {
3221        use crate::x86::{Assembler, JEmitter, JmpEmitter};
3222
3223        let forward = |padding: usize, conditional: bool| {
3224            let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3225            let target = buf.get_label();
3226            {
3227                let mut asm = Assembler::new(&mut buf);
3228                if conditional {
3229                    asm.jz(target);
3230                } else {
3231                    asm.jmp(target);
3232                }
3233            }
3234            for _ in 0..padding {
3235                buf.write_u8(0x90);
3236            }
3237            buf.bind_label(target);
3238            buf.finish().unwrap().data().to_vec()
3239        };
3240        for conditional in [false, true] {
3241            let short = forward(127, conditional);
3242            assert_eq!(short[0], if conditional { 0x74 } else { 0xEB });
3243            assert_eq!(short[1], 127);
3244
3245            let near = forward(128, conditional);
3246            if conditional {
3247                assert_eq!(&near[..2], &[0x0F, 0x84]);
3248            } else {
3249                assert_eq!(near[0], 0xE9);
3250            }
3251        }
3252
3253        let backward = |padding: usize, conditional: bool| {
3254            let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3255            let target = buf.get_label();
3256            buf.bind_label(target);
3257            for _ in 0..padding {
3258                buf.write_u8(0x90);
3259            }
3260            {
3261                let mut asm = Assembler::new(&mut buf);
3262                if conditional {
3263                    asm.jz(target);
3264                } else {
3265                    asm.jmp(target);
3266                }
3267            }
3268            buf.finish().unwrap().data().to_vec()
3269        };
3270        for conditional in [false, true] {
3271            let short = backward(126, conditional);
3272            assert_eq!(
3273                &short[126..],
3274                &[if conditional { 0x74 } else { 0xEB }, 0x80]
3275            );
3276
3277            let near = backward(127, conditional);
3278            assert_eq!(near[127], if conditional { 0x0F } else { 0xE9 });
3279        }
3280    }
3281
3282    #[cfg(feature = "x86")]
3283    #[test]
3284    fn x86_branch_relaxation_reaches_a_bounded_fixed_point() {
3285        use crate::x86::{Assembler, JmpEmitter};
3286
3287        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3288        let target = buf.get_label();
3289        {
3290            let mut asm = Assembler::new(&mut buf);
3291            asm.jmp(target);
3292        }
3293        for _ in 0..120 {
3294            buf.write_u8(0x90);
3295        }
3296        {
3297            let mut asm = Assembler::new(&mut buf);
3298            asm.jmp(target);
3299        }
3300        for _ in 0..3 {
3301            buf.write_u8(0x90);
3302        }
3303        buf.bind_label(target);
3304
3305        let code = buf.finish().unwrap();
3306        assert_eq!(&code.data()[..2], &[0xEB, 125]);
3307        assert_eq!(&code.data()[122..124], &[0xEB, 3]);
3308        assert_eq!(code.label_offsets[target.id() as usize], 127);
3309    }
3310
3311    #[cfg(feature = "x86")]
3312    #[test]
3313    fn x86_branch_relaxation_rebases_metadata_and_is_deterministic() {
3314        use crate::x86::{Assembler, JmpEmitter};
3315
3316        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3317        let target = buf.get_label();
3318        {
3319            let mut asm = Assembler::new(&mut buf);
3320            asm.jmp(target);
3321        }
3322        buf.write_u8(0x90);
3323        buf.bind_label(target);
3324        buf.bind_symbol("target", target);
3325
3326        buf.add_reloc(Reloc::Abs4, RelocTarget::Label(target), 0);
3327        buf.write_u32(0);
3328        let block = buf.reserve_patch_block(4, 1).unwrap();
3329        let block_offset = block.offset();
3330        let site = buf
3331            .try_record_patch_site(
3332                block_offset,
3333                LabelUse::X86JmpRel32,
3334                buf.label_offset(target),
3335            )
3336            .unwrap();
3337
3338        let later = buf.get_label();
3339        {
3340            let mut asm = Assembler::new(&mut buf);
3341            asm.long().jmp(later);
3342        }
3343        buf.write_u8(0x90);
3344        buf.bind_label(later);
3345
3346        let first = buf.finish().unwrap();
3347        assert_eq!(&first.data()[..3], &[0xEB, 1, 0x90]);
3348        assert_eq!(first.defined_symbol_str("target"), Some(3));
3349        assert_eq!(first.relocs()[0].offset, 3);
3350        assert_eq!(first.patch_catalog().block(PatchBlockId::from_index(0)).unwrap().offset, 7);
3351        let patch_site = first.patch_catalog().site(site).unwrap();
3352        assert_eq!(patch_site.offset, 7);
3353        assert_eq!(patch_site.current_target, 3);
3354        assert_eq!(&first.data()[11..], &[0xE9, 1, 0, 0, 0, 0x90]);
3355
3356        let second = buf.finish().unwrap();
3357        assert_eq!(second.data(), first.data());
3358        assert_eq!(second.relocs(), first.relocs());
3359        assert_eq!(second.patch_catalog(), first.patch_catalog());
3360    }
3361
3362    #[cfg(feature = "x86")]
3363    #[test]
3364    fn x86_branch_relaxation_preserves_recorded_alignment() {
3365        use crate::x86::{Assembler, JmpEmitter};
3366
3367        let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3368        let target = buf.get_label();
3369        {
3370            let mut asm = Assembler::new(&mut buf);
3371            asm.jmp(target);
3372        }
3373        buf.try_align_to(16).unwrap();
3374        buf.bind_label(target);
3375
3376        let code = buf.finish().unwrap();
3377        assert_eq!(code.data()[0], 0xE9);
3378        assert_eq!(code.label_offsets[target.id() as usize], 16);
3379    }
3380}