1#![cfg(all(target_arch = "x86_64", not(target_os = "windows")))]
7#![doc = include_str!("../README.md")]
8
9use std::fmt;
10use std::mem;
11
12pub mod buffer;
13pub mod jit;
14pub mod patch;
15#[cfg(target_arch = "x86_64")]
16pub mod regalloc;
17pub mod segv_trace;
18pub mod vectorize;
19#[cfg(target_arch = "x86_64")]
20pub mod x64asm;
21mod stencil_model;
22pub use stencil_model::{HoleId, Reloc, RelocKind, Stencil};
23
24include!(concat!(env!("OUT_DIR"), "/stencils.rs"));
27
28#[derive(Debug)]
30pub enum JitError {
31 EmptyCode,
33 Map(std::io::Error),
35 Protect(std::io::Error),
37}
38
39impl fmt::Display for JitError {
40 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41 match self {
42 JitError::EmptyCode => write!(f, "JitPage: empty code"),
43 JitError::Map(e) => write!(f, "executable mapping failed: {e}"),
44 JitError::Protect(e) => write!(f, "marking page executable failed: {e}"),
45 }
46 }
47}
48
49impl std::error::Error for JitError {}
50
51#[derive(Debug)]
53pub struct JitPage {
54 ptr: *mut u8,
55 code_len: usize,
57 alloc_len: usize,
60}
61
62unsafe impl Send for JitPage {}
67unsafe impl Sync for JitPage {}
68
69impl JitPage {
70 pub fn new(code: &[u8]) -> Result<JitPage, JitError> {
72 if code.is_empty() {
73 return Err(JitError::EmptyCode);
74 }
75 let page = page_size();
76 let alloc_len = code.len().div_ceil(page) * page;
77 let ptr = map_executable(alloc_len)?;
78 debug_assert_eq!(
79 ptr as usize % page,
80 0,
81 "executable mapping is not page-aligned"
82 );
83 unsafe {
84 write_code(ptr, code)?;
85 }
86 Ok(JitPage { ptr, code_len: code.len(), alloc_len })
87 }
88
89 #[cfg(all(unix, not(all(target_os = "macos", target_arch = "aarch64"))))]
95 pub fn patch_word(&self, offset: usize, value: u64) -> Result<(), JitError> {
96 assert!(offset + 8 <= self.code_len, "patch outside the mapping");
97 let page = page_size();
98 let start = (self.ptr as usize + offset) & !(page - 1);
99 let end = self.ptr as usize + offset + 8;
100 let len = end - start;
101 unsafe {
102 let rc = libc::mprotect(
103 start as *mut libc::c_void,
104 len,
105 libc::PROT_READ | libc::PROT_WRITE,
106 );
107 if rc != 0 {
108 return Err(JitError::Protect(std::io::Error::last_os_error()));
109 }
110 std::ptr::write_unaligned(self.ptr.add(offset) as *mut u64, value);
111 let rc = libc::mprotect(
112 start as *mut libc::c_void,
113 len,
114 libc::PROT_READ | libc::PROT_EXEC,
115 );
116 if rc != 0 {
117 return Err(JitError::Protect(std::io::Error::last_os_error()));
118 }
119 }
120 Ok(())
121 }
122
123 #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
126 pub fn patch_word(&self, _offset: usize, _value: u64) -> Result<(), JitError> {
127 Err(JitError::EmptyCode)
128 }
129
130 pub fn with_layout(
134 len: usize,
135 fill: impl FnOnce(u64) -> Vec<u8>,
136 ) -> Result<JitPage, JitError> {
137 if len == 0 {
138 return Err(JitError::EmptyCode);
139 }
140 let page = page_size();
141 let alloc_len = len.div_ceil(page) * page;
142 let ptr = map_executable(alloc_len)?;
143 let code = fill(ptr as u64);
144 debug_assert_eq!(code.len(), len, "fill returned a different length");
145 unsafe {
146 write_code(ptr, &code)?;
147 }
148 Ok(JitPage { ptr, code_len: len, alloc_len })
149 }
150
151 pub unsafe fn as_fn_i64_i64(&self) -> extern "C" fn(i64, i64) -> i64 {
158 mem::transmute::<*mut u8, extern "C" fn(i64, i64) -> i64>(self.ptr)
159 }
160
161 pub fn as_ptr(&self) -> *const u8 {
163 self.ptr
164 }
165
166 pub fn len(&self) -> usize {
169 self.code_len
170 }
171
172 pub fn alloc_len(&self) -> usize {
174 self.alloc_len
175 }
176
177 pub fn is_empty(&self) -> bool {
179 self.code_len == 0
180 }
181}
182
183impl Drop for JitPage {
184 fn drop(&mut self) {
185 unsafe {
186 #[cfg(unix)]
187 {
188 let rc = libc::munmap(self.ptr as *mut libc::c_void, self.alloc_len);
189 debug_assert_eq!(rc, 0, "munmap failed for a JitPage mapping");
190 }
191 #[cfg(windows)]
192 {
193 let rc = windows_sys::Win32::System::Memory::VirtualFree(
194 self.ptr as *mut core::ffi::c_void,
195 0,
196 windows_sys::Win32::System::Memory::MEM_RELEASE,
197 );
198 debug_assert_ne!(rc, 0, "VirtualFree failed for a JitPage mapping");
199 }
200 }
201 }
202}
203
204fn page_size() -> usize {
206 #[cfg(unix)]
207 unsafe {
208 let sz = libc::sysconf(libc::_SC_PAGESIZE);
209 if sz > 0 {
210 sz as usize
211 } else {
212 4096
213 }
214 }
215 #[cfg(windows)]
216 unsafe {
217 let mut info = mem::zeroed::<windows_sys::Win32::System::SystemInformation::SYSTEM_INFO>();
218 windows_sys::Win32::System::SystemInformation::GetSystemInfo(&mut info);
219 info.dwPageSize as usize
220 }
221}
222
223#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
228mod sys {
229 pub const MAP_JIT: libc::c_int = 0x800;
232
233 extern "C" {
234 pub fn pthread_jit_write_protect_np(enabled: libc::c_int);
237 pub fn sys_icache_invalidate(start: *mut libc::c_void, len: libc::size_t);
239 }
240}
241
242#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
243fn map_executable(len: usize) -> Result<*mut u8, JitError> {
244 let ptr = unsafe {
245 libc::mmap(
246 std::ptr::null_mut(),
247 len,
248 libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
249 libc::MAP_PRIVATE | libc::MAP_ANON | sys::MAP_JIT,
250 -1,
251 0,
252 )
253 };
254 if ptr == libc::MAP_FAILED {
255 return Err(JitError::Map(std::io::Error::last_os_error()));
256 }
257 Ok(ptr as *mut u8)
258}
259
260#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
261unsafe fn write_code(dst: *mut u8, code: &[u8]) -> Result<(), JitError> {
262 sys::pthread_jit_write_protect_np(0);
264 std::ptr::copy_nonoverlapping(code.as_ptr(), dst, code.len());
265 sys::pthread_jit_write_protect_np(1);
266 sys::sys_icache_invalidate(dst as *mut libc::c_void, code.len());
268 Ok(())
269}
270
271#[cfg(all(unix, not(all(target_os = "macos", target_arch = "aarch64"))))]
276fn map_executable(len: usize) -> Result<*mut u8, JitError> {
277 let ptr = unsafe {
278 libc::mmap(
279 std::ptr::null_mut(),
280 len,
281 libc::PROT_READ | libc::PROT_WRITE,
282 libc::MAP_PRIVATE | libc::MAP_ANON,
283 -1,
284 0,
285 )
286 };
287 if ptr == libc::MAP_FAILED {
288 return Err(JitError::Map(std::io::Error::last_os_error()));
289 }
290 Ok(ptr as *mut u8)
291}
292
293#[cfg(all(unix, not(all(target_os = "macos", target_arch = "aarch64"))))]
294unsafe fn write_code(dst: *mut u8, code: &[u8]) -> Result<(), JitError> {
295 std::ptr::copy_nonoverlapping(code.as_ptr(), dst, code.len());
296 let rc = libc::mprotect(
297 dst as *mut libc::c_void,
298 code.len(),
299 libc::PROT_READ | libc::PROT_EXEC,
300 );
301 if rc != 0 {
302 return Err(JitError::Protect(std::io::Error::last_os_error()));
303 }
304 #[cfg(target_arch = "aarch64")]
306 flush_icache_aarch64(dst, code.len());
307 Ok(())
308}
309
310#[cfg(all(unix, target_arch = "aarch64", not(target_os = "macos")))]
313unsafe fn flush_icache_aarch64(start: *mut u8, len: usize) {
314 const LINE: usize = 64;
315 let begin = start as usize & !(LINE - 1);
316 let end = start as usize + len;
317 let mut addr = begin;
318 while addr < end {
319 core::arch::asm!("dc cvau, {0}", in(reg) addr, options(nostack, preserves_flags));
320 addr += LINE;
321 }
322 core::arch::asm!("dsb ish", options(nostack, preserves_flags));
323 let mut addr = begin;
324 while addr < end {
325 core::arch::asm!("ic ivau, {0}", in(reg) addr, options(nostack, preserves_flags));
326 addr += LINE;
327 }
328 core::arch::asm!("dsb ish", "isb", options(nostack, preserves_flags));
329}
330
331#[cfg(windows)]
336fn map_executable(len: usize) -> Result<*mut u8, JitError> {
337 use windows_sys::Win32::System::Memory::{
338 VirtualAlloc, MEM_COMMIT, MEM_RESERVE, PAGE_READWRITE,
339 };
340 let ptr = unsafe {
341 VirtualAlloc(
342 std::ptr::null(),
343 len,
344 MEM_COMMIT | MEM_RESERVE,
345 PAGE_READWRITE,
346 )
347 };
348 if ptr.is_null() {
349 return Err(JitError::Map(std::io::Error::last_os_error()));
350 }
351 Ok(ptr as *mut u8)
352}
353
354#[cfg(windows)]
355unsafe fn write_code(dst: *mut u8, code: &[u8]) -> Result<(), JitError> {
356 use windows_sys::Win32::System::Diagnostics::Debug::FlushInstructionCache;
357 use windows_sys::Win32::System::Memory::{VirtualProtect, PAGE_EXECUTE_READ};
358 use windows_sys::Win32::System::Threading::GetCurrentProcess;
359
360 std::ptr::copy_nonoverlapping(code.as_ptr(), dst, code.len());
361 let mut old = 0u32;
362 let rc = VirtualProtect(
363 dst as *const core::ffi::c_void,
364 code.len(),
365 PAGE_EXECUTE_READ,
366 &mut old,
367 );
368 if rc == 0 {
369 return Err(JitError::Protect(std::io::Error::last_os_error()));
370 }
371 let rc = FlushInstructionCache(
372 GetCurrentProcess(),
373 dst as *const core::ffi::c_void,
374 code.len(),
375 );
376 if rc == 0 {
377 return Err(JitError::Protect(std::io::Error::last_os_error()));
378 }
379 Ok(())
380}
381
382#[cfg(test)]
383mod tests {
384 use super::*;
385
386 #[cfg(target_arch = "aarch64")]
389 const ADD_CODE: [u8; 8] = [0x00, 0x00, 0x01, 0x8b, 0xc0, 0x03, 0x5f, 0xd6];
390
391 #[cfg(target_arch = "x86_64")]
394 const ADD_CODE: [u8; 5] = [0x48, 0x8d, 0x04, 0x37, 0xc3];
395
396 #[test]
397 fn jit_add_3_5_equals_8() {
398 let page = JitPage::new(&ADD_CODE).expect("could not allocate executable page");
399 let add = unsafe { page.as_fn_i64_i64() };
400 assert_eq!(add(3, 5), 8);
401 assert_eq!(add(100, -1), 99);
402 assert_eq!(add(0, 0), 0);
403 assert_eq!(add(i64::MAX, 1), i64::MIN);
405 }
406
407 #[test]
408 fn jit_page_reports_lengths_and_alignment() {
409 let page = JitPage::new(&ADD_CODE).unwrap();
410 assert_eq!(page.len(), ADD_CODE.len());
411 assert!(!page.is_empty());
412 assert!(!page.as_ptr().is_null());
413 let ps = super::page_size();
415 assert_eq!(page.alloc_len() % ps, 0);
416 assert!(page.alloc_len() >= page.len());
417 assert_eq!(page.as_ptr() as usize % ps, 0);
418 }
419
420 #[test]
421 fn jit_page_empty_code_is_error() {
422 match JitPage::new(&[]) {
423 Err(JitError::EmptyCode) => {}
424 other => panic!("expected EmptyCode error, got {:?}", other.map(|p| p.len())),
425 }
426 }
427
428 #[test]
429 fn jit_page_built_on_one_thread_executes_on_another() {
430 let page = JitPage::new(&ADD_CODE).unwrap();
433 let handle = std::thread::spawn(move || {
434 let add = unsafe { page.as_fn_i64_i64() };
435 add(20, 22)
436 });
437 assert_eq!(handle.join().unwrap(), 42);
438 }
439
440 #[test]
441 fn jit_many_threads_create_write_execute_drop() {
442 let handles: Vec<_> = (0..8)
444 .map(|k| {
445 std::thread::spawn(move || {
446 for _ in 0..200 {
447 let page = JitPage::new(&ADD_CODE).unwrap();
448 let add = unsafe { page.as_fn_i64_i64() };
449 assert_eq!(add(k, 1), k + 1);
450 }
451 })
452 })
453 .collect();
454 for h in handles {
455 h.join().unwrap();
456 }
457 }
458
459 #[test]
460 fn jit_create_and_drop_many_pages() {
461 for _ in 0..1000 {
464 let page = JitPage::new(&ADD_CODE).unwrap();
465 let add = unsafe { page.as_fn_i64_i64() };
466 assert_eq!(add(1, 2), 3);
467 }
468 }
469
470 fn stencil(name: &str) -> &'static Stencil {
473 super::STENCILS
474 .iter()
475 .find(|s| s.name == name)
476 .unwrap_or_else(|| panic!("stencil '{name}' missing from table"))
477 }
478
479 #[test]
480 fn stencil_table_contains_the_full_cps_set() {
481 for name in [
482 "logos_stencil_const",
483 "logos_stencil_addi",
484 "logos_stencil_subi",
485 "logos_stencil_muli",
486 "logos_stencil_lti",
487 "logos_stencil_branch_if",
488 "logos_stencil_return",
489 "logos_stencil_add",
490 ] {
491 let st = stencil(name);
492 assert!(!st.code.is_empty(), "{name} has no code");
493 }
494 }
495
496 #[test]
497 fn checked_div_mod_stencils_have_two_distinct_cont_holes() {
498 for name in ["logos_stencil_divi_checked", "logos_stencil_modi_checked"] {
499 let st = stencil(name);
500 let conts: std::collections::HashSet<_> = st
501 .relocs
502 .iter()
503 .filter_map(|r| match r.target {
504 HoleId::Cont(n) => Some(n),
505 _ => None,
506 })
507 .collect();
508 assert_eq!(
509 conts,
510 [0u8, 1u8].into_iter().collect(),
511 "{name} must expose both the success and side-exit continuations"
512 );
513 }
514 }
515
516 #[test]
517 fn deopt_stencil_has_const_hole_and_no_continuations() {
518 let st = stencil("logos_stencil_deopt");
519 assert!(
520 st.relocs.iter().any(|r| matches!(r.target, HoleId::ConstI64(0))),
521 "deopt stencil lacks its status-cell address hole: {:?}",
522 st.relocs
523 );
524 assert!(
525 !st.relocs.iter().any(|r| matches!(r.target, HoleId::Cont(_))),
526 "deopt is a terminal — it must not continue anywhere"
527 );
528 }
529
530 #[test]
531 fn checked_div_chain_computes_and_side_exits() {
532 use crate::jit::{compile_straightline, ChainOutcome, MicroOp};
533 let prog = [
535 MicroOp::Div { dst: 2, lhs: 0, rhs: 1 },
536 MicroOp::Return { src: 2 },
537 ];
538 let chain = compile_straightline(&prog).expect("compile");
539 let mut frame = [40i64, 5, 0];
540 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Return(8));
541 let mut frame = [i64::MIN, -1, 0];
544 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Deopt(1));
545 let mut frame = [40i64, 0, 0];
547 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Deopt(1));
548 let mut frame = [40i64, 4, 0];
549 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Return(10));
550 }
551
552 #[test]
553 fn checked_mod_chain_computes_and_side_exits() {
554 use crate::jit::{compile_straightline, ChainOutcome, MicroOp};
555 let prog = [
556 MicroOp::Mod { dst: 2, lhs: 0, rhs: 1 },
557 MicroOp::Return { src: 2 },
558 ];
559 let chain = compile_straightline(&prog).expect("compile");
560 let mut frame = [43i64, 5, 0];
561 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Return(3));
562 let mut frame = [i64::MIN, -1, 0];
563 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Return(0));
564 let mut frame = [-7i64, 2, 0];
565 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Return(-1));
566 let mut frame = [43i64, 0, 0];
567 assert_eq!(chain.run_with_frame(&mut frame), ChainOutcome::Deopt(1));
568 }
569
570 #[test]
578 fn float_move_between_xmm_pins_threads_the_register() {
579 use crate::jit::{compile_straightline_pinned_float, ChainOutcome, MicroOp};
580 let prog = [
581 MicroOp::MulF { dst: 2, lhs: 0, rhs: 0 }, MicroOp::Move { dst: 0, src: 2 }, MicroOp::Return { src: 0 },
584 ];
585 let chain = compile_straightline_pinned_float(&prog, &[], &[0, 2], None).expect("compile");
586 let mut frame = [3.0f64.to_bits() as i64, 0, 0];
587 assert_eq!(
588 chain.run_with_frame(&mut frame),
589 ChainOutcome::Return(9.0f64.to_bits() as i64),
590 "Move between XMM pins must copy the register value (9.0), not a stale frame cell"
591 );
592 }
593
594 #[test]
599 fn float_move_frame_into_xmm_pin_reloads() {
600 use crate::jit::{compile_straightline_pinned_float, ChainOutcome, MicroOp};
601 let prog = [
602 MicroOp::Move { dst: 0, src: 2 }, MicroOp::Return { src: 0 },
604 ];
605 let chain = compile_straightline_pinned_float(&prog, &[], &[0], None).expect("compile");
606 let mut frame = [1.0f64.to_bits() as i64, 0, 7.5f64.to_bits() as i64];
607 assert_eq!(
608 chain.run_with_frame(&mut frame),
609 ChainOutcome::Return(7.5f64.to_bits() as i64),
610 "Move from frame into an XMM pin must load 7.5 into the pin"
611 );
612 }
613
614 #[test]
615 fn return_stencil_has_zero_relocs() {
616 assert!(stencil("logos_stencil_return").relocs.is_empty());
617 assert!(stencil("logos_stencil_add").relocs.is_empty());
618 }
619
620 #[test]
621 fn binop_stencils_have_exactly_one_cont_hole() {
622 for name in ["logos_stencil_addi", "logos_stencil_subi", "logos_stencil_muli", "logos_stencil_lti"] {
623 let st = stencil(name);
624 let conts: Vec<_> = st
625 .relocs
626 .iter()
627 .filter(|r| matches!(r.target, HoleId::Cont(0)))
628 .collect();
629 assert_eq!(conts.len(), 1, "{name}: expected exactly one Cont(0) reloc, got {:?}", st.relocs);
630 }
631 }
632
633 #[test]
634 fn const_stencil_has_const_hole_and_cont_hole() {
635 let st = stencil("logos_stencil_const");
636 assert!(
637 st.relocs.iter().any(|r| matches!(r.target, HoleId::ConstI64(0))),
638 "const stencil lacks its ConstI64 hole: {:?}",
639 st.relocs
640 );
641 assert!(
642 st.relocs.iter().any(|r| matches!(r.target, HoleId::Cont(0))),
643 "const stencil lacks its continuation: {:?}",
644 st.relocs
645 );
646 }
647
648 #[test]
649 fn branch_if_has_two_distinct_cont_holes() {
650 let st = stencil("logos_stencil_branch_if");
651 let has0 = st.relocs.iter().any(|r| matches!(r.target, HoleId::Cont(0)));
652 let has1 = st.relocs.iter().any(|r| matches!(r.target, HoleId::Cont(1)));
653 assert!(has0 && has1, "branch_if needs Cont(0) and Cont(1): {:?}", st.relocs);
654 }
655
656 #[test]
657 fn every_reloc_is_in_bounds_and_aligned() {
658 for st in super::STENCILS {
659 for r in st.relocs {
660 assert!((r.offset as usize) < st.code.len(), "{}: reloc out of bounds", st.name);
661 #[cfg(target_arch = "aarch64")]
662 assert_eq!(r.offset % 4, 0, "{}: arm64 reloc misaligned", st.name);
663 }
664 }
665 }
666
667 #[test]
670 fn jit_extracted_rust_stencil_runs() {
671 assert!(!super::ADD_STENCIL.is_empty(), "stencil extraction produced no bytes");
672 let page = JitPage::new(super::ADD_STENCIL).expect("load extracted stencil");
673 let add = unsafe { page.as_fn_i64_i64() };
674 assert_eq!(add(3, 5), 8);
675 assert_eq!(add(40, 2), 42);
676 assert_eq!(add(-10, 7), -3);
677 }
678
679 #[test]
684 fn extracted_add_is_a_relocation_free_leaf() {
685 #[cfg(target_arch = "aarch64")]
686 {
687 let bytes = super::ADD_STENCIL;
688 assert!(bytes.len() <= 16, "leaf stencil unexpectedly large: {} bytes", bytes.len());
690 assert!(
691 bytes.ends_with(&[0xc0, 0x03, 0x5f, 0xd6]),
692 "stencil does not end in `ret`: {:02x?}",
693 bytes
694 );
695 assert_eq!(bytes[3], 0x8b, "first instruction is not a 64-bit ADD");
698 }
699 }
700}