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, fill_with_nops,
10 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(
993 &mut self,
994 size: CodeOffset,
995 align: CodeOffset,
996 ) -> Result<PatchBlockId, AsmError> {
997 if let Some(error) = self.error.clone() {
998 return Err(error);
999 }
1000 let min_align = minimum_patch_alignment(self.env.arch());
1001 let align = align.max(min_align);
1002 if size == 0 || !align.is_power_of_two() {
1003 return Err(AsmError::InvalidArgument);
1004 }
1005 let nop_size = match self.env.arch() {
1006 Arch::X86 | Arch::X64 => 1,
1007 Arch::AArch64 | Arch::RISCV32 | Arch::RISCV64 => 4,
1008 _ => return Err(AsmError::InvalidArch),
1009 };
1010 if size as usize % nop_size != 0 {
1011 return Err(AsmError::InvalidArgument);
1012 }
1013
1014 self.try_align_to(align)?;
1015 let arch = self.env.arch();
1016 let offset = self.cur_offset();
1017 let block = self.get_appended_space(size as usize);
1018 fill_with_nops(arch, block)?;
1019
1020 let id = PatchBlockId::from_index(self.patch_blocks.len());
1021 self.patch_blocks.push(PendingPatchBlock {
1022 offset,
1023 size,
1024 align,
1025 });
1026 Ok(id)
1027 }
1028
1029 pub fn record_patch_block(
1030 &mut self,
1031 offset: CodeOffset,
1032 size: CodeOffset,
1033 align: CodeOffset,
1034 ) -> PatchBlockId {
1035 match self.try_record_patch_block(offset, size, align) {
1036 Ok(id) => id,
1037 Err(error) => {
1038 self.record_error(error);
1039 PatchBlockId::from_index(usize::MAX)
1040 }
1041 }
1042 }
1043
1044 pub fn try_record_patch_block(
1045 &mut self,
1046 offset: CodeOffset,
1047 size: CodeOffset,
1048 align: CodeOffset,
1049 ) -> Result<PatchBlockId, AsmError> {
1050 if let Some(error) = self.error.clone() {
1051 return Err(error);
1052 }
1053 if size == 0 || align == 0 || !align.is_power_of_two() || offset & (align - 1) != 0 {
1054 return Err(AsmError::InvalidArgument);
1055 }
1056 let end = (offset as usize)
1057 .checked_add(size as usize)
1058 .ok_or(AsmError::TooLarge)?;
1059 if end > self.data.len() {
1060 return Err(AsmError::InvalidArgument);
1061 }
1062 let id = PatchBlockId::from_index(self.patch_blocks.len());
1063 self.patch_blocks.push(PendingPatchBlock {
1064 offset,
1065 size,
1066 align,
1067 });
1068 Ok(id)
1069 }
1070
1071 pub fn record_patch_site(
1072 &mut self,
1073 offset: CodeOffset,
1074 kind: LabelUse,
1075 target_offset: CodeOffset,
1076 ) -> PatchSiteId {
1077 match self.try_record_patch_site(offset, kind, target_offset) {
1078 Ok(id) => id,
1079 Err(error) => {
1080 self.record_error(error);
1081 PatchSiteId::from_index(usize::MAX)
1082 }
1083 }
1084 }
1085
1086 pub fn try_record_patch_site(
1087 &mut self,
1088 offset: CodeOffset,
1089 kind: LabelUse,
1090 target_offset: CodeOffset,
1091 ) -> Result<PatchSiteId, AsmError> {
1092 if let Some(error) = self.error.clone() {
1093 return Err(error);
1094 }
1095 self.validate_patch_site_offset(offset, kind)?;
1096 let id = PatchSiteId::from_index(self.patch_sites.len());
1097 self.patch_sites.push(PendingPatchSite {
1098 offset,
1099 kind,
1100 target: PendingPatchTarget::Offset(target_offset),
1101 addend: 0,
1102 });
1103 Ok(id)
1104 }
1105
1106 pub fn record_label_patch_site(
1107 &mut self,
1108 offset: CodeOffset,
1109 label: Label,
1110 kind: LabelUse,
1111 ) -> PatchSiteId {
1112 match self.try_record_label_patch_site(offset, label, kind) {
1113 Ok(id) => id,
1114 Err(error) => {
1115 self.record_error(error);
1116 PatchSiteId::from_index(usize::MAX)
1117 }
1118 }
1119 }
1120
1121 pub fn try_record_label_patch_site(
1122 &mut self,
1123 offset: CodeOffset,
1124 label: Label,
1125 kind: LabelUse,
1126 ) -> Result<PatchSiteId, AsmError> {
1127 if let Some(error) = self.error.clone() {
1128 return Err(error);
1129 }
1130 if self.label_offsets.get(label.id() as usize).is_none() {
1131 return Err(AsmError::InvalidArgument);
1132 }
1133 self.validate_patch_site_offset(offset, kind)?;
1134 let id = PatchSiteId::from_index(self.patch_sites.len());
1135 self.patch_sites.push(PendingPatchSite {
1136 offset,
1137 kind,
1138 target: PendingPatchTarget::Label(label),
1139 addend: 0,
1140 });
1141 Ok(id)
1142 }
1143
1144 fn validate_patch_site_offset(
1145 &self,
1146 offset: CodeOffset,
1147 kind: LabelUse,
1148 ) -> Result<(), AsmError> {
1149 kind.validate_for_arch(self.env.arch())?;
1150 let end = (offset as usize)
1151 .checked_add(kind.patch_size())
1152 .ok_or(AsmError::TooLarge)?;
1153 if end > self.data.len() {
1154 return Err(AsmError::InvalidArgument);
1155 }
1156 Ok(())
1157 }
1158
1159 fn handle_fixup(&mut self, fixup: AsmFixup) -> Result<(), AsmError> {
1160 let AsmFixup {
1161 kind,
1162 label,
1163 offset,
1164 } = fixup;
1165 let start = offset;
1166 let end = (offset as usize)
1167 .checked_add(kind.patch_size())
1168 .ok_or(AsmError::TooLarge)?;
1169 if end > self.data.len() {
1170 return Err(AsmError::InvalidArgument);
1171 }
1172
1173 let label_offset = self.label_offset(label);
1174 if label_offset != u32::MAX {
1175 if !kind.can_reach(offset, label_offset) {
1176 if label_offset < offset {
1177 self.emit_veneer(label, offset, kind)?;
1178 } else {
1179 return Err(AsmError::TooLarge);
1180 }
1181 } else {
1182 let slice = &mut self.data[start as usize..end];
1183
1184 kind.patch(slice, start, label_offset);
1185 }
1186 } else {
1187 self.emit_veneer(label, offset, kind)?;
1191 }
1192 Ok(())
1193 }
1194
1195 pub fn emit_veneer(
1202 &mut self,
1203 label: Label,
1204 offset: CodeOffset,
1205 kind: LabelUse,
1206 ) -> Result<(), AsmError> {
1207 if let Some(error) = self.error.clone() {
1208 return Err(error);
1209 }
1210 kind.validate_for_arch(self.env.arch())?;
1211 if !kind.supports_veneer() {
1212 return Err(AsmError::UnsupportedInstruction {
1213 reason: "branch range requires an unsupported veneer",
1214 });
1215 }
1216 if self.label_offsets.get(label.id() as usize).is_none() {
1217 return Err(AsmError::InvalidArgument);
1218 }
1219 let start = offset as usize;
1220 let end = start
1221 .checked_add(kind.patch_size())
1222 .ok_or(AsmError::TooLarge)?;
1223 if end > self.data.len() {
1224 return Err(AsmError::InvalidArgument);
1225 }
1226 #[cfg(not(feature = "riscv"))]
1227 {
1228 let _ = (label, offset, kind);
1229 Err(AsmError::UnsupportedInstruction {
1230 reason: "RISC-V veneer support is disabled",
1231 })
1232 }
1233
1234 #[cfg(feature = "riscv")]
1235 {
1236 self.try_align_to(kind.align() as _)?;
1237 let veneer_offset = self.cur_offset();
1238 let slice = &mut self.data[start..end];
1239
1240 kind.patch(slice, offset, veneer_offset);
1241 let veneer_slice = self.get_appended_space(kind.veneer_size());
1242 let (veneer_fixup_off, veneer_label_use) =
1243 kind.generate_veneer(veneer_slice, veneer_offset);
1244
1245 self.use_label_at_offset(veneer_fixup_off, label, veneer_label_use);
1250 if let Some(error) = self.error.clone() {
1251 return Err(error);
1252 }
1253 Ok(())
1254 }
1255 }
1256
1257 pub(crate) fn get_appended_space(&mut self, len: usize) -> &mut [u8] {
1259 let off = self.data.len();
1260 let new_len = self.data.len() + len;
1261 self.data.resize(new_len, 0);
1262 &mut self.data[off..]
1263
1264 }
1266
1267 fn worst_case_end_of_island(&self, distance: CodeOffset) -> CodeOffset {
1273 let island_worst_case_size =
1278 ((self.fixup_records.len() + self.pending_fixup_records.len()) as u32) * 20
1279 + self.pending_constants_size;
1280 self.cur_offset()
1281 .saturating_add(distance)
1282 .saturating_add(island_worst_case_size)
1283 }
1284
1285 fn should_apply_fixup(&self, fixup: &AsmFixup, forced_threshold: CodeOffset) -> bool {
1286 let label_offset = self.label_offset(fixup.label);
1287 label_offset != u32::MAX
1288 || fixup.offset.saturating_add(fixup.kind.max_pos_range()) < forced_threshold
1289 }
1290 pub fn island_needed(&mut self, distance: CodeOffset) -> bool {
1292 let deadline = match self.fixup_records.peek() {
1293 Some(fixup) => fixup
1294 .offset
1295 .saturating_add(fixup.kind.max_pos_range())
1296 .min(self.pending_fixup_deadline),
1297 None => self.pending_fixup_deadline,
1298 };
1299
1300 deadline < u32::MAX && self.worst_case_end_of_island(distance) > deadline
1301 }
1302
1303 pub fn emit_island(&mut self, distance: CodeOffset) -> Result<(), AsmError> {
1305 if let Some(error) = self.error.clone() {
1306 return Err(error);
1307 }
1308 let forced_threshold = self.worst_case_end_of_island(distance);
1309
1310 for constant in core::mem::take(&mut self.pending_constants) {
1311 let (_, AsmConstant { align, size, .. }) = self.constants[constant.0 as usize];
1312 let label = self.constants[constant.0 as usize]
1313 .1
1314 .upcoming_label
1315 .take()
1316 .unwrap();
1317 self.try_align_to(align as _)?;
1318 self.try_bind_label(label)?;
1319 self.used_constants.push((constant, self.cur_offset()));
1320 self.get_appended_space(size);
1321 }
1322 for fixup in core::mem::take(&mut self.pending_fixup_records) {
1326 if self.should_apply_fixup(&fixup, forced_threshold) {
1327 self.handle_fixup(fixup)?;
1328 } else {
1329 self.fixup_records.push(fixup);
1330 }
1331 }
1332
1333 self.pending_fixup_deadline = u32::MAX;
1334
1335 while let Some(fixup) = self.fixup_records.peek() {
1336 if !self.should_apply_fixup(fixup, forced_threshold) {
1342 break;
1343 }
1344 let fixup = self.fixup_records.pop().unwrap();
1345 self.handle_fixup(fixup)?;
1346 }
1347 Ok(())
1348 }
1349
1350 fn finish_emission_maybe_forcing_veneers(&mut self) -> Result<(), AsmError> {
1351 while !self.pending_constants.is_empty()
1352 || !self.pending_fixup_records.is_empty()
1353 || !self.fixup_records.is_empty()
1354 {
1355 self.emit_island(u32::MAX)?;
1359 }
1360 Ok(())
1361 }
1362
1363 #[cfg(feature = "x86")]
1364 fn relax_x86_branches(&mut self) -> Result<(), AsmError> {
1365 loop {
1366 let mut selected = None;
1367
1368 for (index, &candidate) in self.x86_branch_relaxations.iter().enumerate() {
1369 let start = candidate.opcode_offset as usize;
1370 let near_size = candidate.near_size as usize;
1371 let end = start.checked_add(near_size).ok_or(AsmError::TooLarge)?;
1372 if end > self.data.len() {
1373 return Err(AsmError::InvalidState);
1374 }
1375 let valid_encoding = match candidate.near_size {
1376 5 => self.data[start] == 0xE9 && candidate.short_opcode == 0xEB,
1377 6 => {
1378 self.data[start] == 0x0F
1379 && self.data[start + 1] == candidate.short_opcode.wrapping_add(0x10)
1380 }
1381 _ => false,
1382 };
1383 if !valid_encoding {
1384 return Err(AsmError::InvalidState);
1385 }
1386
1387 let target = self.label_offset(candidate.label);
1388 if target == u32::MAX {
1389 return Err(AsmError::UnboundLabel);
1390 }
1391 let old_end = candidate
1392 .opcode_offset
1393 .checked_add(candidate.near_size as u32)
1394 .ok_or(AsmError::TooLarge)?;
1395 if target > candidate.opcode_offset && target < old_end {
1396 return Err(AsmError::InvalidState);
1397 }
1398
1399 let removed = candidate.near_size as u32 - 2;
1400 let target_after = if target >= old_end {
1401 target - removed
1402 } else {
1403 target
1404 };
1405 let short_end = candidate
1406 .opcode_offset
1407 .checked_add(2)
1408 .ok_or(AsmError::TooLarge)?;
1409 let displacement = i64::from(target_after) - i64::from(short_end);
1410 let Ok(displacement) = i8::try_from(displacement) else {
1411 continue;
1412 };
1413
1414 let cut_start = short_end;
1415 let cut_end = old_end;
1416 let preserves_alignment =
1417 self.alignment_constraints.iter().all(|&(offset, align)| {
1418 offset < cut_end || (offset - removed) & (align - 1) == 0
1419 }) && self.patch_blocks.iter().all(|block| {
1420 block.offset < cut_end || (block.offset - removed) & (block.align - 1) == 0
1421 });
1422 if !preserves_alignment {
1423 continue;
1424 }
1425
1426 selected = Some((index, candidate, displacement, cut_start, cut_end));
1427 break;
1428 }
1429
1430 let Some((index, candidate, displacement, cut_start, cut_end)) = selected else {
1431 return Ok(());
1432 };
1433 let removed = cut_end - cut_start;
1434 let fixup_offset = cut_end - LabelUse::X86JmpRel32.patch_size() as u32;
1435 let is_candidate_fixup = |fixup: &AsmFixup| {
1436 fixup.offset == fixup_offset
1437 && fixup.label == candidate.label
1438 && fixup.kind == LabelUse::X86JmpRel32
1439 };
1440
1441 let overlaps_cut = |offset: CodeOffset, size: usize| -> Result<bool, AsmError> {
1442 let end = offset
1443 .checked_add(u32::try_from(size).map_err(|_| AsmError::TooLarge)?)
1444 .ok_or(AsmError::TooLarge)?;
1445 Ok(offset < cut_end && end > cut_start)
1446 };
1447 for reloc in &self.relocs {
1448 if overlaps_cut(reloc.offset, relocation_patch_size(reloc.kind)?)? {
1449 return Err(AsmError::InvalidState);
1450 }
1451 }
1452 for fixup in self
1453 .pending_fixup_records
1454 .iter()
1455 .chain(self.fixup_records.iter())
1456 {
1457 if !is_candidate_fixup(fixup)
1458 && overlaps_cut(fixup.offset, fixup.kind.patch_size())?
1459 {
1460 return Err(AsmError::InvalidState);
1461 }
1462 }
1463 for block in &self.patch_blocks {
1464 if overlaps_cut(block.offset, block.size as usize)? {
1465 return Err(AsmError::InvalidState);
1466 }
1467 }
1468 for site in &self.patch_sites {
1469 if overlaps_cut(site.offset, site.kind.patch_size())?
1470 || matches!(site.target, PendingPatchTarget::Offset(offset) if (cut_start..cut_end).contains(&offset))
1471 {
1472 return Err(AsmError::InvalidState);
1473 }
1474 }
1475 if self
1476 .label_offsets
1477 .iter()
1478 .any(|&offset| offset != u32::MAX && (cut_start..cut_end).contains(&offset))
1479 || self
1480 .used_constants
1481 .iter()
1482 .any(|&(_, offset)| (cut_start..cut_end).contains(&offset))
1483 || self
1484 .alignment_constraints
1485 .iter()
1486 .any(|&(offset, _)| (cut_start..cut_end).contains(&offset))
1487 {
1488 return Err(AsmError::InvalidState);
1489 }
1490 for (other_index, other) in self.x86_branch_relaxations.iter().enumerate() {
1491 if other_index != index
1492 && overlaps_cut(other.opcode_offset, other.near_size as usize)?
1493 {
1494 return Err(AsmError::InvalidState);
1495 }
1496 }
1497
1498 let start = candidate.opcode_offset as usize;
1499 self.data[start] = candidate.short_opcode;
1500 self.data[start + 1] = displacement as u8;
1501 self.data.drain(cut_start as usize..cut_end as usize);
1502 self.x86_branch_relaxations.remove(index);
1503
1504 self.pending_fixup_records
1505 .retain(|fixup| !is_candidate_fixup(fixup));
1506 let mut fixups = core::mem::take(&mut self.fixup_records).into_vec();
1507 fixups.retain(|fixup| !is_candidate_fixup(fixup));
1508
1509 let rebase = |offset: &mut CodeOffset| -> Result<(), AsmError> {
1510 if *offset >= cut_end {
1511 *offset -= removed;
1512 } else if *offset >= cut_start {
1513 return Err(AsmError::InvalidState);
1514 }
1515 Ok(())
1516 };
1517
1518 for reloc in &mut self.relocs {
1519 rebase(&mut reloc.offset)?;
1520 }
1521 for offset in &mut self.label_offsets {
1522 if *offset != u32::MAX {
1523 rebase(offset)?;
1524 }
1525 }
1526 for fixup in &mut self.pending_fixup_records {
1527 rebase(&mut fixup.offset)?;
1528 }
1529 for fixup in &mut fixups {
1530 rebase(&mut fixup.offset)?;
1531 }
1532 self.fixup_records = BinaryHeap::from(fixups);
1533 self.pending_fixup_deadline = self
1534 .pending_fixup_records
1535 .iter()
1536 .map(AsmFixup::deadline)
1537 .min()
1538 .unwrap_or(u32::MAX);
1539 for (_, offset) in &mut self.used_constants {
1540 rebase(offset)?;
1541 }
1542 for block in &mut self.patch_blocks {
1543 rebase(&mut block.offset)?;
1544 }
1545 for site in &mut self.patch_sites {
1546 rebase(&mut site.offset)?;
1547 if let PendingPatchTarget::Offset(offset) = &mut site.target {
1548 rebase(offset)?;
1549 }
1550 }
1551 for relaxation in &mut self.x86_branch_relaxations {
1552 rebase(&mut relaxation.opcode_offset)?;
1553 }
1554 for (offset, _) in &mut self.alignment_constraints {
1555 rebase(offset)?;
1556 }
1557 }
1558 }
1559
1560 fn preflight_finalization(&self) -> Result<(), AsmError> {
1563 for fixup in self
1564 .pending_fixup_records
1565 .iter()
1566 .chain(self.fixup_records.iter())
1567 {
1568 let end = (fixup.offset as usize)
1569 .checked_add(fixup.kind.patch_size())
1570 .ok_or(AsmError::TooLarge)?;
1571 if end > self.data.len() {
1572 return Err(AsmError::InvalidArgument);
1573 }
1574 let label_offset = self.label_offset(fixup.label);
1575 if label_offset == u32::MAX {
1576 return Err(AsmError::UnboundLabel);
1577 }
1578 if fixup.kind.can_reach(fixup.offset, label_offset) {
1579 continue;
1580 }
1581 if matches!(
1582 fixup.kind,
1583 LabelUse::A64Branch14 | LabelUse::A64Branch19 | LabelUse::A64Branch26
1584 ) && !fixup.kind.supports_veneer()
1585 {
1586 return Err(AsmError::UnsupportedInstruction {
1587 reason: "AArch64 branch veneers are not implemented",
1588 });
1589 }
1590 if label_offset >= fixup.offset {
1591 return Err(AsmError::TooLarge);
1592 }
1593 if !fixup.kind.supports_veneer() {
1594 return Err(AsmError::UnsupportedInstruction {
1595 reason: "branch range requires an unsupported veneer",
1596 });
1597 }
1598 #[cfg(not(feature = "riscv"))]
1599 return Err(AsmError::UnsupportedInstruction {
1600 reason: "RISC-V veneer support is disabled",
1601 });
1602 }
1603 Ok(())
1604 }
1605
1606 fn finish_constants(&mut self) -> u32 {
1607 let mut alignment = 32;
1608
1609 for (constant, offset) in core::mem::take(&mut self.used_constants) {
1610 let constant = &self.constants[constant.0 as usize].0;
1611 let data = constant.as_slice();
1612 self.data[offset as usize..][..data.len()].copy_from_slice(data);
1613 alignment = constant.alignment().max(alignment);
1614 }
1615
1616 alignment as _
1617 }
1618
1619 fn resolve_patch_catalog(&self, validate_ranges: bool) -> Result<PatchCatalog, AsmError> {
1620 let mut blocks = SmallVec::new();
1621 let mut sites = SmallVec::new();
1622
1623 for block in &self.patch_blocks {
1624 blocks.push(PatchBlock {
1625 offset: block.offset,
1626 size: block.size,
1627 align: block.align,
1628 });
1629 }
1630
1631 for site in &self.patch_sites {
1632 let target_offset = match site.target {
1633 PendingPatchTarget::Offset(offset) => offset,
1634 PendingPatchTarget::Label(label) => self.label_offset(label),
1635 };
1636
1637 if target_offset == u32::MAX {
1638 return Err(AsmError::InvalidState);
1639 }
1640
1641 if validate_ranges && !site.kind.can_reach(site.offset, target_offset) {
1642 return Err(AsmError::TooLarge);
1643 }
1644
1645 sites.push(PatchSite {
1646 offset: site.offset,
1647 kind: site.kind,
1648 current_target: target_offset,
1649 addend: site.addend,
1650 });
1651 }
1652
1653 Ok(PatchCatalog::with_parts(self.env.arch(), blocks, sites))
1654 }
1655
1656 fn resolved_defined_symbols(&self) -> SmallVec<[(ExternalName, CodeOffset); 4]> {
1657 self.defined_symbols
1660 .iter()
1661 .map(|(name, label)| (name.clone(), self.label_offset(*label)))
1662 .collect()
1663 }
1664
1665 pub fn finish_patched(mut self) -> Result<CodeBufferFinalized, AsmError> {
1666 if let Some(error) = self.error.take() {
1667 return Err(error);
1668 }
1669 if self.has_unbound_labels() {
1670 return Err(AsmError::UnboundLabel);
1671 }
1672 self.preflight_finalization()?;
1673 #[cfg(feature = "x86")]
1674 self.relax_x86_branches()?;
1675 #[cfg(feature = "x86")]
1676 self.preflight_finalization()?;
1677 self.finish_emission_maybe_forcing_veneers()?;
1678 let patch_catalog = self.resolve_patch_catalog(true)?;
1679 let alignment = self.finish_constants();
1680 let defined_symbols = self.resolved_defined_symbols();
1681 Ok(CodeBufferFinalized {
1682 data: self.data,
1683 relocs: self.relocs,
1684 symbols: self.symbols,
1685 label_offsets: self.label_offsets,
1686 defined_symbols,
1687 alignment,
1688 patch_catalog,
1689 })
1690 }
1691
1692 pub fn finish(&mut self) -> Result<CodeBufferFinalized, AsmError> {
1693 if let Some(error) = self.error.clone() {
1694 return Err(error);
1695 }
1696 if self.has_unbound_labels() {
1697 return Err(AsmError::UnboundLabel);
1698 }
1699 self.preflight_finalization()?;
1700 #[cfg(feature = "x86")]
1701 self.relax_x86_branches()?;
1702 #[cfg(feature = "x86")]
1703 self.preflight_finalization()?;
1704 self.finish_emission_maybe_forcing_veneers()?;
1705 let patch_catalog = self.resolve_patch_catalog(false)?;
1706 let alignment = self.finish_constants();
1707 Ok(CodeBufferFinalized {
1708 data: self.data.clone(),
1709 relocs: self.relocs.clone(),
1710 symbols: self.symbols.clone(),
1711 label_offsets: self.label_offsets.clone(),
1712 defined_symbols: self.resolved_defined_symbols(),
1713 alignment,
1714 patch_catalog,
1715 })
1716 }
1717
1718 fn has_unbound_labels(&self) -> bool {
1719 let is_unbound = |label: Label| {
1720 self.label_offsets
1721 .get(label.id() as usize)
1722 .is_none_or(|offset| *offset == u32::MAX)
1723 };
1724 self.pending_fixup_records
1725 .iter()
1726 .any(|fixup| is_unbound(fixup.label))
1727 || self
1728 .fixup_records
1729 .iter()
1730 .any(|fixup| is_unbound(fixup.label))
1731 || self
1732 .relocs
1733 .iter()
1734 .any(|reloc| matches!(reloc.target, RelocTarget::Label(label) if is_unbound(label)))
1735 || self
1736 .defined_symbols
1737 .iter()
1738 .any(|(_, label)| is_unbound(*label))
1739 || self.patch_sites.iter().any(
1740 |site| matches!(site.target, PendingPatchTarget::Label(label) if is_unbound(label)),
1741 )
1742 }
1743}
1744
1745#[derive(Clone, PartialEq, Eq, Debug, Hash)]
1746pub enum ConstantData {
1747 WellKnown(&'static [u8]),
1748 U64([u8; 8]),
1749 Bytes(Vec<u8>),
1750}
1751
1752impl ConstantData {
1753 pub fn as_slice(&self) -> &[u8] {
1754 match self {
1755 ConstantData::WellKnown(data) => data,
1756 ConstantData::U64(data) => data.as_ref(),
1757 ConstantData::Bytes(data) => data,
1758 }
1759 }
1760
1761 pub fn alignment(&self) -> usize {
1762 if self.as_slice().len() <= 8 { 8 } else { 16 }
1763 }
1764}
1765
1766#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1768pub struct Constant(pub(crate) u32);
1769
1770impl From<&'static str> for ConstantData {
1771 fn from(value: &'static str) -> Self {
1772 Self::WellKnown(value.as_bytes())
1773 }
1774}
1775
1776impl From<[u8; 8]> for ConstantData {
1777 fn from(value: [u8; 8]) -> Self {
1778 Self::U64(value)
1779 }
1780}
1781
1782impl From<Vec<u8>> for ConstantData {
1783 fn from(value: Vec<u8>) -> Self {
1784 Self::Bytes(value)
1785 }
1786}
1787
1788impl From<&'static [u8]> for ConstantData {
1789 fn from(value: &'static [u8]) -> Self {
1790 Self::WellKnown(value)
1791 }
1792}
1793
1794impl From<u64> for ConstantData {
1795 fn from(value: u64) -> Self {
1796 Self::U64(value.to_ne_bytes())
1797 }
1798}
1799
1800pub type CodeOffset = u32;
1805
1806pub type Addend = i64;
1808
1809#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
1811pub enum Reloc {
1812 Abs4,
1814 Abs8,
1816 X86PCRel4,
1818 X86CallPCRel4,
1820 X86CallPLTRel4,
1822 X86GOTPCRel4,
1824 X86SecRel,
1828 Arm32Call,
1830 Arm64Call,
1834
1835 ElfX86_64TlsGd,
1837
1838 MachOX86_64Tlv,
1840
1841 MachOAarch64TlsAdrPage21,
1844
1845 MachOAarch64TlsAdrPageOff12,
1848
1849 Aarch64TlsDescAdrPage21,
1852
1853 Aarch64TlsDescLd64Lo12,
1856
1857 Aarch64TlsDescAddLo12,
1860
1861 Aarch64TlsDescCall,
1864
1865 Aarch64AdrGotPage21,
1869
1870 Aarch64Ld64GotLo12Nc,
1875
1876 Aarch64AdrPrelPgHi21,
1878 Aarch64AddAbsLo12Nc,
1880
1881 RiscvAbs8,
1883
1884 RiscvCallPlt,
1892
1893 RiscvTlsGdHi20,
1898
1899 RiscvPCRelLo12I,
1904
1905 RiscvGotHi20,
1910}
1911
1912impl Reloc {
1913 const fn supports_arch(self, arch: Arch) -> bool {
1914 match self {
1915 Self::Abs4 | Self::Abs8 => true,
1916 Self::X86PCRel4
1917 | Self::X86CallPCRel4
1918 | Self::X86CallPLTRel4
1919 | Self::X86GOTPCRel4
1920 | Self::X86SecRel
1921 | Self::ElfX86_64TlsGd
1922 | Self::MachOX86_64Tlv => {
1923 cfg!(feature = "x86") && matches!(arch, Arch::X86 | Arch::X64)
1924 }
1925 Self::Arm32Call => false,
1926 Self::Arm64Call
1927 | Self::MachOAarch64TlsAdrPage21
1928 | Self::MachOAarch64TlsAdrPageOff12
1929 | Self::Aarch64TlsDescAdrPage21
1930 | Self::Aarch64TlsDescLd64Lo12
1931 | Self::Aarch64TlsDescAddLo12
1932 | Self::Aarch64TlsDescCall
1933 | Self::Aarch64AdrGotPage21
1934 | Self::Aarch64Ld64GotLo12Nc
1935 | Self::Aarch64AdrPrelPgHi21
1936 | Self::Aarch64AddAbsLo12Nc => {
1937 cfg!(feature = "aarch64") && matches!(arch, Arch::AArch64)
1938 }
1939 Self::RiscvAbs8
1940 | Self::RiscvCallPlt
1941 | Self::RiscvTlsGdHi20
1942 | Self::RiscvPCRelLo12I
1943 | Self::RiscvGotHi20 => {
1944 cfg!(feature = "riscv") && matches!(arch, Arch::RISCV32 | Arch::RISCV64)
1945 }
1946 }
1947 }
1948}
1949
1950#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
1951pub enum LabelUse {
1952 X86JmpRel32,
1953 RVJal20,
1956 RVPCRel32,
1965
1966 RVB12,
1971
1972 RVPCRelHi20,
1975
1976 RVPCRelLo12I,
1983
1984 RVCJump,
1986 RVCB9,
1988 A64Branch14,
1991 A64Branch19,
1994 A64Branch26,
1997 A64Ldr19,
2000 A64Adr21,
2003 A64Adrp21,
2006 A64Ldr12,
2007
2008 A64AddAbsLo12,
2009}
2010
2011impl LabelUse {
2012 fn validate_for_arch(self, arch: Arch) -> Result<(), AsmError> {
2013 if self == Self::A64Ldr12 {
2014 return Err(AsmError::UnsupportedInstruction {
2015 reason: "AArch64 LDR12 label patching is not implemented",
2016 });
2017 }
2018 if !self.supports_arch(arch) {
2019 return Err(AsmError::InvalidArch);
2020 }
2021 Ok(())
2022 }
2023
2024 const fn supports_arch(self, arch: Arch) -> bool {
2025 match self {
2026 Self::X86JmpRel32 => cfg!(feature = "x86") && matches!(arch, Arch::X86 | Arch::X64),
2027 Self::RVJal20
2028 | Self::RVPCRel32
2029 | Self::RVB12
2030 | Self::RVPCRelHi20
2031 | Self::RVPCRelLo12I
2032 | Self::RVCJump
2033 | Self::RVCB9 => {
2034 cfg!(feature = "riscv") && matches!(arch, Arch::RISCV32 | Arch::RISCV64)
2035 }
2036 Self::A64Branch14
2037 | Self::A64Branch19
2038 | Self::A64Branch26
2039 | Self::A64Ldr19
2040 | Self::A64Adr21
2041 | Self::A64Adrp21
2042 | Self::A64Ldr12
2043 | Self::A64AddAbsLo12 => cfg!(feature = "aarch64") && matches!(arch, Arch::AArch64),
2044 }
2045 }
2046
2047 pub fn can_reach(&self, use_offset: CodeOffset, label_offset: CodeOffset) -> bool {
2048 let delta = (label_offset as i64) - (use_offset as i64);
2049
2050 match self {
2051 Self::X86JmpRel32 => {
2052 let disp = delta - 4;
2053 i32::try_from(disp).is_ok()
2054 }
2055 Self::RVJal20 => delta % 2 == 0 && (-(1 << 20)..=((1 << 20) - 2)).contains(&delta),
2056 Self::RVB12 => delta % 2 == 0 && (-(1 << 12)..=((1 << 12) - 2)).contains(&delta),
2057 Self::RVCJump => delta % 2 == 0 && (-(1 << 11)..=((1 << 11) - 2)).contains(&delta),
2058 Self::RVCB9 => delta % 2 == 0 && (-(1 << 8)..=((1 << 8) - 2)).contains(&delta),
2059 Self::RVPCRelHi20 | Self::RVPCRelLo12I | Self::RVPCRel32 => {
2060 i32::try_from(delta).is_ok()
2061 }
2062 Self::A64Branch14 => delta % 4 == 0 && (-(1 << 15)..=((1 << 15) - 4)).contains(&delta),
2063 Self::A64Branch19 | Self::A64Ldr19 => {
2064 delta % 4 == 0 && (-(1 << 20)..=((1 << 20) - 4)).contains(&delta)
2065 }
2066 Self::A64Branch26 => delta % 4 == 0 && (-(1 << 27)..=((1 << 27) - 4)).contains(&delta),
2067 Self::A64Adr21 => (-(1 << 20)..=((1 << 20) - 1)).contains(&delta),
2068 Self::A64Adrp21 => {
2069 let page_delta = ((label_offset & !0xfff) as i64) - ((use_offset & !0xfff) as i64);
2070 page_delta % 4096 == 0 && (-(1 << 32)..=((1 << 32) - 4096)).contains(&page_delta)
2071 }
2072
2073 Self::A64AddAbsLo12 => {
2074 delta % 4096 == delta && (-(1 << 12)..=(1 << 12) - 1).contains(&delta)
2075 }
2076
2077 Self::A64Ldr12 => true,
2078 }
2079 }
2080
2081 pub const fn max_pos_range(self) -> CodeOffset {
2083 match self {
2084 LabelUse::RVJal20 => ((1 << 19) - 1) * 2,
2085 LabelUse::RVPCRelLo12I | LabelUse::RVPCRelHi20 | LabelUse::RVPCRel32 => {
2086 let imm20_max: i64 = ((1 << 19) - 1) << 12;
2087 let imm12_max = (1 << 11) - 1;
2088 (imm20_max + imm12_max) as _
2089 }
2090 LabelUse::RVB12 => ((1 << 11) - 1) * 2,
2091 LabelUse::RVCB9 => ((1 << 8) - 1) * 2,
2092 LabelUse::RVCJump => ((1 << 10) - 1) * 2,
2093 LabelUse::X86JmpRel32 => i32::MAX as _,
2094 _ => u32::MAX,
2095 }
2096 }
2097
2098 pub const fn max_neg_range(self) -> CodeOffset {
2099 match self {
2100 LabelUse::RVPCRel32 => {
2101 let imm20_max: i64 = (1 << 19) << 12;
2102 let imm12_max = 1 << 11;
2103 (-imm20_max - imm12_max) as CodeOffset
2104 }
2105 _ => self.max_pos_range() + 2,
2106 }
2107 }
2108
2109 pub const fn patch_size(&self) -> usize {
2110 match self {
2111 Self::X86JmpRel32 => 4,
2112 Self::RVCJump | Self::RVCB9 => 2,
2113 Self::RVJal20 | Self::RVB12 | Self::RVPCRelHi20 | Self::RVPCRelLo12I => 4,
2114 Self::RVPCRel32 => 8,
2115 _ => 4,
2116 }
2117 }
2118
2119 pub const fn align(&self) -> usize {
2120 match self {
2121 Self::X86JmpRel32 => 1,
2122 Self::RVCJump => 4,
2123 Self::RVJal20 | Self::RVB12 | Self::RVCB9 | Self::RVPCRelHi20 | Self::RVPCRelLo12I => 4,
2124 Self::RVPCRel32 => 4,
2125 _ => 4,
2126 }
2127 }
2128
2129 pub const fn supports_veneer(&self) -> bool {
2130 matches!(self, Self::RVB12 | Self::RVJal20 | Self::RVCJump)
2131 }
2132
2133 #[cfg(feature = "riscv")]
2134 pub(crate) fn veneer_size(&self) -> usize {
2135 debug_assert!(self.supports_veneer());
2136 8
2137 }
2138
2139 #[cfg(feature = "riscv")]
2140 pub(crate) fn generate_veneer(
2141 &self,
2142 buffer: &mut [u8],
2143 veneer_offset: CodeOffset,
2144 ) -> (CodeOffset, Self) {
2145 debug_assert!(self.supports_veneer());
2146 let base = riscv::X31;
2147
2148 {
2149 let x = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2150 .encode()
2151 .set_rd(base.id())
2152 .value
2153 .to_le_bytes();
2154 buffer[0] = x[0];
2155 buffer[1] = x[1];
2156 buffer[2] = x[2];
2157 buffer[3] = x[3];
2158 }
2159
2160 {
2161 let x = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2162 .encode()
2163 .set_rd(riscv::ZERO.id())
2164 .set_rs1(base.id())
2165 .value
2166 .to_le_bytes();
2167 buffer[4] = x[0];
2168 buffer[5] = x[1];
2169 buffer[6] = x[2];
2170 buffer[7] = x[3];
2171 }
2172
2173 (veneer_offset, LabelUse::RVPCRel32)
2174 }
2175 pub(crate) fn patch(
2176 &self,
2177 buffer: &mut [u8],
2178 use_offset: CodeOffset,
2179 label_offset: CodeOffset,
2180 ) {
2181 let addend = match self {
2182 Self::X86JmpRel32 => i64::from(u32::from_le_bytes([
2183 buffer[0], buffer[1], buffer[2], buffer[3],
2184 ])),
2185 _ => 0,
2186 };
2187
2188 self.patch_with_addend(buffer, use_offset, label_offset, addend);
2189 }
2190
2191 pub(crate) fn patch_with_addend(
2192 &self,
2193 buffer: &mut [u8],
2194 use_offset: CodeOffset,
2195 label_offset: CodeOffset,
2196 addend: i64,
2197 ) {
2198 let pc_reli = (label_offset as i64) - (use_offset as i64);
2199
2200 let pc_rel = pc_reli as u32;
2201
2202 match self {
2203 Self::X86JmpRel32 => {
2204 let value = pc_rel.wrapping_add(addend as u32).wrapping_sub(4);
2205
2206 buffer.copy_from_slice(&value.to_le_bytes());
2207 }
2208
2209 Self::RVJal20 => {
2210 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2211 let offset = pc_rel;
2212 let v = ((offset >> 12 & 0b1111_1111) << 12)
2213 | ((offset >> 11 & 0b1) << 20)
2214 | ((offset >> 1 & 0b11_1111_1111) << 21)
2215 | ((offset >> 20 & 0b1) << 31);
2216 buffer[0..4].clone_from_slice(&u32::to_le_bytes(insn | v));
2217 }
2218
2219 Self::RVPCRel32 => {
2220 #[cfg(feature = "riscv")]
2221 {
2222 let (imm20, imm12) = generate_imm(pc_rel as u64);
2223 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2224 let insn2 = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]);
2225
2226 let auipc = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2227 .encode()
2228 .set_imm20(0);
2229 let jalr = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2230 .encode()
2231 .set_rd(0)
2232 .set_rs1(0)
2233 .set_imm12(0);
2234
2235 buffer[0..4].copy_from_slice(&(insn | auipc.value | imm20).to_le_bytes());
2236 buffer[4..8].copy_from_slice(&(insn2 | jalr.value | imm12).to_le_bytes());
2237 }
2238 #[cfg(not(feature = "riscv"))]
2239 {
2240 panic!("RISC-V veneers aren't supported without the `riscv` feature");
2241 }
2242 }
2243
2244 Self::RVB12 => {
2245 #[cfg(feature = "riscv")]
2246 {
2247 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2248 let offset = pc_rel;
2249 let v = ((offset >> 11 & 0b1) << 7)
2250 | ((offset >> 1 & 0b1111) << 8)
2251 | ((offset >> 5 & 0b11_1111) << 25)
2252 | ((offset >> 12 & 0b1) << 31);
2253 buffer[0..4].clone_from_slice(&u32::to_le_bytes(insn | v));
2254 }
2255 #[cfg(not(feature = "riscv"))]
2256 {
2257 panic!("RISC-V veneers aren't supported without the `riscv` feature");
2258 }
2259 }
2260
2261 Self::RVPCRelHi20 => {
2262 let offset = pc_reli as u32;
2268 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2269 let hi20 = offset.wrapping_add(0x800) >> 12;
2270 let insn = (insn & 0xfff) | (hi20 << 12);
2271 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2272 }
2273
2274 Self::RVPCRelLo12I => {
2275 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2286
2287 let lo12 = (pc_reli + 4) as u32 & 0xfff;
2288 let insn = (insn & 0xFFFFF) | (lo12 << 20);
2289 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2290 }
2291
2292 Self::RVCJump => {
2293 debug_assert!(pc_rel & 1 == 0);
2294
2295 #[cfg(feature = "riscv")]
2296 {
2297 let insn = riscv::opcodes::Inst::new(riscv::Opcode::CJ)
2298 .encode()
2299 .set_c_imm12(pc_rel as _);
2300 buffer[0..2].clone_from_slice(&(insn.value as u16).to_le_bytes());
2301 }
2302 #[cfg(not(feature = "riscv"))]
2303 {
2304 panic!("RISC-V jumps aren't supported without the `riscv` feature");
2305 }
2306 }
2307
2308 Self::RVCB9 => {
2309 debug_assert!(pc_rel & 1 == 0);
2310
2311 #[cfg(feature = "riscv")]
2312 {
2313 let insn = riscv::opcodes::Inst::new(riscv::Opcode::BEQZ)
2314 .encode()
2315 .set_c_bimm9lohi(pc_rel as _);
2316 buffer[0..2].clone_from_slice(&(insn.value as u16).to_le_bytes());
2317 }
2318 #[cfg(not(feature = "riscv"))]
2319 {
2320 panic!("RISC-V veneers aren't supported without the `riscv` feature");
2321 }
2322 }
2323
2324 Self::A64Branch14 => {
2325 debug_assert!(pc_reli & 0b11 == 0);
2326
2327 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2328 let imm14 = ((pc_reli >> 2) as i32 as u32) & 0x3fff;
2329 let insn = (insn & !0x0007ffe0) | (imm14 << 5);
2330 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2331 }
2332
2333 Self::A64Branch19 | Self::A64Ldr19 => {
2334 debug_assert!(pc_reli & 0b11 == 0);
2335
2336 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2337 let imm19 = ((pc_reli >> 2) as i32 as u32) & 0x7ffff;
2338 let insn = (insn & !0x00ffffe0) | (imm19 << 5);
2339 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2340 }
2341
2342 Self::A64Branch26 => {
2343 debug_assert!(pc_reli & 0b11 == 0);
2344
2345 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2346 let imm26 = ((pc_reli >> 2) as i32 as u32) & 0x03ff_ffff;
2347 let insn = (insn & !0x03ff_ffff) | imm26;
2348 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2349 }
2350
2351 Self::A64Adr21 => {
2352 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2353 let imm21 = (pc_reli as i32 as u32) & 0x1f_ffff;
2354 let immlo = imm21 & 0x3;
2355 let immhi = (imm21 >> 2) & 0x7ffff;
2356 let insn = (insn & !0x60ff_ffe0) | (immlo << 29) | (immhi << 5);
2357 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2358 }
2359
2360 Self::A64Adrp21 => {
2361 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2362
2363 let pc_page = (use_offset as i64) & !0xFFF;
2365 let target_page = ((label_offset as i64) + addend) & !0xFFF;
2366
2367 let page_offset = (target_page - pc_page) >> 12;
2369
2370 let imm21 = (page_offset as u32) & 0x1F_FFFF;
2372 let immlo = imm21 & 0x3; let immhi = (imm21 >> 2) & 0x7FFFF; let insn = (insn & !0x60FF_FFE0) | (immlo << 29) | (immhi << 5);
2378
2379 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2380 }
2381
2382 Self::A64AddAbsLo12 => {
2383 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2384
2385 let imm12 = ((pc_reli as i32 as u32) & 0xfff) << 10;
2386 let insn = insn | imm12;
2387 buffer[0..4].copy_from_slice(&insn.to_le_bytes());
2388 }
2389
2390 _ => todo!(),
2391 }
2392 }
2393}
2394
2395pub const fn is_imm12(val: i64) -> bool {
2396 val >= -2048 && val <= 2047
2397}
2398
2399#[allow(dead_code)]
2400pub(crate) fn generate_imm(value: u64) -> (u32, u32) {
2401 #[cfg(not(feature = "riscv"))]
2402 {
2403 let _ = value;
2404 panic!("Can't generate RISC-V immediates without the `riscv` feature");
2405 }
2406 #[cfg(feature = "riscv")]
2407 {
2408 if is_imm12(value as _) {
2409 return (
2410 0,
2411 riscv::opcodes::InstructionValue::new(0)
2412 .set_imm12(value as i64 as i32)
2413 .value,
2414 );
2415 }
2416
2417 let value = value as i64;
2418
2419 let mod_num = 4096i64;
2420 let (imm20, imm12) = if value > 0 {
2421 let mut imm20 = value / mod_num;
2422 let mut imm12 = value % mod_num;
2423
2424 if imm12 >= 2048 {
2425 imm12 -= mod_num;
2426 imm20 += 1;
2427 }
2428
2429 (imm20, imm12)
2430 } else {
2431 let value_abs = value.abs();
2432 let imm20 = value_abs / mod_num;
2433 let imm12 = value_abs % mod_num;
2434 let mut imm20 = -imm20;
2435 let mut imm12 = -imm12;
2436 if imm12 < -2048 {
2437 imm12 += mod_num;
2438 imm20 -= 1;
2439 }
2440 (imm20, imm12)
2441 };
2442 (
2443 riscv::opcodes::InstructionValue::new(0)
2444 .set_imm20(imm20 as _)
2445 .value,
2446 riscv::opcodes::InstructionValue::new(0)
2447 .set_imm12(imm12 as _)
2448 .value,
2449 )
2450 }
2451}
2452
2453pub(crate) fn relocation_patch_size(kind: Reloc) -> Result<usize, AsmError> {
2454 match kind {
2455 Reloc::Abs4
2456 | Reloc::X86PCRel4
2457 | Reloc::X86CallPCRel4
2458 | Reloc::X86CallPLTRel4
2459 | Reloc::X86GOTPCRel4
2460 | Reloc::RiscvGotHi20
2461 | Reloc::RiscvPCRelLo12I
2462 | Reloc::Aarch64AdrPrelPgHi21
2463 | Reloc::Aarch64AddAbsLo12Nc
2464 | Reloc::Aarch64AdrGotPage21
2465 | Reloc::Aarch64Ld64GotLo12Nc => Ok(4),
2466 Reloc::Abs8 | Reloc::RiscvAbs8 | Reloc::RiscvCallPlt => Ok(8),
2467 _ => Err(AsmError::InvalidArgument),
2468 }
2469}
2470
2471fn checked_address(base: usize, addend: i64) -> Result<usize, AsmError> {
2472 if addend >= 0 {
2473 base.checked_add(usize::try_from(addend).map_err(|_| AsmError::TooLarge)?)
2474 } else {
2475 base.checked_sub(usize::try_from(addend.unsigned_abs()).map_err(|_| AsmError::TooLarge)?)
2476 }
2477 .ok_or(AsmError::TooLarge)
2478}
2479
2480fn checked_pcrel(target: usize, instruction: usize) -> Result<i64, AsmError> {
2481 i64::try_from(target as i128 - instruction as i128).map_err(|_| AsmError::TooLarge)
2482}
2483
2484pub unsafe fn perform_relocations(
2495 code: *mut u8,
2496 code_rx: *const u8,
2497 code_size: usize,
2498 relocs: &[AsmReloc],
2499 get_address: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2500 get_got_entry: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2501 get_plt_entry: impl Fn(&RelocTarget) -> Result<*const u8, AsmError>,
2502) -> Result<(), AsmError> {
2503 use core::ptr::write_unaligned;
2504
2505 if code.is_null() || code_rx.is_null() {
2506 return Err(AsmError::InvalidArgument);
2507 }
2508
2509 for &AsmReloc {
2510 addend,
2511 kind,
2512 offset,
2513 ref target,
2514 } in relocs
2515 {
2516 let patch_size = relocation_patch_size(kind)?;
2517 let patch_start = offset as usize;
2518 let patch_end = patch_start
2519 .checked_add(patch_size)
2520 .ok_or(AsmError::TooLarge)?;
2521 if patch_end > code_size {
2522 return Err(AsmError::InvalidArgument);
2523 }
2524 let at = unsafe { code.add(patch_start) };
2525 let atrx = code_rx
2526 .addr()
2527 .checked_add(patch_start)
2528 .ok_or(AsmError::TooLarge)?;
2529 let resolve = |resolver: &dyn Fn(&RelocTarget) -> Result<*const u8, AsmError>| {
2530 let base = resolver(target)?;
2531 if base.is_null() {
2532 return Err(AsmError::InvalidArgument);
2533 }
2534 checked_address(base.addr(), addend)
2535 };
2536
2537 match kind {
2538 Reloc::Abs4 => {
2539 let what = resolve(&get_address)?;
2540 let what = u32::try_from(what).map_err(|_| AsmError::TooLarge)?;
2541 unsafe {
2542 write_unaligned(at as *mut u32, what);
2543 }
2544 }
2545
2546 Reloc::Abs8 | Reloc::RiscvAbs8 => {
2547 let what = resolve(&get_address)?;
2548 let what = u64::try_from(what).map_err(|_| AsmError::TooLarge)?;
2549 unsafe {
2550 write_unaligned(at as *mut u64, what);
2551 }
2552 }
2553
2554 Reloc::X86PCRel4 | Reloc::X86CallPCRel4 => {
2555 let what = resolve(&get_address)?;
2556 let pcrel =
2557 i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2558
2559 unsafe {
2560 write_unaligned(at as *mut i32, pcrel);
2561 }
2562 }
2563
2564 Reloc::X86GOTPCRel4 => {
2565 let what = resolve(&get_got_entry)?;
2566 let pcrel =
2567 i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2568
2569 unsafe {
2570 write_unaligned(at as *mut i32, pcrel);
2571 }
2572 }
2573
2574 Reloc::X86CallPLTRel4 => {
2575 let what = resolve(&get_plt_entry)?;
2576 let pcrel =
2577 i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2578 unsafe { write_unaligned(at as *mut i32, pcrel) };
2579 }
2580
2581 Reloc::RiscvGotHi20 => {
2582 let what = resolve(&get_got_entry)?;
2583 let pc_rel =
2584 i32::try_from(checked_pcrel(what, atrx)?).map_err(|_| AsmError::TooLarge)?;
2585 unsafe {
2586 let buffer = core::slice::from_raw_parts_mut(at, 4);
2587 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2588 let hi20 = (pc_rel as u32).wrapping_add(0x800) >> 12;
2589 let insn = (insn & 0xfff) | (hi20 << 12);
2590 buffer.copy_from_slice(&insn.to_le_bytes());
2591 }
2592 }
2593
2594 Reloc::RiscvPCRelLo12I => {
2595 let what = resolve(&get_got_entry)?;
2596 let pc_rel = checked_pcrel(what, atrx)?;
2597 let pc_rel = i32::try_from(pc_rel).map_err(|_| AsmError::TooLarge)?;
2598
2599 unsafe {
2600 let buffer = core::slice::from_raw_parts_mut(at, 4);
2601 let insn = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
2602 let lo12 = ((pc_rel as i64 + 4) as u32) & 0xfff;
2603 let insn = (insn & 0xFFFFF) | (lo12 << 20);
2604 buffer.copy_from_slice(&insn.to_le_bytes());
2605 }
2606 }
2607
2608 Reloc::RiscvCallPlt => {
2609 #[cfg(not(feature = "riscv"))]
2610 {
2611 return Err(AsmError::InvalidArgument);
2612 }
2613 #[cfg(feature = "riscv")]
2614 {
2615 let what = resolve(&get_address)?;
2621 let pcrel = i32::try_from(checked_pcrel(what, atrx)?)
2622 .map_err(|_| AsmError::TooLarge)?;
2623
2624 let hi20 = pcrel.wrapping_add(0x800) as u32 & 0xFFFFF000u32;
2638 let lo12 = (pcrel as u32).wrapping_sub(hi20) & 0xFFF;
2639
2640 unsafe {
2641 let auipc_addr = at as *mut u32;
2642 let auipc = riscv::opcodes::Inst::new(riscv::Opcode::AUIPC)
2643 .encode()
2644 .set_imm20(hi20 as _)
2645 .value;
2646 write_unaligned(auipc_addr, auipc_addr.read_unaligned() | auipc);
2647
2648 let jalr_addr = at.add(4) as *mut u32;
2649 let jalr = riscv::opcodes::Inst::new(riscv::Opcode::JALR)
2650 .encode()
2651 .set_imm12(lo12 as _)
2652 .value;
2653 write_unaligned(jalr_addr, jalr_addr.read_unaligned() | jalr);
2654 }
2655 }
2656 }
2657
2658 Reloc::Aarch64AdrPrelPgHi21 => {
2659 let what = resolve(&get_address)?;
2660 let pages = ((what & !0xfff) as i128 - (atrx & !0xfff) as i128) >> 12;
2661 if !(-(1i128 << 20)..(1i128 << 20)).contains(&pages) {
2662 return Err(AsmError::TooLarge);
2663 }
2664 let iptr = at as *mut u32;
2665 let imm21 = pages as u32 & 0x1f_ffff;
2666 let lo = (imm21 & 0x3) << 29;
2667 let hi = ((imm21 >> 2) & 0x7ffff) << 5;
2668 unsafe {
2669 let insn = iptr.read_unaligned();
2670 write_unaligned(iptr, insn | lo | hi);
2671 }
2672 }
2673
2674 Reloc::Aarch64AddAbsLo12Nc => {
2675 let what = resolve(&get_address)?;
2676 let iptr = at as *mut u32;
2677 let imm12 = (what as u32 & 0xfff) << 10;
2678 unsafe {
2679 let insn = iptr.read_unaligned();
2680 write_unaligned(iptr, insn | imm12);
2681 }
2682 }
2683
2684 Reloc::Aarch64AdrGotPage21 => {
2685 let what = resolve(&get_got_entry)?;
2686 let pages = ((what & !0xfff) as i128 - (atrx & !0xfff) as i128) >> 12;
2687 if !(-(1i128 << 20)..(1i128 << 20)).contains(&pages) {
2688 return Err(AsmError::TooLarge);
2689 }
2690 let iptr = at as *mut u32;
2691 let imm21 = pages as u32 & 0x1f_ffff;
2692 let lo = (imm21 & 0x3) << 29;
2693 let hi = ((imm21 >> 2) & 0x7ffff) << 5;
2694 unsafe {
2695 let insn = iptr.read_unaligned();
2696 write_unaligned(iptr, insn | lo | hi);
2697 }
2698 }
2699
2700 Reloc::Aarch64Ld64GotLo12Nc => {
2701 let what = resolve(&get_got_entry)?;
2702 if what & 7 != 0 {
2703 return Err(AsmError::InvalidArgument);
2704 }
2705 let iptr = at as *mut u32;
2706 let imm12 = ((what as u32 & 0xfff) >> 3) << 10;
2707 unsafe {
2708 let insn = iptr.read_unaligned();
2709 write_unaligned(iptr, insn | imm12);
2710 }
2711 }
2712
2713 _ => return Err(AsmError::InvalidArgument),
2714 }
2715 }
2716 Ok(())
2717}
2718
2719#[cfg(test)]
2720mod tests {
2721 use super::*;
2722
2723 fn resolved(address: usize) -> impl Fn(&RelocTarget) -> Result<*const u8, AsmError> {
2724 move |_| Ok(address as *const u8)
2725 }
2726
2727 fn unresolved(_: &RelocTarget) -> Result<*const u8, AsmError> {
2728 Ok(core::ptr::null())
2729 }
2730
2731 #[test]
2732 fn relocation_rejects_patch_outside_code() {
2733 let mut code = [0u8; 4];
2734 let relocs = [AsmReloc {
2735 offset: 1,
2736 kind: Reloc::Abs4,
2737 addend: 0,
2738 target: RelocTarget::Label(Label::from_id(0)),
2739 }];
2740
2741 let result = unsafe {
2742 perform_relocations(
2743 code.as_mut_ptr(),
2744 code.as_ptr(),
2745 code.len(),
2746 &relocs,
2747 resolved(1),
2748 resolved(1),
2749 resolved(1),
2750 )
2751 };
2752
2753 assert_eq!(result, Err(AsmError::InvalidArgument));
2754 assert_eq!(code, [0; 4]);
2755 }
2756
2757 #[test]
2758 fn relocation_rejects_null_and_overflowing_targets() {
2759 let mut code = [0u8; 8];
2760 let reloc = AsmReloc {
2761 offset: 0,
2762 kind: Reloc::Abs8,
2763 addend: 0,
2764 target: RelocTarget::Label(Label::from_id(0)),
2765 };
2766
2767 let null_result = unsafe {
2768 perform_relocations(
2769 code.as_mut_ptr(),
2770 code.as_ptr(),
2771 code.len(),
2772 core::slice::from_ref(&reloc),
2773 unresolved,
2774 unresolved,
2775 unresolved,
2776 )
2777 };
2778 assert_eq!(null_result, Err(AsmError::InvalidArgument));
2779
2780 let overflowing = AsmReloc {
2781 addend: -2,
2782 ..reloc
2783 };
2784 let overflow_result = unsafe {
2785 perform_relocations(
2786 code.as_mut_ptr(),
2787 code.as_ptr(),
2788 code.len(),
2789 core::slice::from_ref(&overflowing),
2790 resolved(1),
2791 resolved(1),
2792 resolved(1),
2793 )
2794 };
2795 assert_eq!(overflow_result, Err(AsmError::TooLarge));
2796 }
2797
2798 #[test]
2799 fn poisoned_buffer_rejects_raw_mutation_and_patch_metadata() {
2800 let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2801 buffer.write_u32(0);
2802 let bytes = buffer.data().to_vec();
2803 buffer.record_error(AsmError::InvalidOperand);
2804
2805 buffer.write_u8(1);
2806 buffer.put8(2);
2807 assert!(
2808 !buffer
2809 .add_symbol(ExternalName::user(0, 1), RelocDistance::Far)
2810 .is_valid()
2811 );
2812 assert_eq!(buffer.add_constant(3u64), Constant(u32::MAX));
2813 assert!(!buffer.get_label().is_valid());
2814 buffer.add_reloc(Reloc::Abs4, RelocTarget::Label(Label::from_id(0)), 0);
2815 let patch = buffer.try_record_patch_site(0, LabelUse::X86JmpRel32, 0);
2816
2817 assert_eq!(patch, Err(AsmError::InvalidOperand));
2818 assert_eq!(buffer.data(), bytes);
2819 assert!(buffer.relocs().is_empty());
2820 assert!(buffer.patch_sites.is_empty());
2821 assert!(matches!(buffer.finish(), Err(AsmError::InvalidOperand)));
2822 }
2823
2824 #[test]
2825 fn target_rejects_foreign_relocations_and_patch_kinds() {
2826 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2827 let label = buffer.get_label();
2828 buffer.add_reloc(Reloc::X86PCRel4, RelocTarget::Label(label), 0);
2829 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2830 assert!(buffer.relocs().is_empty());
2831
2832 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2833 buffer.write_u32(0);
2834 assert_eq!(
2835 buffer.try_record_patch_site(0, LabelUse::X86JmpRel32, 0),
2836 Err(AsmError::InvalidArch)
2837 );
2838 assert!(buffer.patch_sites.is_empty());
2839
2840 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2841 let label = buffer.get_label();
2842 buffer.use_label_at_offset(0, label, LabelUse::RVJal20);
2843 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2844 assert!(buffer.pending_fixup_records.is_empty());
2845
2846 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2847 let label = buffer.get_label();
2848 assert_eq!(
2849 buffer.emit_veneer(label, 0, LabelUse::RVJal20),
2850 Err(AsmError::InvalidArch)
2851 );
2852 assert!(buffer.data().is_empty());
2853
2854 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2855 buffer.write_u32(0);
2856 assert_eq!(
2857 buffer.try_record_patch_site(0, LabelUse::A64Ldr12, 0),
2858 Err(AsmError::UnsupportedInstruction {
2859 reason: "AArch64 LDR12 label patching is not implemented",
2860 })
2861 );
2862 }
2863
2864 #[cfg(not(feature = "x86"))]
2865 #[test]
2866 fn disabled_x86_label_use_is_rejected_without_a_fixup() {
2867 let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2868 let label = buffer.get_label();
2869 buffer.write_u32(0);
2870 let bytes = buffer.data().to_vec();
2871
2872 buffer.use_label_at_offset(0, label, LabelUse::X86JmpRel32);
2873
2874 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2875 assert_eq!(buffer.data(), bytes);
2876 assert!(buffer.pending_fixup_records.is_empty());
2877 assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2878
2879 let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2880 let label = buffer.get_label();
2881 buffer.add_reloc(Reloc::X86PCRel4, RelocTarget::Label(label), 0);
2882 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2883 assert!(buffer.relocs().is_empty());
2884 }
2885
2886 #[cfg(not(feature = "riscv"))]
2887 #[test]
2888 fn disabled_riscv_label_use_is_rejected_without_a_fixup() {
2889 let mut buffer = CodeBuffer::new(Environment::new(Arch::RISCV64));
2890 let label = buffer.get_label();
2891 buffer.write_u32(0);
2892 let bytes = buffer.data().to_vec();
2893
2894 buffer.use_label_at_offset(0, label, LabelUse::RVPCRel32);
2895
2896 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2897 assert_eq!(buffer.data(), bytes);
2898 assert!(buffer.pending_fixup_records.is_empty());
2899 assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2900
2901 let mut buffer = CodeBuffer::new(Environment::new(Arch::RISCV64));
2902 let label = buffer.get_label();
2903 buffer.add_reloc(Reloc::RiscvCallPlt, RelocTarget::Label(label), 0);
2904 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2905 assert!(buffer.relocs().is_empty());
2906 }
2907
2908 #[cfg(not(feature = "aarch64"))]
2909 #[test]
2910 fn disabled_aarch64_label_use_is_rejected_without_a_fixup() {
2911 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2912 let label = buffer.get_label();
2913 buffer.write_u32(0);
2914 let bytes = buffer.data().to_vec();
2915
2916 buffer.use_label_at_offset(0, label, LabelUse::A64Branch26);
2917
2918 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2919 assert_eq!(buffer.data(), bytes);
2920 assert!(buffer.pending_fixup_records.is_empty());
2921 assert_eq!(buffer.finish().err(), Some(AsmError::InvalidArch));
2922
2923 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2924 let label = buffer.get_label();
2925 buffer.add_reloc(Reloc::Arm64Call, RelocTarget::Label(label), 0);
2926 assert_eq!(buffer.error(), Some(&AsmError::InvalidArch));
2927 assert!(buffer.relocs().is_empty());
2928 }
2929
2930 #[test]
2931 fn poisoned_buffer_does_not_consume_pending_island_state() {
2932 let mut buffer = CodeBuffer::new(Environment::new(Arch::AArch64));
2933 buffer.add_constant(7u64);
2934 let pending_constants = buffer.pending_constants.len();
2935 buffer.record_error(AsmError::InvalidOperand);
2936
2937 assert_eq!(buffer.emit_island(0), Err(AsmError::InvalidOperand));
2938 assert_eq!(buffer.pending_constants.len(), pending_constants);
2939 assert!(buffer.data().is_empty());
2940 }
2941
2942 #[cfg(feature = "jit")]
2943 #[test]
2944 fn allocate_resolved_rejects_unresolved_symbols() {
2945 let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2946 let symbol = buffer.add_symbol(ExternalName::user(0, 7), RelocDistance::Far);
2947 buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(symbol), 0);
2948 buffer.write_u64(0);
2949 let code = buffer.finish().unwrap();
2950 let mut allocator = JitAllocator::new(Default::default());
2951
2952 let result = code.allocate_resolved(&mut allocator, |_| core::ptr::null());
2953
2954 assert_eq!(result.err(), Some(AsmError::InvalidArgument));
2955 }
2956
2957 #[cfg(feature = "jit")]
2958 #[test]
2959 fn allocate_resolved_resolves_user_external_names() {
2960 let mut buffer = CodeBuffer::new(Environment::new(Arch::X64));
2961 let symbol = buffer.extern_user(1, 2, RelocDistance::Far);
2962 buffer.add_reloc(Reloc::Abs8, RelocTarget::Sym(symbol), 0);
2963 buffer.write_u64(0);
2964 let code = buffer.finish().unwrap();
2965 let mut allocator = JitAllocator::new(Default::default());
2966
2967 let target = 0x1000usize as *const u8;
2969 let loaded = code
2970 .allocate_resolved(&mut allocator, |name| match name {
2971 ExternalName::User(u) if u.namespace == 1 && u.index == 2 => target,
2972 _ => core::ptr::null(),
2973 })
2974 .unwrap();
2975
2976 unsafe {
2977 let patched = core::ptr::read_unaligned(loaded.rx() as *const usize);
2978 assert_eq!(patched, target as usize);
2979 }
2980 }
2981
2982 #[test]
2983 fn extern_sym_deduplicates_by_name() {
2984 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
2985 let first = buf.extern_sym("puts", RelocDistance::Far);
2986 let other = buf.extern_sym("printf", RelocDistance::Near);
2987 let again = buf.extern_sym("puts", RelocDistance::Near);
2988
2989 assert_eq!(first, again);
2990 assert_ne!(first, other);
2991 assert_eq!(buf.symbol_distance(first), Some(RelocDistance::Far));
2993 assert_eq!(buf.symbol_distance(other), Some(RelocDistance::Near));
2994 }
2995
2996 #[test]
2997 fn extern_user_deduplicates_by_namespace_and_index() {
2998 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
2999 let first = buf.extern_user(0, 1, RelocDistance::Far);
3000 let other_ns = buf.extern_user(1, 1, RelocDistance::Near);
3001 let other_idx = buf.extern_user(0, 2, RelocDistance::Near);
3002 let again = buf.extern_user(0, 1, RelocDistance::Near);
3003
3004 assert_eq!(first, again);
3005 assert_ne!(first, other_ns);
3006 assert_ne!(first, other_idx);
3007 assert_eq!(buf.symbol_distance(first), Some(RelocDistance::Far));
3008 assert_eq!(
3009 buf.symbol_name(first),
3010 Some(&ExternalName::user(0, 1))
3011 );
3012 }
3013
3014 #[test]
3015 fn unknown_symbols_are_fallible() {
3016 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3017 let missing = Sym::from_id(u32::MAX);
3018
3019 assert_eq!(buf.symbol_name(missing), None);
3020 assert_eq!(buf.symbol_distance(missing), None);
3021
3022 let finalized = buf.finish().unwrap();
3023 assert_eq!(finalized.symbol_name(missing), None);
3024 assert_eq!(finalized.symbol_distance(missing), None);
3025 }
3026
3027 #[test]
3028 fn defined_symbols_are_resolved_at_finish() {
3029 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3030 buf.write_u8(0x90);
3031 let entry = buf.get_label();
3032 buf.bind_label(entry);
3033 buf.bind_symbol("entry", entry);
3034 buf.write_u8(0xC3);
3035
3036 let result = buf.finish().unwrap();
3037 assert_eq!(result.defined_symbol_str("entry"), Some(1));
3038 assert_eq!(result.defined_symbol_str("missing"), None);
3039 assert_eq!(
3040 result.defined_symbol_offset(&ExternalName::from("entry")),
3041 Some(1)
3042 );
3043 }
3044
3045 #[test]
3046 fn invalid_label_binding_is_reported_without_mutating_labels() {
3047 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3048 let invalid = Label::from_id(0);
3049
3050 assert_eq!(buf.try_bind_label(invalid), Err(AsmError::InvalidArgument));
3051 assert_eq!(buf.label_count(), 0);
3052
3053 buf.bind_label(invalid);
3054 assert_eq!(buf.error(), Some(&AsmError::InvalidArgument));
3055 assert!(matches!(buf.finish(), Err(AsmError::InvalidArgument)));
3056 }
3057
3058 #[cfg(feature = "x86")]
3059 #[test]
3060 fn invalid_patch_registration_does_not_create_metadata() {
3061 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3062
3063 assert_eq!(
3064 buf.try_record_patch_site(0, LabelUse::X86JmpRel32, 0),
3065 Err(AsmError::InvalidArgument)
3066 );
3067 assert!(buf.patch_sites.is_empty());
3068
3069 buf.record_patch_site(0, LabelUse::X86JmpRel32, 0);
3070 assert_eq!(buf.error(), Some(&AsmError::InvalidArgument));
3071 assert!(buf.patch_sites.is_empty());
3072 }
3073
3074 #[cfg(feature = "x86")]
3075 #[test]
3076 fn finish_rejects_an_unbound_fixup() {
3077 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3078 let label = buf.get_label();
3079 buf.write_u32(0);
3080 buf.use_label_at_offset(0, label, LabelUse::X86JmpRel32);
3081
3082 assert!(matches!(buf.finish(), Err(AsmError::UnboundLabel)));
3083 }
3084
3085 #[cfg(feature = "x86")]
3086 #[test]
3087 fn unsupported_veneer_leaves_buffer_unchanged() {
3088 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3089 let label = buf.get_label();
3090 buf.write_u32(0);
3091 let original = buf.data().to_vec();
3092
3093 assert_eq!(
3094 buf.emit_veneer(label, 0, LabelUse::X86JmpRel32),
3095 Err(AsmError::UnsupportedInstruction {
3096 reason: "branch range requires an unsupported veneer",
3097 })
3098 );
3099 assert_eq!(buf.data(), original);
3100 }
3101
3102 #[test]
3103 fn branch_ranges_include_the_last_encodable_offset() {
3104 for (kind, range) in [
3105 (LabelUse::A64Branch14, 1 << 15),
3106 (LabelUse::A64Branch19, 1 << 20),
3107 (LabelUse::A64Branch26, 1 << 27),
3108 ] {
3109 assert!(kind.can_reach(0, range - 4));
3110 assert!(!kind.can_reach(0, range));
3111 assert!(kind.can_reach(range, 0));
3112 assert!(!kind.can_reach(range + 4, 0));
3113 }
3114
3115 assert!(LabelUse::RVB12.can_reach(0, (1 << 12) - 2));
3116 assert!(!LabelUse::RVB12.can_reach(0, 1 << 12));
3117 assert!(LabelUse::RVB12.can_reach(1 << 12, 0));
3118 assert!(!LabelUse::RVB12.can_reach((1 << 12) + 2, 0));
3119 }
3120
3121 #[test]
3122 fn supported_veneer_ranges_cover_both_exact_boundaries() {
3123 for (kind, positive, negative, step) in [
3124 (LabelUse::RVJal20, (1 << 20) - 2, 1 << 20, 2),
3125 (LabelUse::RVB12, (1 << 12) - 2, 1 << 12, 2),
3126 (LabelUse::RVCJump, (1 << 11) - 2, 1 << 11, 2),
3127 ] {
3128 assert!(kind.supports_veneer());
3129 assert!(kind.can_reach(0, positive));
3130 assert!(!kind.can_reach(0, positive + step));
3131 assert!(kind.can_reach(negative, 0));
3132 assert!(!kind.can_reach(negative + step, 0));
3133 }
3134 }
3135
3136 #[cfg(feature = "riscv")]
3137 #[test]
3138 fn riscv_veneer_is_finalized_as_a_fixup() {
3139 let mut buf = CodeBuffer::new(Environment::new(Arch::RISCV64));
3140 let label = buf.get_label();
3141 buf.bind_label(label);
3142 buf.write_u32(0);
3143
3144 buf.emit_veneer(label, 0, LabelUse::RVB12).unwrap();
3145 let code = buf.finish().unwrap();
3146 assert_eq!(code.data().len(), 12);
3147 assert_eq!(&code.data()[0..4], &[0x00, 0x02, 0x00, 0x00]);
3149 }
3150
3151 #[cfg(feature = "riscv")]
3152 #[test]
3153 fn riscv_supported_veneers_are_inserted_by_islands() {
3154 for kind in [LabelUse::RVJal20, LabelUse::RVB12, LabelUse::RVCJump] {
3155 let mut buf = CodeBuffer::new(Environment::new(Arch::RISCV64));
3156 let label = buf.get_label();
3157 let patch_size = kind.patch_size();
3158 buf.get_appended_space(patch_size);
3159 buf.use_label_at_offset(0, label, kind);
3160
3161 buf.emit_island(u32::MAX).unwrap();
3162 buf.bind_label(label);
3163 let code = buf.finish().unwrap();
3164
3165 assert_eq!(code.data().len(), 12);
3166 assert!(kind.can_reach(0, 4));
3167 }
3168 }
3169
3170 #[cfg(feature = "riscv")]
3171 #[test]
3172 fn riscv_large_images_hit_exact_branch_boundaries() {
3173 for (kind, positive, step) in [
3174 (LabelUse::RVJal20, (1 << 20) - 2, 2),
3175 (LabelUse::RVB12, (1 << 12) - 2, 2),
3176 (LabelUse::RVCJump, (1 << 11) - 2, 2),
3177 ] {
3178 let mut exact = CodeBuffer::new(Environment::new(Arch::RISCV64));
3179 let label = exact.get_label();
3180 exact.get_appended_space(kind.patch_size());
3181 exact.use_label_at_offset(0, label, kind);
3182 exact.get_appended_space(positive as usize - kind.patch_size());
3183 exact.bind_label(label);
3184 assert!(exact.finish().is_ok());
3185
3186 let mut outside = CodeBuffer::new(Environment::new(Arch::RISCV64));
3187 let label = outside.get_label();
3188 outside.get_appended_space(kind.patch_size());
3189 outside.use_label_at_offset(0, label, kind);
3190 outside.get_appended_space((positive + step) as usize - kind.patch_size());
3191 outside.bind_label(label);
3192 assert_eq!(outside.finish().err(), Some(AsmError::TooLarge));
3193 }
3194 }
3195
3196 #[cfg(feature = "aarch64")]
3197 #[test]
3198 fn out_of_range_aarch64_branch_reports_unsupported_veneer() {
3199 let mut buf = CodeBuffer::new(Environment::new(Arch::AArch64));
3200 let label = buf.get_label();
3201 buf.write_u32(0);
3202 buf.use_label_at_offset(0, label, LabelUse::A64Branch14);
3203 buf.get_appended_space((1 << 15) - 4);
3204 buf.bind_label(label);
3205
3206 assert_eq!(
3207 buf.finish().err(),
3208 Some(AsmError::UnsupportedInstruction {
3209 reason: "AArch64 branch veneers are not implemented",
3210 })
3211 );
3212 }
3213
3214 #[cfg(feature = "x86")]
3215 #[test]
3216 fn x86_branch_relaxation_honors_rel8_boundaries() {
3217 use crate::x86::{Assembler, JEmitter, JmpEmitter};
3218
3219 let forward = |padding: usize, conditional: bool| {
3220 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3221 let target = buf.get_label();
3222 {
3223 let mut asm = Assembler::new(&mut buf);
3224 if conditional {
3225 asm.jz(target);
3226 } else {
3227 asm.jmp(target);
3228 }
3229 }
3230 for _ in 0..padding {
3231 buf.write_u8(0x90);
3232 }
3233 buf.bind_label(target);
3234 buf.finish().unwrap().data().to_vec()
3235 };
3236 for conditional in [false, true] {
3237 let short = forward(127, conditional);
3238 assert_eq!(short[0], if conditional { 0x74 } else { 0xEB });
3239 assert_eq!(short[1], 127);
3240
3241 let near = forward(128, conditional);
3242 if conditional {
3243 assert_eq!(&near[..2], &[0x0F, 0x84]);
3244 } else {
3245 assert_eq!(near[0], 0xE9);
3246 }
3247 }
3248
3249 let backward = |padding: usize, conditional: bool| {
3250 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3251 let target = buf.get_label();
3252 buf.bind_label(target);
3253 for _ in 0..padding {
3254 buf.write_u8(0x90);
3255 }
3256 {
3257 let mut asm = Assembler::new(&mut buf);
3258 if conditional {
3259 asm.jz(target);
3260 } else {
3261 asm.jmp(target);
3262 }
3263 }
3264 buf.finish().unwrap().data().to_vec()
3265 };
3266 for conditional in [false, true] {
3267 let short = backward(126, conditional);
3268 assert_eq!(
3269 &short[126..],
3270 &[if conditional { 0x74 } else { 0xEB }, 0x80]
3271 );
3272
3273 let near = backward(127, conditional);
3274 assert_eq!(near[127], if conditional { 0x0F } else { 0xE9 });
3275 }
3276 }
3277
3278 #[cfg(feature = "x86")]
3279 #[test]
3280 fn x86_branch_relaxation_reaches_a_bounded_fixed_point() {
3281 use crate::x86::{Assembler, JmpEmitter};
3282
3283 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3284 let target = buf.get_label();
3285 {
3286 let mut asm = Assembler::new(&mut buf);
3287 asm.jmp(target);
3288 }
3289 for _ in 0..120 {
3290 buf.write_u8(0x90);
3291 }
3292 {
3293 let mut asm = Assembler::new(&mut buf);
3294 asm.jmp(target);
3295 }
3296 for _ in 0..3 {
3297 buf.write_u8(0x90);
3298 }
3299 buf.bind_label(target);
3300
3301 let code = buf.finish().unwrap();
3302 assert_eq!(&code.data()[..2], &[0xEB, 125]);
3303 assert_eq!(&code.data()[122..124], &[0xEB, 3]);
3304 assert_eq!(code.label_offsets[target.id() as usize], 127);
3305 }
3306
3307 #[cfg(feature = "x86")]
3308 #[test]
3309 fn x86_branch_relaxation_rebases_metadata_and_is_deterministic() {
3310 use crate::x86::{Assembler, JmpEmitter};
3311
3312 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3313 let target = buf.get_label();
3314 {
3315 let mut asm = Assembler::new(&mut buf);
3316 asm.jmp(target);
3317 }
3318 buf.write_u8(0x90);
3319 buf.bind_label(target);
3320 buf.bind_symbol("target", target);
3321
3322 buf.add_reloc(Reloc::Abs4, RelocTarget::Label(target), 0);
3323 buf.write_u32(0);
3324 let block = buf.reserve_patch_block(4, 1).unwrap();
3325 let block_offset = buf.patch_blocks[block.index()].offset;
3326 let site = buf
3327 .try_record_patch_site(
3328 block_offset,
3329 LabelUse::X86JmpRel32,
3330 buf.label_offset(target),
3331 )
3332 .unwrap();
3333
3334 let later = buf.get_label();
3335 {
3336 let mut asm = Assembler::new(&mut buf);
3337 asm.long().jmp(later);
3338 }
3339 buf.write_u8(0x90);
3340 buf.bind_label(later);
3341
3342 let first = buf.finish().unwrap();
3343 assert_eq!(&first.data()[..3], &[0xEB, 1, 0x90]);
3344 assert_eq!(first.defined_symbol_str("target"), Some(3));
3345 assert_eq!(first.relocs()[0].offset, 3);
3346 assert_eq!(first.patch_catalog().block(block).unwrap().offset, 7);
3347 let patch_site = first.patch_catalog().site(site).unwrap();
3348 assert_eq!(patch_site.offset, 7);
3349 assert_eq!(patch_site.current_target, 3);
3350 assert_eq!(&first.data()[11..], &[0xE9, 1, 0, 0, 0, 0x90]);
3351
3352 let second = buf.finish().unwrap();
3353 assert_eq!(second.data(), first.data());
3354 assert_eq!(second.relocs(), first.relocs());
3355 assert_eq!(second.patch_catalog(), first.patch_catalog());
3356 }
3357
3358 #[cfg(feature = "x86")]
3359 #[test]
3360 fn x86_branch_relaxation_preserves_recorded_alignment() {
3361 use crate::x86::{Assembler, JmpEmitter};
3362
3363 let mut buf = CodeBuffer::new(Environment::new(Arch::X64));
3364 let target = buf.get_label();
3365 {
3366 let mut asm = Assembler::new(&mut buf);
3367 asm.jmp(target);
3368 }
3369 buf.try_align_to(16).unwrap();
3370 buf.bind_label(target);
3371
3372 let code = buf.finish().unwrap();
3373 assert_eq!(code.data()[0], 0xE9);
3374 assert_eq!(code.label_offsets[target.id() as usize], 16);
3375 }
3376}