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
23pub 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#[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#[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 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#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
140pub enum ExternalName {
141 Symbol(Cow<'static, str>),
142 User(UserExternalName),
143}
144
145impl ExternalName {
146 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#[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#[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#[derive(Clone, Copy)]
228struct AsmConstant {
229 upcoming_label: Option<Label>,
235 align: CodeOffset,
237 size: usize,
239}
240
241pub 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#[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 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 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 #[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 #[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 #[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 RelocTarget::Label(_) => core::ptr::null(),
496 };
497
498 self.allocate_relocated(jit_allocator, by_name, by_name)
499 }
500}
501
502impl CodeBuffer {
503 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 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 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 #[cfg(feature = "x86")]
627 pub(crate) fn byte_at(&self, offset: CodeOffset) -> u8 {
628 self.data[offset as usize]
629 }
630
631 #[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 #[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 #[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 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 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 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 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 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 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 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 self.emit_veneer(label, offset, kind)?;
1195 }
1196 Ok(())
1197 }
1198
1199 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 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 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 }
1270
1271 fn worst_case_end_of_island(&self, distance: CodeOffset) -> CodeOffset {
1277 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 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 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 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 !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 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 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 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#[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
1804pub type CodeOffset = u32;
1809
1810pub type Addend = i64;
1812
1813#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
1815pub enum Reloc {
1816 Abs4,
1818 Abs8,
1820 X86PCRel4,
1822 X86CallPCRel4,
1824 X86CallPLTRel4,
1826 X86GOTPCRel4,
1828 X86SecRel,
1832 Arm32Call,
1834 Arm64Call,
1838
1839 ElfX86_64TlsGd,
1841
1842 MachOX86_64Tlv,
1844
1845 MachOAarch64TlsAdrPage21,
1848
1849 MachOAarch64TlsAdrPageOff12,
1852
1853 Aarch64TlsDescAdrPage21,
1856
1857 Aarch64TlsDescLd64Lo12,
1860
1861 Aarch64TlsDescAddLo12,
1864
1865 Aarch64TlsDescCall,
1868
1869 Aarch64AdrGotPage21,
1873
1874 Aarch64Ld64GotLo12Nc,
1879
1880 Aarch64AdrPrelPgHi21,
1882 Aarch64AddAbsLo12Nc,
1884
1885 RiscvAbs8,
1887
1888 RiscvCallPlt,
1896
1897 RiscvTlsGdHi20,
1902
1903 RiscvPCRelLo12I,
1908
1909 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 RVJal20,
1960 RVPCRel32,
1969
1970 RVB12,
1975
1976 RVPCRelHi20,
1979
1980 RVPCRelLo12I,
1987
1988 RVCJump,
1990 RVCB9,
1992 A64Branch14,
1995 A64Branch19,
1998 A64Branch26,
2001 A64Ldr19,
2004 A64Adr21,
2007 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 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 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 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 let pc_page = (use_offset as i64) & !0xFFF;
2369 let target_page = ((label_offset as i64) + addend) & !0xFFF;
2370
2371 let page_offset = (target_page - pc_page) >> 12;
2373
2374 let imm21 = (page_offset as u32) & 0x1F_FFFF;
2376 let immlo = imm21 & 0x3; let immhi = (imm21 >> 2) & 0x7FFFF; 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
2488pub 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 let what = resolve(&get_address)?;
2625 let pcrel = i32::try_from(checked_pcrel(what, atrx)?)
2626 .map_err(|_| AsmError::TooLarge)?;
2627
2628 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 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 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 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}