1use std::collections::HashMap;
14use std::mem;
15
16use cranelift_codegen::ir::{self, AbiParam, InstBuilder, types};
17use cranelift_codegen::settings::{self, Configurable};
18use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
19use cranelift_jit::{JITBuilder, JITModule};
20use cranelift_module::{Linkage, Module};
21
22use crate::ast::PolydatNode;
23use crate::ast::SlotShape;
24
25use super::kernels::{JitCore, JitKernelPull, JitKernelPush, JitKernelPushPull, JitKernelRaw};
26
27extern "C" fn jit_xxh3_hash(value: u64) -> u64 {
31 guarded(|| xxhash_rust::xxh3::xxh3_64(&value.to_le_bytes()))
32}
33
34extern "C" fn jit_interleave(a: u64, b: u64) -> u64 {
36 guarded(|| {
37 let mut result: u64 = 0;
38 for i in 0..32 {
39 result |= ((a >> i) & 1) << (2 * i);
40 result |= ((b >> i) & 1) << (2 * i + 1);
41 }
42 result
43 })
44}
45
46extern "C" fn jit_lut_sample(input_bits: u64, lut_ptr: u64, lut_len: u64) -> u64 {
51 guarded(|| {
52 let u = f64::from_bits(input_bits).clamp(0.0, 1.0);
53 let n = (lut_len - 1) as f64;
54 let pos = u * n;
55 let idx = (pos as usize).min(lut_len as usize - 2);
56 let frac = pos - idx as f64;
57 let result = unsafe {
58 let ptr = lut_ptr as *const f64;
59 let a = *ptr.add(idx);
60 let b = *ptr.add(idx + 1);
61 a * (1.0 - frac) + b * frac
62 };
63 result.to_bits()
64 })
65}
66
67extern "C" fn jit_shuffle(input: u64, feedback: u64, size: u64, min: u64) -> u64 {
69 guarded(|| {
70 let mut register = (input % size) + 1;
71 loop {
72 let lsb = register & 1;
73 register >>= 1;
74 if lsb != 0 {
75 register ^= feedback;
76 }
77 if register <= size {
78 break;
79 }
80 }
81 (register - 1) + min
82 })
83}
84
85extern "C" fn jit_sin(bits: u64) -> u64 {
87 guarded(|| f64::from_bits(bits).sin().to_bits())
88}
89extern "C" fn jit_cos(bits: u64) -> u64 {
90 guarded(|| f64::from_bits(bits).cos().to_bits())
91}
92extern "C" fn jit_tan(bits: u64) -> u64 {
93 guarded(|| f64::from_bits(bits).tan().to_bits())
94}
95extern "C" fn jit_asin(bits: u64) -> u64 {
96 guarded(|| f64::from_bits(bits).asin().to_bits())
97}
98extern "C" fn jit_acos(bits: u64) -> u64 {
99 guarded(|| f64::from_bits(bits).acos().to_bits())
100}
101extern "C" fn jit_atan(bits: u64) -> u64 {
102 guarded(|| f64::from_bits(bits).atan().to_bits())
103}
104extern "C" fn jit_sqrt(bits: u64) -> u64 {
105 guarded(|| f64::from_bits(bits).sqrt().to_bits())
106}
107extern "C" fn jit_abs_f64(bits: u64) -> u64 {
108 guarded(|| f64::from_bits(bits).abs().to_bits())
109}
110extern "C" fn jit_ln(bits: u64) -> u64 {
111 guarded(|| f64::from_bits(bits).ln().to_bits())
112}
113extern "C" fn jit_exp(bits: u64) -> u64 {
114 guarded(|| f64::from_bits(bits).exp().to_bits())
115}
116extern "C" fn jit_floor_base10(bits: u64) -> u64 {
117 guarded(|| {
118 use crate::numeric::round_numbers::*;
119 let x = f64::from_bits(bits);
120 let r = if !positive_finite(x) {
121 0.0
122 } else {
123 floor_pow10(x)
124 };
125 r.to_bits()
126 })
127}
128extern "C" fn jit_ceiling_base10(bits: u64) -> u64 {
129 guarded(|| {
130 use crate::numeric::round_numbers::*;
131 let x = f64::from_bits(bits);
132 let r = if !positive_finite(x) {
133 0.0
134 } else {
135 let lo = floor_pow10(x);
136 if lo == x { lo } else { lo * 10.0 }
137 };
138 r.to_bits()
139 })
140}
141extern "C" fn jit_closest_base10(bits: u64) -> u64 {
142 guarded(|| {
143 use crate::numeric::round_numbers::*;
144 let x = f64::from_bits(bits);
145 let r = if !positive_finite(x) {
146 0.0
147 } else {
148 let lo = floor_pow10(x);
149 let hi = if lo == x { lo } else { lo * 10.0 };
150 pick_closest(x, lo, hi)
151 };
152 r.to_bits()
153 })
154}
155extern "C" fn jit_floor_decade(bits: u64) -> u64 {
156 guarded(|| {
157 use crate::numeric::round_numbers::*;
158 let x = f64::from_bits(bits);
159 let r = if !positive_finite(x) {
160 0.0
161 } else {
162 let base = floor_pow10(x);
163 (x / base).floor() * base
164 };
165 r.to_bits()
166 })
167}
168extern "C" fn jit_ceiling_decade(bits: u64) -> u64 {
169 guarded(|| {
170 use crate::numeric::round_numbers::*;
171 let x = f64::from_bits(bits);
172 let r = if !positive_finite(x) {
173 0.0
174 } else {
175 let base = floor_pow10(x);
176 (x / base).ceil() * base
177 };
178 r.to_bits()
179 })
180}
181extern "C" fn jit_closest_decade(bits: u64) -> u64 {
182 guarded(|| {
183 use crate::numeric::round_numbers::*;
184 let x = f64::from_bits(bits);
185 let r = if !positive_finite(x) {
186 0.0
187 } else {
188 let base = floor_pow10(x);
189 (x / base).round() * base
190 };
191 r.to_bits()
192 })
193}
194extern "C" fn jit_floor_binomial(bits: u64) -> u64 {
195 guarded(|| {
196 use crate::numeric::round_numbers::*;
197 let x = f64::from_bits(bits);
198 let r = if !positive_finite(x) {
199 0.0
200 } else {
201 floor_pow2(x)
202 };
203 r.to_bits()
204 })
205}
206extern "C" fn jit_ceiling_binomial(bits: u64) -> u64 {
207 guarded(|| {
208 use crate::numeric::round_numbers::*;
209 let x = f64::from_bits(bits);
210 let r = if !positive_finite(x) {
211 0.0
212 } else {
213 let lo = floor_pow2(x);
214 if lo == x { lo } else { lo * 2.0 }
215 };
216 r.to_bits()
217 })
218}
219extern "C" fn jit_closest_binomial(bits: u64) -> u64 {
220 guarded(|| {
221 use crate::numeric::round_numbers::*;
222 let x = f64::from_bits(bits);
223 let r = if !positive_finite(x) {
224 0.0
225 } else {
226 let lo = floor_pow2(x);
227 let hi = if lo == x { lo } else { lo * 2.0 };
228 pick_closest(x, lo, hi)
229 };
230 r.to_bits()
231 })
232}
233extern "C" fn jit_floor_fibonacci(bits: u64) -> u64 {
234 guarded(|| {
235 use crate::numeric::round_numbers::*;
236 let x = f64::from_bits(bits);
237 let r = if !positive_finite(x) {
238 0.0
239 } else {
240 floor_fibonacci_val(x)
241 };
242 r.to_bits()
243 })
244}
245extern "C" fn jit_ceiling_fibonacci(bits: u64) -> u64 {
246 guarded(|| {
247 use crate::numeric::round_numbers::*;
248 let x = f64::from_bits(bits);
249 let r = if !positive_finite(x) {
250 0.0
251 } else {
252 ceiling_fibonacci_val(x)
253 };
254 r.to_bits()
255 })
256}
257extern "C" fn jit_closest_fibonacci(bits: u64) -> u64 {
258 guarded(|| {
259 use crate::numeric::round_numbers::*;
260 let x = f64::from_bits(bits);
261 let r = if !positive_finite(x) {
262 0.0
263 } else {
264 pick_closest(x, floor_fibonacci_val(x), ceiling_fibonacci_val(x))
265 };
266 r.to_bits()
267 })
268}
269
270extern "C" fn jit_atan2(y_bits: u64, x_bits: u64) -> u64 {
271 guarded(|| {
272 f64::from_bits(y_bits)
273 .atan2(f64::from_bits(x_bits))
274 .to_bits()
275 })
276}
277extern "C" fn jit_pow(base_bits: u64, exp_bits: u64) -> u64 {
278 guarded(|| {
279 f64::from_bits(base_bits)
280 .powf(f64::from_bits(exp_bits))
281 .to_bits()
282 })
283}
284extern "C" fn jit_round_nearest(x_bits: u64, iv_bits: u64) -> u64 {
285 guarded(|| {
286 let x = f64::from_bits(x_bits);
287 let interval = f64::from_bits(iv_bits);
288 let r = if !(interval.is_finite() && interval > 0.0) {
289 x
290 } else {
291 (x / interval).round() * interval
292 };
293 r.to_bits()
294 })
295}
296extern "C" fn jit_round_floor(x_bits: u64, iv_bits: u64) -> u64 {
297 guarded(|| {
298 let x = f64::from_bits(x_bits);
299 let interval = f64::from_bits(iv_bits);
300 let r = if !(interval.is_finite() && interval > 0.0) {
301 x
302 } else {
303 (x / interval).floor() * interval
304 };
305 r.to_bits()
306 })
307}
308extern "C" fn jit_round_ceiling(x_bits: u64, iv_bits: u64) -> u64 {
309 guarded(|| {
310 let x = f64::from_bits(x_bits);
311 let interval = f64::from_bits(iv_bits);
312 let r = if !(interval.is_finite() && interval > 0.0) {
313 x
314 } else {
315 (x / interval).ceil() * interval
316 };
317 r.to_bits()
318 })
319}
320
321extern "C" fn jit_pcg(input: u64, seed: u64, stream: u64) -> u64 {
322 guarded(|| {
323 let inc = 2u64.wrapping_mul(stream).wrapping_add(1);
324 crate::numeric::pcg::pcg_seek(seed, inc, input)
325 })
326}
327extern "C" fn jit_pcg_stream(input: u64, stream: u64, seed: u64) -> u64 {
328 guarded(|| {
329 let inc = 2u64.wrapping_mul(stream).wrapping_add(1);
330 crate::numeric::pcg::pcg_seek(seed, inc, input)
331 })
332}
333extern "C" fn jit_n_of(input: u64, n: u64, m: u64) -> u64 {
334 guarded(|| {
335 if m == 0 {
336 return 0;
337 }
338 crate::numeric::n_of_m::n_of_m_eval(input, n, m)
339 })
340}
341
342extern "C" fn jit_cycle_walk(pos: u64, range: u64, seed: u64, inc: u64) -> u64 {
343 guarded(|| {
344 let stream = inc.saturating_sub(1) / 2;
345 let state = crate::numeric::pcg::build_cycle_walk_state(range, seed, stream);
346 crate::numeric::pcg::cycle_walk_inner(
347 pos,
348 range,
349 state.half_bits,
350 state.half_mask,
351 &state.round_keys,
352 )
353 })
354}
355
356extern "C" fn jit_perlin_1d(input: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
357 guarded(|| {
358 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
359 let freq = f64::from_bits(freq_bits);
360 let r = crate::numeric::noise::perlin_1d_algo(perm, input as f64 * freq);
361 r.to_bits()
362 })
363}
364
365extern "C" fn jit_perlin_2d(x: u64, y: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
366 guarded(|| {
367 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
368 let freq = f64::from_bits(freq_bits);
369 let r = crate::numeric::noise::perlin_2d_algo(perm, x as f64 * freq, y as f64 * freq);
370 r.to_bits()
371 })
372}
373
374extern "C" fn jit_simplex_2d(x: u64, y: u64, perm_ptr: u64, freq_bits: u64) -> u64 {
375 guarded(|| {
376 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
377 let freq = f64::from_bits(freq_bits);
378 let r = crate::numeric::noise::simplex_2d_algo(perm, x as f64 * freq, y as f64 * freq);
379 r.to_bits()
380 })
381}
382
383extern "C" fn jit_fractal_noise_1d(input: u64, perm_ptr: u64, freq_bits: u64, octaves: u64) -> u64 {
384 guarded(|| {
385 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
386 let freq = f64::from_bits(freq_bits);
387 let r = crate::numeric::noise::fbm_1d(perm, input as f64, freq, octaves as u32);
388 r.to_bits()
389 })
390}
391
392extern "C" fn jit_fractal_noise_2d(
393 x: u64,
394 y: u64,
395 perm_ptr: u64,
396 freq_bits: u64,
397 octaves: u64,
398) -> u64 {
399 guarded(|| {
400 let perm = unsafe { &*(perm_ptr as *const crate::numeric::noise::PermTable) };
401 let freq = f64::from_bits(freq_bits);
402 let r = crate::numeric::noise::fbm_2d(perm, x as f64, y as f64, freq, octaves as u32);
403 r.to_bits()
404 })
405}
406
407extern "C" fn jit_thread_id() -> u64 {
408 guarded(|| {
409 let id = std::thread::current().id();
410 let id_str = format!("{id:?}");
411 let num = id_str.trim_start_matches("ThreadId(").trim_end_matches(')');
412 num.parse().unwrap_or(0)
413 })
414}
415
416extern "C" fn jit_current_epoch_millis() -> u64 {
417 guarded(|| {
418 std::time::SystemTime::now()
419 .duration_since(std::time::UNIX_EPOCH)
420 .unwrap()
421 .as_millis() as u64
422 })
423}
424
425#[repr(C, align(16))]
463struct JitJmpBuf([u8; 512]);
464
465#[cfg(not(windows))]
466unsafe extern "C" {
467 fn _setjmp(env: *mut JitJmpBuf) -> i32;
468 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
469}
470
471#[cfg(windows)]
484unsafe extern "C" {
485 fn _setjmp(env: *mut JitJmpBuf, frame: *mut std::ffi::c_void) -> i32;
486 #[link_name = "longjmp"]
487 fn _longjmp(env: *mut JitJmpBuf, val: i32) -> !;
488}
489
490use std::cell::{Cell, RefCell};
491thread_local! {
492 static JIT_JMP_BUF: Cell<Option<*mut JitJmpBuf>> = const { Cell::new(None) };
499 static JIT_VIOLATION_MSG: RefCell<Option<String>> = const { RefCell::new(None) };
503}
504
505fn jit_violation_longjmp(msg: String) -> ! {
512 JIT_VIOLATION_MSG.with(|m| *m.borrow_mut() = Some(msg.clone()));
513 let buf_ptr: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.get());
514 match buf_ptr {
515 Some(ptr) => unsafe { _longjmp(ptr, 1) },
516 None => {
517 let mut err = std::io::stderr().lock();
518 use std::io::Write;
519 let _ = writeln!(err, "{msg}");
520 let _ = err.flush();
521 std::process::abort();
522 }
523 }
524}
525
526struct JmpBufGuard {
534 prev: Option<*mut JitJmpBuf>,
535}
536
537impl Drop for JmpBufGuard {
538 fn drop(&mut self) {
539 JIT_JMP_BUF.with(|b| b.set(self.prev));
540 }
541}
542
543pub(crate) fn invoke_with_catch<F: FnOnce()>(f: F) {
553 use std::mem::MaybeUninit;
554 let mut buf: MaybeUninit<JitJmpBuf> = MaybeUninit::uninit();
555 let buf_ptr = buf.as_mut_ptr();
556 let prev: Option<*mut JitJmpBuf> = JIT_JMP_BUF.with(|b| b.replace(Some(buf_ptr)));
561 let _guard = JmpBufGuard { prev };
562 #[cfg(not(windows))]
563 let jmpval = unsafe { _setjmp(buf_ptr) };
564 #[cfg(windows)]
566 let jmpval = unsafe { _setjmp(buf_ptr, std::ptr::null_mut()) };
567 if jmpval == 0 {
568 f();
569 } else {
570 let msg = JIT_VIOLATION_MSG
575 .with(|m| m.borrow_mut().take())
576 .unwrap_or_else(|| "JIT predicate violation (no message)".into());
577 std::panic::resume_unwind(Box::new(msg));
581 }
582}
583
584extern "C" fn jit_is_positive_fail(value: u64, name_ptr: u64, name_len: u64) -> u64 {
588 let name = if name_ptr != 0 {
593 unsafe {
594 std::str::from_utf8_unchecked(std::slice::from_raw_parts(
595 name_ptr as *const u8,
596 name_len as usize,
597 ))
598 }
599 } else {
600 "value"
601 };
602 jit_violation_longjmp(format!(
603 "is_positive({name}): value must be > 0, got {value}"
604 ));
605}
606
607extern "C" fn jit_in_range_fail(value: u64, lo: u64, hi: u64) -> u64 {
610 jit_violation_longjmp(format!("in_range: value {value} outside [{lo}, {hi}]"));
611}
612
613extern "C" fn jit_div_zero_fail(kind: u64) -> u64 {
617 jit_violation_longjmp(
618 if kind == 0 {
619 "attempt to divide by zero"
620 } else {
621 "attempt to calculate the remainder with a divisor of zero"
622 }
623 .to_string(),
624 );
625}
626
627extern "C" fn jit_f64_mod(a_bits: u64, b_bits: u64) -> u64 {
631 let (a, b) = (f64::from_bits(a_bits), f64::from_bits(b_bits));
632 (if b != 0.0 { a % b } else { 0.0 }).to_bits()
633}
634
635extern "C" fn jit_is_one_of_fail(value: u64, set_ptr: u64, set_len: u64) -> u64 {
643 let msg = if set_ptr != 0 {
644 let set = unsafe { std::slice::from_raw_parts(set_ptr as *const u64, set_len as usize) };
645 format!("is_one_of: value {value} not in allowed set {set:?}")
646 } else {
647 format!("is_one_of: value {value} not in allowed set [..]")
648 };
649 jit_violation_longjmp(msg);
650}
651
652extern "C" fn jit_weighted_pick(
657 input: u64,
658 values_ptr: u64,
659 biases_ptr: u64,
660 primaries_ptr: u64,
661 aliases_ptr: u64,
662 n: u64,
663) -> u64 {
664 guarded(|| {
665 let n = n as usize;
666 let slot = (input as usize) % n;
667 let bias_test = ((input >> 32) as f64) / (u32::MAX as f64);
668 unsafe {
669 let biases = std::slice::from_raw_parts(biases_ptr as *const f64, n);
670 let primaries = std::slice::from_raw_parts(primaries_ptr as *const u64, n);
671 let aliases = std::slice::from_raw_parts(aliases_ptr as *const u64, n);
672 let values = std::slice::from_raw_parts(values_ptr as *const u64, n);
673 let index = if bias_test < biases[slot] {
674 primaries[slot]
675 } else {
676 aliases[slot]
677 };
678 values[index as usize]
679 }
680 })
681}
682
683fn guarded<T>(body: impl FnOnce() -> T) -> T {
688 match std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)) {
689 Ok(v) => v,
690 Err(payload) => {
691 let msg = payload
692 .downcast_ref::<String>()
693 .cloned()
694 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
695 .unwrap_or_else(|| "panic in a compiled helper".to_string());
696 jit_violation_longjmp(msg)
697 }
698 }
699}
700
701#[derive(Clone)]
704pub struct SlotKitRef(pub std::sync::Arc<crate::ast::CompiledSlotKit>);
705
706impl SlotKitRef {
707 fn new(kit: crate::ast::CompiledSlotKit) -> Self {
708 SlotKitRef(std::sync::Arc::new(kit))
709 }
710}
711
712impl std::fmt::Debug for SlotKitRef {
713 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
714 write!(
715 f,
716 "SlotKitRef({:p}, {} scratch)",
717 std::sync::Arc::as_ptr(&self.0),
718 self.0.scratch.len()
719 )
720 }
721}
722
723impl PartialEq for SlotKitRef {
724 fn eq(&self, other: &Self) -> bool {
725 std::sync::Arc::ptr_eq(&self.0, &other.0)
726 }
727}
728
729extern "C" fn jit_slot_call(
741 kit: *const crate::ast::CompiledSlotKit,
742 inputs: *const u64,
743 n_in: u64,
744 outputs: *mut u64,
745 n_out: u64,
746 scratch: *mut crate::ast::ScratchBuf,
747 base: u64,
748 n_scratch: u64,
749) {
750 guarded(|| unsafe {
751 let kit = &*kit;
752 let ins = std::slice::from_raw_parts(inputs, n_in as usize);
753 let outs = std::slice::from_raw_parts_mut(outputs, n_out as usize);
754 let sc = std::slice::from_raw_parts_mut(scratch.add(base as usize), n_scratch as usize);
755 (kit.op)(ins, outs, sc)
756 })
757}
758
759impl JitOp {
760 pub(crate) fn slot_kit(&self) -> Option<&SlotKitRef> {
762 match self {
763 JitOp::SlotCall { kit, .. } => Some(kit),
764 _ => None,
765 }
766 }
767
768 pub(crate) fn scratch_elems(&self) -> &[crate::ast::ScratchElem] {
770 const STR_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::Str];
771 const F32_ENTRY: [crate::ast::ScratchElem; 1] = [crate::ast::ScratchElem::F32];
772 match self {
773 JitOp::SlotCall { kit, .. } => &kit.0.scratch,
774 JitOp::U64ToStr { .. }
775 | JitOp::I64ToStr { .. }
776 | JitOp::F64ToStr { .. }
777 | JitOp::StrConcat { .. }
778 | JitOp::JsonToStr { .. } => &STR_ENTRY,
779 JitOp::VecProduce { .. } => &F32_ENTRY,
780 _ => &[],
781 }
782 }
783
784 pub(crate) fn place_scratch(&mut self, base: usize) {
787 match self {
788 JitOp::SlotCall { scratch_base, .. }
789 | JitOp::U64ToStr { scratch_base }
790 | JitOp::I64ToStr { scratch_base }
791 | JitOp::F64ToStr { scratch_base }
792 | JitOp::StrConcat { scratch_base }
793 | JitOp::JsonToStr { scratch_base }
794 | JitOp::VecProduce { scratch_base, .. } => *scratch_base = base,
795 _ => {}
796 }
797 }
798}
799
800#[derive(Debug, Clone, Copy, PartialEq, Eq)]
804pub enum VecProducer {
805 Add,
807 Scale,
809 Norm,
811 HashVec,
813 XxHash3Vec,
815 RegToVec,
817}
818
819#[derive(Debug, Clone, Copy, PartialEq, Eq)]
822pub enum VecReducer {
823 Dot,
825 L2,
827 Cosine,
829 LidMle,
831}
832
833#[derive(Debug, Clone, Copy, PartialEq, Eq)]
836pub enum RegLaneRead {
837 F32,
839 I16,
841 I64,
843}
844
845#[derive(Debug, Clone, Copy, PartialEq, Eq)]
849pub enum RegProducer {
850 WithLaneF32,
852 GatherF32,
854 VecToRegF32,
856 MulI8,
858}
859
860unsafe fn vec_f32_of<'a>(ptr: u64, len: u64) -> &'a [f32] {
866 if len == 0 {
867 &[]
868 } else {
869 unsafe { std::slice::from_raw_parts(ptr as usize as *const f32, len as usize) }
870 }
871}
872
873unsafe fn write_f32_entry(
880 scratch: *mut crate::ast::ScratchBuf,
881 base: u64,
882 buffer: *mut u64,
883 out_slot: u64,
884 f: impl FnOnce(&mut Vec<f32>),
885) {
886 unsafe {
887 let entry = &mut *scratch.add(base as usize);
888 let crate::ast::ScratchBuf::F32(v) = entry else {
889 panic!("a vector lowering's scratch entry is not an f32 vector");
890 };
891 f(v);
892 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
893 *buffer.add(out_slot as usize + 1) = v.len() as u64;
894 }
895}
896
897macro_rules! vec_producer {
901 ($name:ident, |$out:ident, $w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
902 extern "C" fn $name(
903 scratch: *mut crate::ast::ScratchBuf,
904 base: u64,
905 buffer: *mut u64,
906 out_slot: u64,
907 $w0: u64,
908 $w1: u64,
909 $w2: u64,
910 $w3: u64,
911 ) {
912 guarded(|| unsafe {
913 let _ = ($w2, $w3);
914 write_f32_entry(scratch, base, buffer, out_slot, |$out| $body)
915 })
916 }
917 };
918}
919
920vec_producer!(jit_vec_add, |out, a_ptr, a_len, b_ptr, b_len| {
921 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
922 crate::numeric::vector::check_lens("vec_add", a.len(), b.len());
923 crate::numeric::vector::add_f32_into(a, b, out)
924});
925vec_producer!(jit_vec_scale, |out, a_ptr, a_len, k_bits, _z| {
926 let a = vec_f32_of(a_ptr, a_len);
927 crate::numeric::vector::scale_f32_into(a, f64::from_bits(k_bits) as f32, out)
928});
929vec_producer!(jit_vec_norm, |out, a_ptr, a_len, _y, _z| {
930 crate::numeric::vector::norm_f32_into(vec_f32_of(a_ptr, a_len), out)
931});
932vec_producer!(jit_hash_vec, |out, seed, dim, _y, _z| {
933 crate::numeric::vector::hash_vec_into(seed, dim, out)
934});
935vec_producer!(jit_xxhash3_vec, |out, seed, dim, _y, _z| {
936 crate::numeric::vector::xxhash3_vec_into(seed, dim, out)
937});
938vec_producer!(jit_reg_to_vec_f32, |out, lo, hi, _y, _z| {
939 out.clear();
940 out.extend_from_slice(&crate::ast::Bits128([lo, hi]).lanes_f32())
941});
942
943macro_rules! vec_reducer {
946 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
947 extern "C" fn $name($w0: u64, $w1: u64, $w2: u64, $w3: u64) -> u64 {
948 guarded(|| unsafe {
949 let _ = ($w2, $w3);
950 let r: f64 = $body;
951 r.to_bits()
952 })
953 }
954 };
955}
956
957vec_reducer!(jit_vec_dot, |a_ptr, a_len, b_ptr, b_len| {
958 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
959 crate::numeric::vector::check_lens("vec_dot", a.len(), b.len());
960 crate::numeric::vector::dot_f32(a, b) as f64
961});
962vec_reducer!(jit_vec_l2, |a_ptr, a_len, b_ptr, b_len| {
963 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
964 crate::numeric::vector::check_lens("vec_l2", a.len(), b.len());
965 (crate::numeric::vector::l2sq_f32(a, b) as f64).sqrt()
966});
967vec_reducer!(jit_vec_cosine, |a_ptr, a_len, b_ptr, b_len| {
968 let (a, b) = (vec_f32_of(a_ptr, a_len), vec_f32_of(b_ptr, b_len));
969 crate::numeric::vector::check_lens("vec_cosine", a.len(), b.len());
970 crate::numeric::vector::cosine_f32(a, b)
971});
972vec_reducer!(jit_lid_mle, |d_ptr, d_len, k_bits, _z| {
973 crate::numeric::vector::lid_mle_of(vec_f32_of(d_ptr, d_len), f64::from_bits(k_bits))
974});
975
976extern "C" fn jit_reg_lane_f32(lo: u64, hi: u64, i: u64) -> u64 {
978 guarded(|| crate::numeric::register::lane_f32(crate::ast::Bits128([lo, hi]), i).to_bits())
979}
980
981extern "C" fn jit_reg_lane_i16(lo: u64, hi: u64, i: u64) -> u64 {
983 guarded(|| crate::numeric::register::lane_i16(crate::ast::Bits128([lo, hi]), i) as i64 as u64)
984}
985
986extern "C" fn jit_reg_lane_i64(lo: u64, hi: u64, i: u64) -> u64 {
988 guarded(|| crate::numeric::register::lane_i64(crate::ast::Bits128([lo, hi]), i) as u64)
989}
990
991macro_rules! reg_producer {
994 ($name:ident, |$w0:ident, $w1:ident, $w2:ident, $w3:ident| $body:expr) => {
995 extern "C" fn $name(
996 buffer: *mut u64,
997 out_slot: u64,
998 $w0: u64,
999 $w1: u64,
1000 $w2: u64,
1001 $w3: u64,
1002 ) {
1003 guarded(|| unsafe {
1004 let _ = ($w2, $w3);
1005 let r: crate::ast::Bits128 = $body;
1006 *buffer.add(out_slot as usize) = r.0[0];
1007 *buffer.add(out_slot as usize + 1) = r.0[1];
1008 })
1009 }
1010 };
1011}
1012
1013reg_producer!(jit_reg_with_lane_f32, |lo, hi, i, v_bits| {
1014 crate::numeric::register::with_lane_f32(
1015 crate::ast::Bits128([lo, hi]),
1016 i,
1017 f64::from_bits(v_bits),
1018 )
1019});
1020reg_producer!(jit_reg_gather_f32, |v_ptr, v_len, offset, _z| {
1021 crate::numeric::register::gather_f32(vec_f32_of(v_ptr, v_len), offset)
1022});
1023reg_producer!(jit_vec_to_reg_f32, |v_ptr, v_len, _y, _z| {
1024 crate::numeric::register::to_reg_f32(vec_f32_of(v_ptr, v_len))
1025});
1026reg_producer!(jit_reg_mul_i8, |a_lo, a_hi, b_lo, b_hi| {
1027 crate::numeric::register::mul_i8(
1028 crate::ast::Bits128([a_lo, a_hi]),
1029 crate::ast::Bits128([b_lo, b_hi]),
1030 )
1031});
1032
1033unsafe fn write_str_entry(
1042 scratch: *mut crate::ast::ScratchBuf,
1043 base: u64,
1044 buffer: *mut u64,
1045 out_slot: u64,
1046 f: impl FnOnce(&mut Vec<u8>),
1047) {
1048 unsafe {
1049 let entry = &mut *scratch.add(base as usize);
1050 let crate::ast::ScratchBuf::Str(v) = entry else {
1051 panic!("a string lowering's scratch entry is not a string");
1052 };
1053 v.clear();
1054 f(v);
1055 *buffer.add(out_slot as usize) = v.as_ptr() as usize as u64;
1056 *buffer.add(out_slot as usize + 1) = v.len() as u64;
1057 }
1058}
1059
1060extern "C" fn jit_u64_to_str(
1062 scratch: *mut crate::ast::ScratchBuf,
1063 base: u64,
1064 buffer: *mut u64,
1065 out_slot: u64,
1066 value: u64,
1067) {
1068 use std::io::Write;
1069 guarded(|| unsafe {
1070 write_str_entry(scratch, base, buffer, out_slot, |v| {
1071 write!(v, "{value}").expect("a vector accepts every write")
1072 })
1073 })
1074}
1075
1076extern "C" fn jit_i64_to_str(
1078 scratch: *mut crate::ast::ScratchBuf,
1079 base: u64,
1080 buffer: *mut u64,
1081 out_slot: u64,
1082 value: u64,
1083) {
1084 use std::io::Write;
1085 guarded(|| unsafe {
1086 write_str_entry(scratch, base, buffer, out_slot, |v| {
1087 write!(v, "{}", value as i64).expect("a vector accepts every write")
1088 })
1089 })
1090}
1091
1092extern "C" fn jit_f64_to_str(
1095 scratch: *mut crate::ast::ScratchBuf,
1096 base: u64,
1097 buffer: *mut u64,
1098 out_slot: u64,
1099 bits: u64,
1100) {
1101 use std::io::Write;
1102 guarded(|| unsafe {
1103 write_str_entry(scratch, base, buffer, out_slot, |v| {
1104 write!(v, "{}", f64::from_bits(bits)).expect("a vector accepts every write")
1105 })
1106 })
1107}
1108
1109extern "C" fn jit_str_concat(
1113 scratch: *mut crate::ast::ScratchBuf,
1114 base: u64,
1115 buffer: *mut u64,
1116 out_slot: u64,
1117 pairs: *const u64,
1118 n: u64,
1119) {
1120 guarded(|| unsafe {
1121 let words = std::slice::from_raw_parts(pairs, 2 * n as usize);
1122 write_str_entry(scratch, base, buffer, out_slot, |v| {
1123 for pair in words.as_chunks::<2>().0 {
1124 let bytes =
1128 std::slice::from_raw_parts(pair[0] as usize as *const u8, pair[1] as usize);
1129 v.extend_from_slice(bytes);
1130 }
1131 })
1132 })
1133}
1134
1135extern "C" fn jit_json_to_str(
1138 scratch: *mut crate::ast::ScratchBuf,
1139 base: u64,
1140 buffer: *mut u64,
1141 out_slot: u64,
1142 ptr: u64,
1143 len: u64,
1144) {
1145 guarded(|| unsafe {
1146 let pair = [ptr, len];
1147 let value = crate::derive_support::ref_value(&pair);
1148 let json = match value {
1149 crate::ast::Value::Json(j) => j.as_ref(),
1150 other => panic!("expected Json wire, got {other:?}"),
1151 };
1152 write_str_entry(scratch, base, buffer, out_slot, |v| {
1153 serde_json::to_writer(v, json).expect("a vector accepts every write")
1154 })
1155 })
1156}
1157
1158pub fn classify_node_typed(node: &dyn PolydatNode, wire_types: &[crate::ast::PortType]) -> JitOp {
1167 use crate::ast::PortType as PT;
1168 let is_ref = |t: &crate::ast::PortType| t.slot_color() == crate::ast::SlotColor::Ref2;
1169 let vec_produce = |kind: VecProducer| JitOp::VecProduce {
1170 kind,
1171 scratch_base: 0,
1172 };
1173 let ref_copy = |ty: crate::ast::PortType| {
1174 crate::compile::assembly::ref_copy_kit(ty)
1175 .map(|kit| JitOp::SlotCall {
1176 kit: SlotKitRef::new(kit),
1177 scratch_base: 0,
1178 })
1179 .unwrap_or(JitOp::Fallback)
1180 };
1181 let meta = node.meta();
1182 let named = match meta.name.as_str() {
1183 "__u64_to_string" => JitOp::U64ToStr { scratch_base: 0 },
1186 "__i64_to_string" => JitOp::I64ToStr { scratch_base: 0 },
1187 "__f64_to_string" => JitOp::F64ToStr { scratch_base: 0 },
1188 "json_to_str" if wire_types == [crate::ast::PortType::Json] => {
1189 JitOp::JsonToStr { scratch_base: 0 }
1190 }
1191 "str_concat"
1192 if !wire_types.is_empty()
1193 && wire_types.iter().all(|t| *t == crate::ast::PortType::Str) =>
1194 {
1195 JitOp::StrConcat { scratch_base: 0 }
1196 }
1197 "vec_add" if wire_types == [PT::VecF32, PT::VecF32] => vec_produce(VecProducer::Add),
1202 "vec_scale" if wire_types == [PT::VecF32, PT::F64] => vec_produce(VecProducer::Scale),
1203 "vec_norm" if wire_types == [PT::VecF32] => vec_produce(VecProducer::Norm),
1204 "hash_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::HashVec),
1205 "xxhash3_vec" if wire_types == [PT::U64, PT::U64] => vec_produce(VecProducer::XxHash3Vec),
1206 "reg_to_vec_f32" if wire_types == [PT::RegF32x4] => vec_produce(VecProducer::RegToVec),
1207 "vec_dot" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::Dot),
1208 "vec_l2" if wire_types == [PT::VecF32, PT::VecF32] => JitOp::VecReduce(VecReducer::L2),
1209 "vec_cosine" if wire_types == [PT::VecF32, PT::VecF32] => {
1210 JitOp::VecReduce(VecReducer::Cosine)
1211 }
1212 "lid_mle" if wire_types == [PT::VecF32, PT::F64] => JitOp::VecReduce(VecReducer::LidMle),
1213 "reg_lane_f32" if wire_types == [PT::RegF32x4, PT::U64] => JitOp::RegLane(RegLaneRead::F32),
1217 "reg_lane_i16" if wire_types == [PT::RegI16x8, PT::U64] => JitOp::RegLane(RegLaneRead::I16),
1218 "reg_lane_i64" if wire_types == [PT::RegI64x2, PT::U64] => JitOp::RegLane(RegLaneRead::I64),
1219 "reg_with_lane_f32" if wire_types == [PT::RegF32x4, PT::U64, PT::F64] => {
1220 JitOp::RegProduce(RegProducer::WithLaneF32)
1221 }
1222 "reg_gather_f32" if wire_types == [PT::VecF32, PT::U64] => {
1223 JitOp::RegProduce(RegProducer::GatherF32)
1224 }
1225 "vec_to_reg_f32" if wire_types == [PT::VecF32] => {
1226 JitOp::RegProduce(RegProducer::VecToRegF32)
1227 }
1228 "reg_mul_i8" if wire_types == [PT::RegI8x16, PT::RegI8x16] => {
1229 JitOp::RegProduce(RegProducer::MulI8)
1230 }
1231 "reg_dot_f32" if wire_types == [PT::RegF32x4, PT::RegF32x4] => JitOp::RegDotF32,
1232 n if n.starts_with("__port_") || n == "default_or" => match meta.outs.first() {
1238 Some(o) if is_ref(&o.typ) => return ref_copy(o.typ),
1239 _ => JitOp::Identity,
1240 },
1241 "select" | "select_u64" if wire_types.iter().skip(1).any(|t| t.slot_width() != 1) => {
1244 JitOp::Fallback
1245 }
1246 _ if wire_types.iter().any(is_ref) || meta.outs.iter().any(|o| is_ref(&o.typ)) => {
1247 JitOp::Fallback
1248 }
1249 _ => classify_node(node),
1250 };
1251 if !matches!(named, JitOp::Fallback) {
1252 return named;
1253 }
1254 if let Some(kit) = node.compiled_slot(wire_types) {
1255 return JitOp::SlotCall {
1256 kit: SlotKitRef::new(kit),
1257 scratch_base: 0,
1258 };
1259 }
1260 if let Some(op) = node.compiled_u64() {
1261 return JitOp::SlotCall {
1262 kit: SlotKitRef::new(crate::ast::CompiledSlotKit {
1263 scratch: Vec::new(),
1264 op: Box::new(move |inputs, outputs, _| op(inputs, outputs)),
1265 }),
1266 scratch_base: 0,
1267 };
1268 }
1269 JitOp::Fallback
1270}
1271
1272#[derive(Debug, Clone, PartialEq)]
1279pub enum JitOp {
1280 Identity,
1283 AddConst(u64),
1285 MulConst(u64),
1287 DivConst(u64),
1289 ModConst(u64),
1291 ClampConst(u64, u64),
1293 Interleave,
1295 MixedRadixConst(Vec<u64>),
1297 Hash,
1299 SplitMix64,
1301 Popcnt,
1303 Clz,
1305 Ctz,
1307 Bswap,
1309 ShuffleConst(u64, u64, u64),
1311
1312 UnitInterval,
1315 F64ToU64,
1317 RoundToU64,
1320 FloorToU64,
1322 CeilToU64,
1324 ClampF64Const(u64, u64), LerpConst(u64, u64), ScaleRangeConst(u64, u64), QuantizeConst(u64), DiscretizeConst(u64, u64), LutSampleConst(u64, u64), WeightedPickConst(u64, u64, u64, u64, u64), MathUnary(u8),
1342 MathBinary(u8),
1345
1346 U64Add2,
1349 U64Sub2,
1351 U64Mul2,
1353 U64Div2,
1355 U64Mod2,
1357 U64And,
1359 U64Or,
1361 U64Xor,
1363 U64Shl,
1365 U64Shr,
1367 U64Not,
1369
1370 ToF64,
1373
1374 F64Add,
1376 F64Sub,
1378 F64Mul,
1380 F64Div,
1382 F64Mod,
1384 U64DivWire,
1387 U64ModWire,
1390
1391 SlotCall {
1399 kit: SlotKitRef,
1402 scratch_base: usize,
1405 },
1406
1407 U64ToStr {
1413 scratch_base: usize,
1415 },
1416 I64ToStr {
1418 scratch_base: usize,
1420 },
1421 F64ToStr {
1423 scratch_base: usize,
1425 },
1426 StrConcat {
1429 scratch_base: usize,
1431 },
1432 JsonToStr {
1434 scratch_base: usize,
1436 },
1437
1438 VecProduce {
1442 kind: VecProducer,
1444 scratch_base: usize,
1446 },
1447 VecReduce(VecReducer),
1449 RegLane(RegLaneRead),
1452 RegProduce(RegProducer),
1454 RegDotF32,
1458 RegShuffleConst([u8; 16]),
1461
1462 IsPositiveCheck {
1469 name_ptr: u64,
1471 name_len: u64,
1473 },
1474 InRangeCheck(u64, u64),
1478 IsOneOfCheck {
1486 allowed: Vec<u64>,
1488 set_ptr: u64,
1490 set_len: u64,
1492 },
1493
1494 RegBinOp(u8, u8),
1502 RegCopy,
1506 RegSplat(u8),
1510
1511 U64Cmp(ir::condcodes::IntCC),
1514 F64Cmp(ir::condcodes::FloatCC),
1516 SelectU64,
1518 SelectF64,
1520
1521 I64ToF64,
1524 F64ToI64,
1526 SignExtendI32,
1528 SignExtendI16,
1530 SignExtendI8,
1532 ZeroExtendU32,
1534 ZeroExtendU16,
1536 ZeroExtendU8,
1538 ToBool,
1540 ConstU64(u64),
1542 ConstF64(u64),
1544
1545 HashRangeConst(u64),
1548 HashIntervalConst(u64, u64),
1550 InvLerpConst(u64, u64),
1552 RemapConst(u64, u64, u64, u64),
1554
1555 EpochOffsetConst(u64),
1558 EpochScaleConst(u64),
1560 ThreadId,
1562 CurrentEpochMillis,
1564
1565 Perlin1dConst(u64, u64),
1568 Perlin2dConst(u64, u64),
1570 Simplex2dConst(u64, u64),
1572 FractalNoise1dConst(u64, u64, u64),
1574 FractalNoise2dConst(u64, u64, u64),
1576
1577 VariadicSum,
1580 VariadicProduct,
1582 VariadicMin,
1584 VariadicMax,
1586 CheckedAdd,
1588 CheckedSub,
1590 CheckedMul,
1592 CeilToMultiple,
1594 MultiplesAtLeast,
1596
1597 FairCoin,
1600 BlendConst(u64),
1602 LfsrStepConst(u64),
1605 PcgConst(u64, u64),
1607 PcgStreamConst(u64),
1609 CycleWalkConst(u64, u64, u64),
1611 UnfairCoinConst(u64),
1613 CoinFlipConst(u64),
1616 ChanceConst(u64),
1618 NOfConst(u64, u64),
1620
1621 Fallback,
1623}
1624
1625pub fn classify_node(node: &dyn PolydatNode) -> JitOp {
1632 let name = node.meta().name.as_str();
1633 let consts = node.jit_constants();
1634
1635 match name {
1636 "identity" => JitOp::Identity,
1637 "hash" | "splitmix64" | "scatter" => JitOp::SplitMix64,
1638 "fair_coin" => JitOp::FairCoin,
1639 "unfair_coin" | "bernoulli" => {
1640 if let Some(&p) = consts.first() {
1641 JitOp::UnfairCoinConst(p)
1642 } else {
1643 JitOp::Fallback
1644 }
1645 }
1646 "chance" => {
1647 if let Some(&p) = consts.first() {
1648 JitOp::ChanceConst(p)
1649 } else {
1650 JitOp::Fallback
1651 }
1652 }
1653 "popcnt" | "count_ones" | "popcount" => JitOp::Popcnt,
1654 "clz" | "leading_zeros" => JitOp::Clz,
1655 "ctz" | "trailing_zeros" => JitOp::Ctz,
1656 "bswap" | "swap_bytes" => JitOp::Bswap,
1657 "xxhash3" | "xxh3" => JitOp::Hash,
1658 "hash_range" => {
1659 if let Some(&c) = consts.first() {
1660 JitOp::HashRangeConst(c)
1661 } else {
1662 JitOp::Fallback
1663 }
1664 }
1665 "hash_interval" => {
1666 if consts.len() >= 2 {
1667 JitOp::HashIntervalConst(consts[0], consts[1])
1668 } else {
1669 JitOp::Fallback
1670 }
1671 }
1672 "add" => {
1673 if let Some(&c) = consts.first() {
1674 JitOp::AddConst(c)
1675 } else {
1676 JitOp::Fallback
1677 }
1678 }
1679 "mul" => {
1680 if let Some(&c) = consts.first() {
1681 JitOp::MulConst(c)
1682 } else {
1683 JitOp::Fallback
1684 }
1685 }
1686 "div" => {
1687 if let Some(&c) = consts.first() {
1688 JitOp::DivConst(c)
1689 } else {
1690 JitOp::Fallback
1691 }
1692 }
1693 "mod" => {
1694 if let Some(&c) = consts.first() {
1695 JitOp::ModConst(c)
1696 } else {
1697 JitOp::Fallback
1698 }
1699 }
1700 "clamp" => {
1701 if consts.len() >= 2 {
1702 JitOp::ClampConst(consts[0], consts[1])
1703 } else {
1704 JitOp::Fallback
1705 }
1706 }
1707 "interleave" => JitOp::Interleave,
1708 "mixed_radix" => {
1709 if consts.is_empty() {
1710 JitOp::Fallback
1711 } else {
1712 JitOp::MixedRadixConst(consts)
1713 }
1714 }
1715 "shuffle" => {
1716 if consts.len() >= 3 {
1717 JitOp::ShuffleConst(consts[0], consts[1], consts[2])
1718 } else {
1719 JitOp::Fallback
1720 }
1721 }
1722 "unit_interval" => JitOp::UnitInterval,
1724 "f64_to_u64" => JitOp::F64ToU64,
1725 "round_to_u64" => JitOp::RoundToU64,
1726 "floor_to_u64" => JitOp::FloorToU64,
1727 "ceil_to_u64" => JitOp::CeilToU64,
1728 "clamp_f64" => {
1729 if consts.len() >= 2 {
1730 JitOp::ClampF64Const(consts[0], consts[1])
1731 } else {
1732 JitOp::Fallback
1733 }
1734 }
1735 "lerp" => {
1736 if consts.len() >= 2 {
1737 JitOp::LerpConst(consts[0], consts[1])
1738 } else {
1739 JitOp::Fallback
1740 }
1741 }
1742 "scale_range" => {
1743 if consts.len() >= 2 {
1744 JitOp::ScaleRangeConst(consts[0], consts[1])
1745 } else {
1746 JitOp::Fallback
1747 }
1748 }
1749 "quantize" => {
1750 if let Some(&c) = consts.first() {
1751 JitOp::QuantizeConst(c)
1752 } else {
1753 JitOp::Fallback
1754 }
1755 }
1756 "discretize" => {
1757 if consts.len() >= 2 {
1758 JitOp::DiscretizeConst(consts[0], consts[1])
1759 } else {
1760 JitOp::Fallback
1761 }
1762 }
1763 "lut_sample" | "dist_normal" | "icd_normal" | "dist_exponential" | "icd_exponential"
1764 | "dist_uniform" | "dist_pareto" | "dist_zipf" | "dist_empirical" => {
1765 if consts.len() >= 2 {
1766 JitOp::LutSampleConst(consts[0], consts[1])
1767 } else {
1768 JitOp::Fallback
1769 }
1770 }
1771 "sin" => JitOp::MathUnary(0),
1773 "cos" => JitOp::MathUnary(1),
1774 "tan" => JitOp::MathUnary(2),
1775 "asin" => JitOp::MathUnary(3),
1776 "acos" => JitOp::MathUnary(4),
1777 "atan" => JitOp::MathUnary(5),
1778 "sqrt" => JitOp::MathUnary(6),
1779 "abs_f64" => JitOp::MathUnary(7),
1780 "ln" => JitOp::MathUnary(8),
1781 "exp" => JitOp::MathUnary(9),
1782 "floor_base10" => JitOp::MathUnary(10),
1783 "ceiling_base10" => JitOp::MathUnary(11),
1784 "closest_base10" => JitOp::MathUnary(12),
1785 "floor_decade" => JitOp::MathUnary(13),
1786 "ceiling_decade" => JitOp::MathUnary(14),
1787 "closest_decade" => JitOp::MathUnary(15),
1788 "floor_binomial" => JitOp::MathUnary(16),
1789 "ceiling_binomial" => JitOp::MathUnary(17),
1790 "closest_binomial" => JitOp::MathUnary(18),
1791 "floor_fibonacci" => JitOp::MathUnary(19),
1792 "ceiling_fibonacci" => JitOp::MathUnary(20),
1793 "closest_fibonacci" => JitOp::MathUnary(21),
1794 "atan2" => JitOp::MathBinary(0),
1795 "pow" => JitOp::MathBinary(1),
1796 "round_nearest" => JitOp::MathBinary(2),
1797 "round_floor" => JitOp::MathBinary(3),
1798 "round_ceiling" => JitOp::MathBinary(4),
1799 "to_f64" => JitOp::ToF64,
1800 "u64_add" => JitOp::U64Add2,
1802 "u64_sub" => JitOp::U64Sub2,
1803 "u64_mul" => JitOp::U64Mul2,
1804 "u64_div" => JitOp::U64Div2,
1805 "u64_mod" => JitOp::U64Mod2,
1806 "u64_and" => JitOp::U64And,
1807 "u64_or" => JitOp::U64Or,
1808 "u64_xor" => JitOp::U64Xor,
1809 "u64_shl" => JitOp::U64Shl,
1810 "u64_shr" => JitOp::U64Shr,
1811 "u64_not" => JitOp::U64Not,
1812
1813 "reg_add_i8" => JitOp::RegBinOp(0, 0),
1815 "reg_sub_i8" => JitOp::RegBinOp(0, 1),
1816 "reg_shuffle_bytes" => {
1821 let mut mask = [0u8; 16];
1822 if consts.len() == 16 && consts.iter().all(|&m| m < 16) {
1823 for (m, &c) in mask.iter_mut().zip(consts.iter()) {
1824 *m = c as u8;
1825 }
1826 JitOp::RegShuffleConst(mask)
1827 } else {
1828 JitOp::Fallback
1829 }
1830 }
1831 "reg_add_i16" => JitOp::RegBinOp(1, 0),
1832 "reg_sub_i16" => JitOp::RegBinOp(1, 1),
1833 "reg_mul_i16" => JitOp::RegBinOp(1, 2),
1834 "reg_add_i32" => JitOp::RegBinOp(2, 0),
1835 "reg_sub_i32" => JitOp::RegBinOp(2, 1),
1836 "reg_mul_i32" => JitOp::RegBinOp(2, 2),
1837 "reg_add_i64" => JitOp::RegBinOp(3, 0),
1838 "reg_sub_i64" => JitOp::RegBinOp(3, 1),
1839 "reg_mul_i64" => JitOp::RegBinOp(3, 2),
1840 "reg_add_f32" => JitOp::RegBinOp(4, 0),
1841 "reg_sub_f32" => JitOp::RegBinOp(4, 1),
1842 "reg_mul_f32" => JitOp::RegBinOp(4, 2),
1843 "reg_add_f64" => JitOp::RegBinOp(5, 0),
1844 "reg_sub_f64" => JitOp::RegBinOp(5, 1),
1845 "reg_mul_f64" => JitOp::RegBinOp(5, 2),
1846 "__reg_view_raw" | "__reg_view_i8x16" | "__reg_view_i16x8" | "__reg_view_i32x4"
1847 | "__reg_view_i64x2" | "__reg_view_f16x8" | "__reg_view_f32x4" | "__reg_view_f64x2" => {
1848 JitOp::RegCopy
1849 }
1850 "reg_splat_i8" => JitOp::RegSplat(0),
1851 "reg_splat_i16" => JitOp::RegSplat(1),
1852 "reg_splat_i32" => JitOp::RegSplat(2),
1853 "reg_splat_i64" => JitOp::RegSplat(3),
1854 "reg_splat_f32" => JitOp::RegSplat(4),
1855 "reg_splat_f64" => JitOp::RegSplat(5),
1856
1857 "f64_add" => JitOp::F64Add,
1858 "f64_sub" => JitOp::F64Sub,
1859 "f64_mul" => JitOp::F64Mul,
1860 "f64_div" => JitOp::F64Div,
1861 "f64_mod" => JitOp::F64Mod,
1862
1863 "u64_eq" => JitOp::U64Cmp(ir::condcodes::IntCC::Equal),
1865 "u64_ne" => JitOp::U64Cmp(ir::condcodes::IntCC::NotEqual),
1866 "u64_lt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThan),
1867 "u64_le" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedLessThanOrEqual),
1868 "u64_gt" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThan),
1869 "u64_ge" => JitOp::U64Cmp(ir::condcodes::IntCC::UnsignedGreaterThanOrEqual),
1870 "f64_eq" => JitOp::F64Cmp(ir::condcodes::FloatCC::Equal),
1871 "f64_ne" => JitOp::F64Cmp(ir::condcodes::FloatCC::NotEqual),
1872 "f64_lt" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThan),
1873 "f64_le" => JitOp::F64Cmp(ir::condcodes::FloatCC::LessThanOrEqual),
1874 "f64_gt" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThan),
1875 "f64_ge" => JitOp::F64Cmp(ir::condcodes::FloatCC::GreaterThanOrEqual),
1876 "select_u64" | "select" => JitOp::SelectU64,
1877 "select_f64" => JitOp::SelectF64,
1878
1879 "div_wire" => JitOp::U64DivWire,
1881 "mod_wire" => JitOp::U64ModWire,
1882 "ceil_to_multiple" => JitOp::CeilToMultiple,
1883 "multiples_at_least" => JitOp::MultiplesAtLeast,
1884 "checked_add" => JitOp::CheckedAdd,
1885 "checked_sub" => JitOp::CheckedSub,
1886 "checked_mul" => JitOp::CheckedMul,
1887
1888 "sum" => JitOp::VariadicSum,
1890 "product" => JitOp::VariadicProduct,
1891 "min" => JitOp::VariadicMin,
1892 "max" => JitOp::VariadicMax,
1893
1894 "blend" => {
1896 if let Some(&c) = consts.first() {
1897 JitOp::BlendConst(c)
1898 } else {
1899 JitOp::Fallback
1900 }
1901 }
1902 "lfsr_step" => {
1903 if let Some(&fb) = consts.first() {
1904 JitOp::LfsrStepConst(fb)
1905 } else {
1906 JitOp::Fallback
1907 }
1908 }
1909 "pcg" => {
1910 if consts.len() >= 2 {
1911 JitOp::PcgConst(consts[0], consts[1])
1912 } else {
1913 JitOp::Fallback
1914 }
1915 }
1916 "pcg_stream" => {
1917 if let Some(&seed) = consts.first() {
1918 JitOp::PcgStreamConst(seed)
1919 } else {
1920 JitOp::Fallback
1921 }
1922 }
1923 "n_of" => {
1924 if consts.len() >= 2 {
1925 JitOp::NOfConst(consts[0], consts[1])
1926 } else {
1927 JitOp::Fallback
1928 }
1929 }
1930
1931 "cycle_walk" => {
1932 if consts.len() >= 3 {
1933 JitOp::CycleWalkConst(consts[0], consts[1], consts[2])
1934 } else {
1935 JitOp::Fallback
1936 }
1937 }
1938 "coin_flip" => {
1939 if let Some(&threshold) = consts.first() {
1944 JitOp::CoinFlipConst(threshold)
1945 } else {
1946 JitOp::Fallback
1947 }
1948 }
1949 "default_or" => JitOp::Identity,
1952 "const_u64" | "const_bool" | "session_start_millis" => {
1953 if let Some(&c) = consts.first() {
1954 JitOp::ConstU64(c)
1955 } else {
1956 JitOp::Fallback
1957 }
1958 }
1959 "const_f64" => {
1960 if let Some(&c) = consts.first() {
1961 JitOp::ConstF64(c)
1962 } else {
1963 JitOp::Fallback
1964 }
1965 }
1966 "inv_lerp" => {
1967 if consts.len() >= 2 {
1968 JitOp::InvLerpConst(consts[0], consts[1])
1969 } else {
1970 JitOp::Fallback
1971 }
1972 }
1973 "remap" => {
1974 if consts.len() >= 4 {
1975 JitOp::RemapConst(consts[0], consts[1], consts[2], consts[3])
1976 } else {
1977 JitOp::Fallback
1978 }
1979 }
1980 "epoch_offset" => {
1981 if let Some(&c) = consts.first() {
1982 JitOp::EpochOffsetConst(c)
1983 } else {
1984 JitOp::Fallback
1985 }
1986 }
1987 "epoch_scale" => {
1988 if let Some(&c) = consts.first() {
1989 JitOp::EpochScaleConst(c)
1990 } else {
1991 JitOp::Fallback
1992 }
1993 }
1994 "thread_id" => JitOp::ThreadId,
1995 "current_epoch_millis" => JitOp::CurrentEpochMillis,
1996 "perlin_1d" => {
1997 if consts.len() >= 2 {
1998 JitOp::Perlin1dConst(consts[0], consts[1])
1999 } else {
2000 JitOp::Fallback
2001 }
2002 }
2003 "perlin_2d" => {
2004 if consts.len() >= 2 {
2005 JitOp::Perlin2dConst(consts[0], consts[1])
2006 } else {
2007 JitOp::Fallback
2008 }
2009 }
2010 "simplex_2d" => {
2011 if consts.len() >= 2 {
2012 JitOp::Simplex2dConst(consts[0], consts[1])
2013 } else {
2014 JitOp::Fallback
2015 }
2016 }
2017 "fractal_noise_1d" => {
2018 if consts.len() >= 3 {
2019 JitOp::FractalNoise1dConst(consts[0], consts[1], consts[2])
2020 } else {
2021 JitOp::Fallback
2022 }
2023 }
2024 "fractal_noise_2d" => {
2025 if consts.len() >= 3 {
2026 JitOp::FractalNoise2dConst(consts[0], consts[1], consts[2])
2027 } else {
2028 JitOp::Fallback
2029 }
2030 }
2031
2032 "u64_to_f64" | "__u64_to_f64" | "u32_to_f64" | "__u32_to_f64" | "bool_to_f64"
2034 | "__bool_to_f64" | "bool_to_f32" | "__bool_to_f32" | "__f32_to_f64" | "f32_to_f64"
2035 | "__u64_to_f32" | "u64_to_f32" | "__u32_to_f32" | "u32_to_f32" | "__u16_to_f32"
2036 | "u16_to_f32" | "__u8_to_f32" | "u8_to_f32" | "__u16_to_f64" | "u16_to_f64"
2037 | "__u8_to_f64" | "u8_to_f64" | "__u128_to_f64" | "__u128_to_f32" | "__u128_to_f16" => {
2038 JitOp::ToF64
2039 }
2040
2041 "i64_to_f64" | "__i64_to_f64" | "i32_to_f64" | "__i32_to_f64" | "__i64_to_f32"
2042 | "i64_to_f32" | "__i32_to_f32" | "i32_to_f32" | "__i16_to_f32" | "i16_to_f32"
2043 | "__i8_to_f32" | "i8_to_f32" | "__i16_to_f64" | "i16_to_f64" | "__i8_to_f64"
2044 | "i8_to_f64" | "__i128_to_f64" | "__i128_to_f32" | "__i128_to_f16" => JitOp::I64ToF64,
2045
2046 "__f64_to_u64"
2047 | "__f64_to_u64_checked"
2048 | "f64_to_u32"
2049 | "__f64_to_u32"
2050 | "f32_to_u64"
2051 | "__f32_to_u64"
2052 | "f32_to_u32"
2053 | "__f32_to_u32"
2054 | "__f64_to_u16"
2055 | "f64_to_u16"
2056 | "__f64_to_u8"
2057 | "f64_to_u8"
2058 | "__f32_to_u16"
2059 | "f32_to_u16"
2060 | "__f32_to_u8"
2061 | "f32_to_u8"
2062 | "__f16_to_u64"
2063 | "__f16_to_u32"
2064 | "__f16_to_u16"
2065 | "__f16_to_u8"
2066 | "__f64_to_u128"
2067 | "__f32_to_u128"
2068 | "__f16_to_u128"
2069 | "trunc_u64" => JitOp::F64ToU64,
2070 "round_u64" => JitOp::RoundToU64,
2073
2074 "f64_to_i64" | "__f64_to_i64" | "f64_to_i32" | "__f64_to_i32" | "f32_to_i64"
2075 | "__f32_to_i64" | "f32_to_i32" | "__f32_to_i32" | "__f64_to_i16" | "f64_to_i16"
2076 | "__f64_to_i8" | "f64_to_i8" | "__f32_to_i16" | "f32_to_i16" | "__f32_to_i8"
2077 | "f32_to_i8" | "__f16_to_i64" | "__f16_to_i32" | "__f16_to_i16" | "__f16_to_i8"
2078 | "__f64_to_i128" | "__f32_to_i128" | "__f16_to_i128" => JitOp::F64ToI64,
2079
2080 "f64_to_f32" | "__f64_to_f32" | "__f16_to_f32" | "__f16_to_f64" | "__f32_to_f16"
2081 | "__f64_to_f16" => JitOp::Identity,
2082
2083 "u32_to_u64" | "__u32_to_u64" | "u64_to_u32" | "__u64_to_u32" | "u32_to_i32"
2084 | "__u32_to_i32" | "i32_to_u32" | "__i32_to_u32" | "u64_to_i64" | "__u64_to_i64"
2085 | "i64_to_u64" | "__i64_to_u64" | "bool_to_u64" | "__bool_to_u64" | "bool_to_i64"
2086 | "__bool_to_i64" | "bool_to_u32" | "__bool_to_u32" | "bool_to_i32" | "__bool_to_i32"
2087 | "__u64_to_u16" | "__u64_to_u8" | "__u64_to_i16" | "__u64_to_i8" | "__i64_to_u32"
2088 | "__i64_to_u16" | "__i64_to_u8" | "__i64_to_i16" | "__i64_to_i8" | "__u32_to_u16"
2089 | "__u32_to_u8" | "__u32_to_i16" | "__u32_to_i8" | "__i32_to_u16" | "__i32_to_u8"
2090 | "__i32_to_i16" | "__i32_to_i8" | "__u16_to_u8" | "__u16_to_i8" | "__i16_to_u8"
2091 | "__i16_to_i8" | "__u128_to_u64" | "__u128_to_i64" | "__i128_to_u64" | "__i128_to_i64"
2092 | "__u128_to_u32" | "__u128_to_u16" | "__u128_to_u8" | "__u128_to_i32"
2093 | "__u128_to_i16" | "__u128_to_i8" | "__i128_to_u32" | "__i128_to_u16" | "__i128_to_u8"
2094 | "__i128_to_i32" | "__i128_to_i16" | "__i128_to_i8" | "__u64_to_u128"
2095 | "__u64_to_i128" | "__i64_to_u128" | "__i64_to_i128" | "__u128_to_i128"
2096 | "__i128_to_u128" | "__bool_to_u16" | "__bool_to_u8" | "__bool_to_i16"
2097 | "__bool_to_i8" | "__bool_to_u128" | "__bool_to_i128" | "__u8_to_f16" | "__u16_to_f16"
2098 | "__i8_to_f16" | "__i16_to_f16" | "__u64_to_f16" | "__i64_to_f16" | "__u32_to_f16"
2099 | "__i32_to_f16" | "__bool_to_f16" => JitOp::Identity,
2100
2101 "i32_to_i64" | "__i32_to_i64" | "i32_to_u64" | "__i32_to_u64" | "__u32_to_i64"
2102 | "u32_to_i64" | "__u32_to_u128" | "__u32_to_i128" | "__i32_to_u128" | "__i32_to_i128" => {
2103 JitOp::SignExtendI32
2104 }
2105
2106 "__i16_to_i32" | "i16_to_i32" | "__i16_to_i64" | "i16_to_i64" | "__i16_to_u32"
2107 | "i16_to_u32" | "__i16_to_u64" | "i16_to_u64" | "__i16_to_u128" | "__i16_to_i128" => {
2108 JitOp::SignExtendI16
2109 }
2110
2111 "__i8_to_i16" | "i8_to_i16" | "__i8_to_i32" | "i8_to_i32" | "__i8_to_i64" | "i8_to_i64"
2112 | "__i8_to_u16" | "i8_to_u16" | "__i8_to_u32" | "i8_to_u32" | "__i8_to_u64"
2113 | "i8_to_u64" | "__i8_to_u128" | "__i8_to_i128" => JitOp::SignExtendI8,
2114
2115 "__u16_to_u32" | "u16_to_u32" | "__u16_to_u64" | "u16_to_u64" | "__u16_to_i32"
2116 | "u16_to_i32" | "__u16_to_i64" | "u16_to_i64" | "__u16_to_u128" | "__u16_to_i128"
2117 | "__u16_to_i16" | "u16_to_i16" | "__i16_to_u16" | "i16_to_u16" => JitOp::ZeroExtendU16,
2118
2119 "__u8_to_u16" | "u8_to_u16" | "__u8_to_u32" | "u8_to_u32" | "__u8_to_u64" | "u8_to_u64"
2120 | "__u8_to_i16" | "u8_to_i16" | "__u8_to_i32" | "u8_to_i32" | "__u8_to_i64"
2121 | "u8_to_i64" | "__u8_to_u128" | "__u8_to_i128" | "__u8_to_i8" | "u8_to_i8"
2122 | "__i8_to_u8" | "i8_to_u8" => JitOp::ZeroExtendU8,
2123
2124 "u64_to_i32" | "__u64_to_i32" | "i64_to_i32" | "__i64_to_i32" => JitOp::ZeroExtendU32,
2125
2126 "u64_to_bool" | "__u64_to_bool" | "i64_to_bool" | "__i64_to_bool" | "u32_to_bool"
2127 | "__u32_to_bool" | "i32_to_bool" | "__i32_to_bool" | "f64_to_bool" | "__f64_to_bool"
2128 | "f32_to_bool" | "__f32_to_bool" | "__f16_to_bool" | "f16_to_bool" | "__u8_to_bool"
2129 | "__u16_to_bool" | "__i8_to_bool" | "__i16_to_bool" | "__u128_to_bool"
2130 | "__i128_to_bool" => JitOp::ToBool,
2131
2132 "weighted_pick" => {
2133 if consts.len() >= 5 {
2134 JitOp::WeightedPickConst(consts[0], consts[1], consts[2], consts[3], consts[4])
2135 } else {
2136 JitOp::Fallback
2137 }
2138 }
2139
2140 "is_positive" => {
2146 let name = node.meta().ins.iter().find_map(|slot| match slot {
2147 crate::ast::Slot::Const {
2148 name,
2149 value: crate::ast::ConstValue::Str(v),
2150 } if name == "name" => Some(v),
2151 _ => None,
2152 });
2153 match name {
2154 Some(v) => JitOp::IsPositiveCheck {
2155 name_ptr: v.as_ptr() as u64,
2156 name_len: v.len() as u64,
2157 },
2158 None => JitOp::IsPositiveCheck {
2159 name_ptr: 0,
2160 name_len: 0,
2161 },
2162 }
2163 }
2164 "in_range" => {
2165 if consts.len() >= 2 {
2166 JitOp::InRangeCheck(consts[0], consts[1])
2167 } else {
2168 JitOp::Fallback
2169 }
2170 }
2171 "is_one_of" => {
2172 if consts.is_empty() {
2173 JitOp::Fallback
2174 } else {
2175 let set = node.meta().ins.iter().find_map(|slot| match slot {
2176 crate::ast::Slot::Const {
2177 name,
2178 value: crate::ast::ConstValue::VecU64(v),
2179 } if name == "allowed" => Some(v),
2180 _ => None,
2181 });
2182 let (set_ptr, set_len) = match set {
2183 Some(v) => (v.as_ptr() as u64, v.len() as u64),
2184 None => (0, 0),
2185 };
2186 JitOp::IsOneOfCheck {
2187 allowed: consts,
2188 set_ptr,
2189 set_len,
2190 }
2191 }
2192 }
2193 _ => JitOp::Fallback,
2206 }
2207}
2208
2209#[doc(hidden)]
2217pub fn compile_jit_raw(
2218 coord_count: usize,
2219 total_slots: usize,
2220 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2221 output_map: HashMap<String, usize>,
2222 nodes: Vec<Box<dyn PolydatNode>>,
2223) -> Result<JitKernelRaw, String> {
2224 compile_jit_raw_with(
2225 coord_count,
2226 total_slots,
2227 steps,
2228 output_map,
2229 nodes,
2230 crate::compile::externs::Externs::default(),
2231 super::kernels::ScratchPlan::default(),
2232 Vec::new(),
2233 )
2234}
2235
2236#[allow(clippy::too_many_arguments)]
2239pub(crate) fn compile_jit_raw_with(
2240 coord_count: usize,
2241 total_slots: usize,
2242 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2243 output_map: HashMap<String, usize>,
2244 nodes: Vec<Box<dyn PolydatNode>>,
2245 externs: crate::compile::externs::Externs,
2246 scratch: super::kernels::ScratchPlan,
2247 volatile: Vec<usize>,
2248) -> Result<JitKernelRaw, String> {
2249 let (raw_fn, _, code) = compile_jit_impl(&steps, false, Some(total_slots))?;
2250 let mut core = JitCore::new(
2251 total_slots,
2252 coord_count,
2253 output_map,
2254 code,
2255 nodes,
2256 scratch,
2257 volatile,
2258 );
2259 core.set_externs(externs);
2260 Ok(JitKernelRaw {
2261 core,
2262 code_fn: raw_fn,
2263 })
2264}
2265
2266pub(crate) type JitSegmentCode = (NativeFn, super::kernels::JitCode);
2270
2271pub(crate) fn compile_jit_entry(
2275 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2276 tracker: Option<usize>,
2277) -> Result<JitSegmentCode, String> {
2278 let (raw_fn, _, code) = compile_jit_impl(steps, false, tracker)?;
2279 Ok((raw_fn, code))
2280}
2281
2282#[allow(clippy::too_many_arguments)]
2284pub(crate) fn compile_jit_push(
2285 coord_count: usize,
2286 total_slots: usize,
2287 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2288 output_map: HashMap<String, usize>,
2289 nodes: Vec<Box<dyn PolydatNode>>,
2290 input_dependents: Vec<Vec<usize>>,
2291 externs: crate::compile::externs::Externs,
2292 scratch: super::kernels::ScratchPlan,
2293 volatile: Vec<usize>,
2294) -> Result<JitKernelPush, String> {
2295 let step_count = steps.len();
2296 let (_, prov_fn, code) = compile_jit_impl(&steps, true, Some(total_slots))?;
2297 let mut core = JitCore::new(
2298 total_slots,
2299 coord_count,
2300 output_map,
2301 code,
2302 nodes,
2303 scratch,
2304 volatile,
2305 );
2306 core.set_externs(externs);
2307 Ok(JitKernelPush {
2308 core,
2309 code_fn_prov: prov_fn,
2310 node_clean: vec![0u8; step_count],
2311 input_dependents,
2312 })
2313}
2314
2315#[allow(clippy::too_many_arguments)]
2317pub(crate) fn compile_jit_pull(
2318 coord_count: usize,
2319 total_slots: usize,
2320 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2321 output_map: HashMap<String, usize>,
2322 nodes: Vec<Box<dyn PolydatNode>>,
2323 input_dependents: &[Vec<usize>],
2324 externs: crate::compile::externs::Externs,
2325 scratch: super::kernels::ScratchPlan,
2326 volatile: Vec<usize>,
2327) -> Result<JitKernelPull, String> {
2328 let buffer_len = total_slots;
2329 let (raw_fn, _, code) = compile_jit_impl(&steps, false, Some(total_slots))?;
2331 let step_outs: Vec<&[usize]> = steps.iter().map(|(_, _, o)| o.as_slice()).collect();
2332 let slot_provenance =
2333 crate::compile::slot_provenance(coord_count, buffer_len, &step_outs, input_dependents);
2334 let mut core = JitCore::new(
2335 total_slots,
2336 coord_count,
2337 output_map,
2338 code,
2339 nodes,
2340 scratch,
2341 volatile,
2342 );
2343 core.set_externs(externs);
2344 Ok(JitKernelPull {
2345 core,
2346 code_fn: raw_fn,
2347 slot_provenance,
2348 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
2349 force_run: false,
2350 })
2351}
2352
2353#[allow(clippy::too_many_arguments)]
2355pub(crate) fn compile_jit_push_pull(
2356 coord_count: usize,
2357 total_slots: usize,
2358 steps: Vec<(JitOp, Vec<usize>, Vec<usize>)>,
2359 output_map: HashMap<String, usize>,
2360 nodes: Vec<Box<dyn PolydatNode>>,
2361 input_dependents: Vec<Vec<usize>>,
2362 externs: crate::compile::externs::Externs,
2363 scratch: super::kernels::ScratchPlan,
2364 volatile: Vec<usize>,
2365) -> Result<JitKernelPushPull, String> {
2366 let step_count = steps.len();
2367 let buffer_len = total_slots;
2368 let (_, prov_fn, code) = compile_jit_impl(&steps, true, Some(total_slots))?;
2369 let step_outs: Vec<&[usize]> = steps.iter().map(|(_, _, o)| o.as_slice()).collect();
2370 let slot_provenance =
2371 crate::compile::slot_provenance(coord_count, buffer_len, &step_outs, &input_dependents);
2372 let mut core = JitCore::new(
2373 total_slots,
2374 coord_count,
2375 output_map,
2376 code,
2377 nodes,
2378 scratch,
2379 volatile,
2380 );
2381 core.set_externs(externs);
2382 Ok(JitKernelPushPull {
2383 core,
2384 code_fn_prov: prov_fn,
2385 node_clean: vec![0u8; step_count],
2386 input_dependents,
2387 slot_provenance,
2388 changed_mask: crate::kernel::ProvMask::all_below(coord_count),
2389 force_run: false,
2390 })
2391}
2392
2393pub type NativeFn = unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf);
2397pub type NativeProvFn = unsafe fn(*const u64, *mut u64, *mut crate::ast::ScratchBuf, *mut u8);
2399
2400type JitCompiled = (NativeFn, NativeProvFn, super::kernels::JitCode);
2404
2405fn compile_jit_impl(
2409 steps: &[(JitOp, Vec<usize>, Vec<usize>)],
2410 provenance: bool,
2411 tracker: Option<usize>,
2412) -> Result<JitCompiled, String> {
2413 let mut flag_builder = settings::builder();
2414 flag_builder.set("opt_level", "speed").unwrap();
2415 flag_builder.set("unwind_info", "true").unwrap();
2421 flag_builder.set("preserve_frame_pointers", "true").unwrap();
2422 let isa = super::host_isa::build_host_isa(flag_builder)?;
2423
2424 let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
2425
2426 jit_builder.symbol("jit_xxh3_hash", jit_xxh3_hash as *const u8);
2428 jit_builder.symbol("jit_interleave", jit_interleave as *const u8);
2429 jit_builder.symbol("jit_shuffle", jit_shuffle as *const u8);
2430 jit_builder.symbol("jit_lut_sample", jit_lut_sample as *const u8);
2431 jit_builder.symbol("jit_weighted_pick", jit_weighted_pick as *const u8);
2432 jit_builder.symbol("jit_pcg", jit_pcg as *const u8);
2433 jit_builder.symbol("jit_pcg_stream", jit_pcg_stream as *const u8);
2434 jit_builder.symbol("jit_n_of", jit_n_of as *const u8);
2435 jit_builder.symbol("jit_cycle_walk", jit_cycle_walk as *const u8);
2436 jit_builder.symbol("jit_perlin_1d", jit_perlin_1d as *const u8);
2437 jit_builder.symbol("jit_perlin_2d", jit_perlin_2d as *const u8);
2438 jit_builder.symbol("jit_simplex_2d", jit_simplex_2d as *const u8);
2439 jit_builder.symbol("jit_fractal_noise_1d", jit_fractal_noise_1d as *const u8);
2440 jit_builder.symbol("jit_fractal_noise_2d", jit_fractal_noise_2d as *const u8);
2441 jit_builder.symbol("jit_thread_id", jit_thread_id as *const u8);
2442 jit_builder.symbol(
2443 "jit_current_epoch_millis",
2444 jit_current_epoch_millis as *const u8,
2445 );
2446 jit_builder.symbol("jit_is_positive_fail", jit_is_positive_fail as *const u8);
2450 jit_builder.symbol("jit_in_range_fail", jit_in_range_fail as *const u8);
2451 jit_builder.symbol("jit_is_one_of_fail", jit_is_one_of_fail as *const u8);
2452 jit_builder.symbol("jit_slot_call", jit_slot_call as *const u8);
2455 jit_builder.symbol("jit_u64_to_str", jit_u64_to_str as *const u8);
2456 jit_builder.symbol("jit_i64_to_str", jit_i64_to_str as *const u8);
2457 jit_builder.symbol("jit_f64_to_str", jit_f64_to_str as *const u8);
2458 jit_builder.symbol("jit_str_concat", jit_str_concat as *const u8);
2459 jit_builder.symbol("jit_json_to_str", jit_json_to_str as *const u8);
2460 jit_builder.symbol("jit_vec_add", jit_vec_add as *const u8);
2461 jit_builder.symbol("jit_vec_scale", jit_vec_scale as *const u8);
2462 jit_builder.symbol("jit_vec_norm", jit_vec_norm as *const u8);
2463 jit_builder.symbol("jit_hash_vec", jit_hash_vec as *const u8);
2464 jit_builder.symbol("jit_xxhash3_vec", jit_xxhash3_vec as *const u8);
2465 jit_builder.symbol("jit_reg_to_vec_f32", jit_reg_to_vec_f32 as *const u8);
2466 jit_builder.symbol("jit_vec_dot", jit_vec_dot as *const u8);
2467 jit_builder.symbol("jit_vec_l2", jit_vec_l2 as *const u8);
2468 jit_builder.symbol("jit_vec_cosine", jit_vec_cosine as *const u8);
2469 jit_builder.symbol("jit_lid_mle", jit_lid_mle as *const u8);
2470 jit_builder.symbol("jit_reg_lane_f32", jit_reg_lane_f32 as *const u8);
2471 jit_builder.symbol("jit_reg_lane_i16", jit_reg_lane_i16 as *const u8);
2472 jit_builder.symbol("jit_reg_lane_i64", jit_reg_lane_i64 as *const u8);
2473 jit_builder.symbol("jit_reg_with_lane_f32", jit_reg_with_lane_f32 as *const u8);
2474 jit_builder.symbol("jit_reg_gather_f32", jit_reg_gather_f32 as *const u8);
2475 jit_builder.symbol("jit_vec_to_reg_f32", jit_vec_to_reg_f32 as *const u8);
2476 jit_builder.symbol("jit_reg_mul_i8", jit_reg_mul_i8 as *const u8);
2477 jit_builder.symbol("jit_sin", jit_sin as *const u8);
2479 jit_builder.symbol("jit_cos", jit_cos as *const u8);
2480 jit_builder.symbol("jit_tan", jit_tan as *const u8);
2481 jit_builder.symbol("jit_asin", jit_asin as *const u8);
2482 jit_builder.symbol("jit_acos", jit_acos as *const u8);
2483 jit_builder.symbol("jit_atan", jit_atan as *const u8);
2484 jit_builder.symbol("jit_sqrt", jit_sqrt as *const u8);
2485 jit_builder.symbol("jit_abs_f64", jit_abs_f64 as *const u8);
2486 jit_builder.symbol("jit_ln", jit_ln as *const u8);
2487 jit_builder.symbol("jit_exp", jit_exp as *const u8);
2488 jit_builder.symbol("jit_floor_base10", jit_floor_base10 as *const u8);
2489 jit_builder.symbol("jit_ceiling_base10", jit_ceiling_base10 as *const u8);
2490 jit_builder.symbol("jit_closest_base10", jit_closest_base10 as *const u8);
2491 jit_builder.symbol("jit_floor_decade", jit_floor_decade as *const u8);
2492 jit_builder.symbol("jit_ceiling_decade", jit_ceiling_decade as *const u8);
2493 jit_builder.symbol("jit_closest_decade", jit_closest_decade as *const u8);
2494 jit_builder.symbol("jit_floor_binomial", jit_floor_binomial as *const u8);
2495 jit_builder.symbol("jit_ceiling_binomial", jit_ceiling_binomial as *const u8);
2496 jit_builder.symbol("jit_closest_binomial", jit_closest_binomial as *const u8);
2497 jit_builder.symbol("jit_floor_fibonacci", jit_floor_fibonacci as *const u8);
2498 jit_builder.symbol("jit_ceiling_fibonacci", jit_ceiling_fibonacci as *const u8);
2499 jit_builder.symbol("jit_closest_fibonacci", jit_closest_fibonacci as *const u8);
2500 jit_builder.symbol("jit_atan2", jit_atan2 as *const u8);
2501 jit_builder.symbol("jit_pow", jit_pow as *const u8);
2502 jit_builder.symbol("jit_round_nearest", jit_round_nearest as *const u8);
2503 jit_builder.symbol("jit_round_floor", jit_round_floor as *const u8);
2504 jit_builder.symbol("jit_round_ceiling", jit_round_ceiling as *const u8);
2505 jit_builder.symbol("jit_f64_mod", jit_f64_mod as *const u8);
2506 jit_builder.symbol("jit_div_zero_fail", jit_div_zero_fail as *const u8);
2507
2508 let mut module = JITModule::new(jit_builder);
2509
2510 let hash_func_id = {
2512 let mut sig = module.make_signature();
2513 sig.params.push(AbiParam::new(types::I64));
2514 sig.returns.push(AbiParam::new(types::I64));
2515 module
2516 .declare_function("jit_xxh3_hash", Linkage::Import, &sig)
2517 .map_err(|e| format!("declare hash: {e}"))?
2518 };
2519
2520 let interleave_func_id = {
2522 let mut sig = module.make_signature();
2523 sig.params.push(AbiParam::new(types::I64));
2524 sig.params.push(AbiParam::new(types::I64));
2525 sig.returns.push(AbiParam::new(types::I64));
2526 module
2527 .declare_function("jit_interleave", Linkage::Import, &sig)
2528 .map_err(|e| format!("declare interleave: {e}"))?
2529 };
2530
2531 let shuffle_func_id = {
2533 let mut sig = module.make_signature();
2534 for _ in 0..4 {
2535 sig.params.push(AbiParam::new(types::I64));
2536 }
2537 sig.returns.push(AbiParam::new(types::I64));
2538 module
2539 .declare_function("jit_shuffle", Linkage::Import, &sig)
2540 .map_err(|e| format!("declare shuffle: {e}"))?
2541 };
2542
2543 let lut_sample_func_id = {
2545 let mut sig = module.make_signature();
2546 for _ in 0..3 {
2547 sig.params.push(AbiParam::new(types::I64));
2548 }
2549 sig.returns.push(AbiParam::new(types::I64));
2550 module
2551 .declare_function("jit_lut_sample", Linkage::Import, &sig)
2552 .map_err(|e| format!("declare lut_sample: {e}"))?
2553 };
2554
2555 let weighted_pick_func_id = {
2557 let mut sig = module.make_signature();
2558 for _ in 0..6 {
2559 sig.params.push(AbiParam::new(types::I64));
2560 }
2561 sig.returns.push(AbiParam::new(types::I64));
2562 module
2563 .declare_function("jit_weighted_pick", Linkage::Import, &sig)
2564 .map_err(|e| format!("declare weighted_pick: {e}"))?
2565 };
2566
2567 let pcg_func_id = {
2568 let mut sig = module.make_signature();
2569 for _ in 0..3 {
2570 sig.params.push(AbiParam::new(types::I64));
2571 }
2572 sig.returns.push(AbiParam::new(types::I64));
2573 module
2574 .declare_function("jit_pcg", Linkage::Import, &sig)
2575 .map_err(|e| format!("declare pcg: {e}"))?
2576 };
2577 let pcg_stream_func_id = {
2578 let mut sig = module.make_signature();
2579 for _ in 0..3 {
2580 sig.params.push(AbiParam::new(types::I64));
2581 }
2582 sig.returns.push(AbiParam::new(types::I64));
2583 module
2584 .declare_function("jit_pcg_stream", Linkage::Import, &sig)
2585 .map_err(|e| format!("declare pcg_stream: {e}"))?
2586 };
2587 let n_of_func_id = {
2588 let mut sig = module.make_signature();
2589 for _ in 0..3 {
2590 sig.params.push(AbiParam::new(types::I64));
2591 }
2592 sig.returns.push(AbiParam::new(types::I64));
2593 module
2594 .declare_function("jit_n_of", Linkage::Import, &sig)
2595 .map_err(|e| format!("declare n_of: {e}"))?
2596 };
2597 let cycle_walk_func_id = {
2598 let mut sig = module.make_signature();
2599 for _ in 0..4 {
2600 sig.params.push(AbiParam::new(types::I64));
2601 }
2602 sig.returns.push(AbiParam::new(types::I64));
2603 module
2604 .declare_function("jit_cycle_walk", Linkage::Import, &sig)
2605 .map_err(|e| format!("declare cycle_walk: {e}"))?
2606 };
2607 let perlin_1d_func_id = {
2608 let mut sig = module.make_signature();
2609 for _ in 0..3 {
2610 sig.params.push(AbiParam::new(types::I64));
2611 }
2612 sig.returns.push(AbiParam::new(types::I64));
2613 module
2614 .declare_function("jit_perlin_1d", Linkage::Import, &sig)
2615 .map_err(|e| format!("declare perlin_1d: {e}"))?
2616 };
2617 let perlin_2d_func_id = {
2618 let mut sig = module.make_signature();
2619 for _ in 0..4 {
2620 sig.params.push(AbiParam::new(types::I64));
2621 }
2622 sig.returns.push(AbiParam::new(types::I64));
2623 module
2624 .declare_function("jit_perlin_2d", Linkage::Import, &sig)
2625 .map_err(|e| format!("declare perlin_2d: {e}"))?
2626 };
2627 let simplex_2d_func_id = {
2628 let mut sig = module.make_signature();
2629 for _ in 0..4 {
2630 sig.params.push(AbiParam::new(types::I64));
2631 }
2632 sig.returns.push(AbiParam::new(types::I64));
2633 module
2634 .declare_function("jit_simplex_2d", Linkage::Import, &sig)
2635 .map_err(|e| format!("declare simplex_2d: {e}"))?
2636 };
2637 let fractal_noise_1d_func_id = {
2638 let mut sig = module.make_signature();
2639 for _ in 0..4 {
2640 sig.params.push(AbiParam::new(types::I64));
2641 }
2642 sig.returns.push(AbiParam::new(types::I64));
2643 module
2644 .declare_function("jit_fractal_noise_1d", Linkage::Import, &sig)
2645 .map_err(|e| format!("declare fractal_noise_1d: {e}"))?
2646 };
2647 let fractal_noise_2d_func_id = {
2648 let mut sig = module.make_signature();
2649 for _ in 0..5 {
2650 sig.params.push(AbiParam::new(types::I64));
2651 }
2652 sig.returns.push(AbiParam::new(types::I64));
2653 module
2654 .declare_function("jit_fractal_noise_2d", Linkage::Import, &sig)
2655 .map_err(|e| format!("declare fractal_noise_2d: {e}"))?
2656 };
2657 let thread_id_func_id = {
2658 let mut sig = module.make_signature();
2659 sig.returns.push(AbiParam::new(types::I64));
2660 module
2661 .declare_function("jit_thread_id", Linkage::Import, &sig)
2662 .map_err(|e| format!("declare thread_id: {e}"))?
2663 };
2664 let current_epoch_millis_func_id = {
2665 let mut sig = module.make_signature();
2666 sig.returns.push(AbiParam::new(types::I64));
2667 module
2668 .declare_function("jit_current_epoch_millis", Linkage::Import, &sig)
2669 .map_err(|e| format!("declare current_epoch_millis: {e}"))?
2670 };
2671
2672 let math_unary_names = [
2674 "jit_sin",
2675 "jit_cos",
2676 "jit_tan",
2677 "jit_asin",
2678 "jit_acos",
2679 "jit_atan",
2680 "jit_sqrt",
2681 "jit_abs_f64",
2682 "jit_ln",
2683 "jit_exp",
2684 "jit_floor_base10",
2685 "jit_ceiling_base10",
2686 "jit_closest_base10",
2687 "jit_floor_decade",
2688 "jit_ceiling_decade",
2689 "jit_closest_decade",
2690 "jit_floor_binomial",
2691 "jit_ceiling_binomial",
2692 "jit_closest_binomial",
2693 "jit_floor_fibonacci",
2694 "jit_ceiling_fibonacci",
2695 "jit_closest_fibonacci",
2696 ];
2697 let mut math_unary_ids = Vec::new();
2698 for name in &math_unary_names {
2699 let mut sig = module.make_signature();
2700 sig.params.push(AbiParam::new(types::I64));
2701 sig.returns.push(AbiParam::new(types::I64));
2702 math_unary_ids.push(
2703 module
2704 .declare_function(name, Linkage::Import, &sig)
2705 .map_err(|e| format!("declare {name}: {e}"))?,
2706 );
2707 }
2708
2709 let is_positive_fail_id = {
2713 let mut sig = module.make_signature();
2714 sig.params.push(AbiParam::new(types::I64));
2715 sig.params.push(AbiParam::new(types::I64));
2716 sig.params.push(AbiParam::new(types::I64));
2717 sig.returns.push(AbiParam::new(types::I64));
2718 module
2719 .declare_function("jit_is_positive_fail", Linkage::Import, &sig)
2720 .map_err(|e| format!("declare is_positive_fail: {e}"))?
2721 };
2722
2723 let in_range_fail_id = {
2725 let mut sig = module.make_signature();
2726 for _ in 0..3 {
2727 sig.params.push(AbiParam::new(types::I64));
2728 }
2729 sig.returns.push(AbiParam::new(types::I64));
2730 module
2731 .declare_function("jit_in_range_fail", Linkage::Import, &sig)
2732 .map_err(|e| format!("declare in_range_fail: {e}"))?
2733 };
2734
2735 let is_one_of_fail_id = {
2738 let mut sig = module.make_signature();
2739 sig.params.push(AbiParam::new(types::I64));
2740 sig.params.push(AbiParam::new(types::I64));
2741 sig.params.push(AbiParam::new(types::I64));
2742 sig.returns.push(AbiParam::new(types::I64));
2743 module
2744 .declare_function("jit_is_one_of_fail", Linkage::Import, &sig)
2745 .map_err(|e| format!("declare is_one_of_fail: {e}"))?
2746 };
2747
2748 let math_binary_names = [
2750 "jit_atan2",
2751 "jit_pow",
2752 "jit_round_nearest",
2753 "jit_round_floor",
2754 "jit_round_ceiling",
2755 "jit_f64_mod",
2756 ];
2757 const F64_MOD_HELPER: usize = 5;
2758
2759 let div_zero_fail_id = {
2761 let mut sig = module.make_signature();
2762 sig.params.push(AbiParam::new(types::I64));
2763 sig.returns.push(AbiParam::new(types::I64));
2764 module
2765 .declare_function("jit_div_zero_fail", Linkage::Import, &sig)
2766 .map_err(|e| format!("declare div_zero_fail: {e}"))?
2767 };
2768 let mut math_binary_ids = Vec::new();
2769 for name in &math_binary_names {
2770 let mut sig = module.make_signature();
2771 sig.params.push(AbiParam::new(types::I64));
2772 sig.params.push(AbiParam::new(types::I64));
2773 sig.returns.push(AbiParam::new(types::I64));
2774 math_binary_ids.push(
2775 module
2776 .declare_function(name, Linkage::Import, &sig)
2777 .map_err(|e| format!("declare {name}: {e}"))?,
2778 );
2779 }
2780
2781 let slot_call_id = {
2784 let mut sig = module.make_signature();
2785 for _ in 0..8 {
2786 sig.params.push(AbiParam::new(types::I64));
2787 }
2788 module
2789 .declare_function("jit_slot_call", Linkage::Import, &sig)
2790 .map_err(|e| format!("declare jit_slot_call: {e}"))?
2791 };
2792
2793 let mut declare_str = |name: &str, args: usize| -> Result<cranelift_module::FuncId, String> {
2797 let mut sig = module.make_signature();
2798 for _ in 0..args {
2799 sig.params.push(AbiParam::new(types::I64));
2800 }
2801 module
2802 .declare_function(name, Linkage::Import, &sig)
2803 .map_err(|e| format!("declare {name}: {e}"))
2804 };
2805 let u64_to_str_id = declare_str("jit_u64_to_str", 5)?;
2806 let i64_to_str_id = declare_str("jit_i64_to_str", 5)?;
2807 let f64_to_str_id = declare_str("jit_f64_to_str", 5)?;
2808 let str_concat_id = declare_str("jit_str_concat", 6)?;
2809 let json_to_str_id = declare_str("jit_json_to_str", 6)?;
2810
2811 let mut declare_words =
2816 |name: &str, args: usize, returns: bool| -> Result<cranelift_module::FuncId, String> {
2817 let mut sig = module.make_signature();
2818 for _ in 0..args {
2819 sig.params.push(AbiParam::new(types::I64));
2820 }
2821 if returns {
2822 sig.returns.push(AbiParam::new(types::I64));
2823 }
2824 module
2825 .declare_function(name, Linkage::Import, &sig)
2826 .map_err(|e| format!("declare {name}: {e}"))
2827 };
2828 let vec_producer_ids = [
2829 (VecProducer::Add, declare_words("jit_vec_add", 8, false)?),
2830 (
2831 VecProducer::Scale,
2832 declare_words("jit_vec_scale", 8, false)?,
2833 ),
2834 (VecProducer::Norm, declare_words("jit_vec_norm", 8, false)?),
2835 (
2836 VecProducer::HashVec,
2837 declare_words("jit_hash_vec", 8, false)?,
2838 ),
2839 (
2840 VecProducer::XxHash3Vec,
2841 declare_words("jit_xxhash3_vec", 8, false)?,
2842 ),
2843 (
2844 VecProducer::RegToVec,
2845 declare_words("jit_reg_to_vec_f32", 8, false)?,
2846 ),
2847 ];
2848 let vec_reducer_ids = [
2849 (VecReducer::Dot, declare_words("jit_vec_dot", 4, true)?),
2850 (VecReducer::L2, declare_words("jit_vec_l2", 4, true)?),
2851 (
2852 VecReducer::Cosine,
2853 declare_words("jit_vec_cosine", 4, true)?,
2854 ),
2855 (VecReducer::LidMle, declare_words("jit_lid_mle", 4, true)?),
2856 ];
2857 let reg_lane_ids = [
2858 (
2859 RegLaneRead::F32,
2860 declare_words("jit_reg_lane_f32", 3, true)?,
2861 ),
2862 (
2863 RegLaneRead::I16,
2864 declare_words("jit_reg_lane_i16", 3, true)?,
2865 ),
2866 (
2867 RegLaneRead::I64,
2868 declare_words("jit_reg_lane_i64", 3, true)?,
2869 ),
2870 ];
2871 let reg_producer_ids = [
2872 (
2873 RegProducer::WithLaneF32,
2874 declare_words("jit_reg_with_lane_f32", 6, false)?,
2875 ),
2876 (
2877 RegProducer::GatherF32,
2878 declare_words("jit_reg_gather_f32", 6, false)?,
2879 ),
2880 (
2881 RegProducer::VecToRegF32,
2882 declare_words("jit_vec_to_reg_f32", 6, false)?,
2883 ),
2884 (
2885 RegProducer::MulI8,
2886 declare_words("jit_reg_mul_i8", 6, false)?,
2887 ),
2888 ];
2889
2890 let mut sig = module.make_signature();
2894 sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); sig.params.push(AbiParam::new(types::I64)); if provenance {
2898 sig.params.push(AbiParam::new(types::I64)); }
2900 let func_id = module
2901 .declare_function("polydat_kernel", Linkage::Local, &sig)
2902 .map_err(|e| format!("declare kernel: {e}"))?;
2903
2904 let mut ctx = module.make_context();
2905 ctx.func.signature = sig;
2906
2907 let mut fb_ctx = FunctionBuilderContext::new();
2908 {
2909 let mut builder = FunctionBuilder::new(&mut ctx.func, &mut fb_ctx);
2910 let block = builder.create_block();
2911 builder.append_block_params_for_function_params(block);
2912 builder.switch_to_block(block);
2913 builder.seal_block(block);
2914
2915 let _coords_ptr = builder.block_params(block)[0];
2916 let buffer_ptr = builder.block_params(block)[1];
2917 let scratch_ptr = builder.block_params(block)[2];
2918 let clean_ptr = if provenance {
2919 Some(builder.block_params(block)[3])
2920 } else {
2921 None
2922 };
2923
2924 let hash_func_ref = module.declare_func_in_func(hash_func_id, builder.func);
2926 let interleave_func_ref = module.declare_func_in_func(interleave_func_id, builder.func);
2927 let shuffle_func_ref = module.declare_func_in_func(shuffle_func_id, builder.func);
2928 let lut_sample_func_ref = module.declare_func_in_func(lut_sample_func_id, builder.func);
2929 let weighted_pick_func_ref =
2930 module.declare_func_in_func(weighted_pick_func_id, builder.func);
2931 let is_positive_fail_ref = module.declare_func_in_func(is_positive_fail_id, builder.func);
2932 let in_range_fail_ref = module.declare_func_in_func(in_range_fail_id, builder.func);
2933 let div_zero_fail_ref = module.declare_func_in_func(div_zero_fail_id, builder.func);
2934 let is_one_of_fail_ref = module.declare_func_in_func(is_one_of_fail_id, builder.func);
2935 let slot_call_ref = module.declare_func_in_func(slot_call_id, builder.func);
2936 let u64_to_str_ref = module.declare_func_in_func(u64_to_str_id, builder.func);
2937 let i64_to_str_ref = module.declare_func_in_func(i64_to_str_id, builder.func);
2938 let f64_to_str_ref = module.declare_func_in_func(f64_to_str_id, builder.func);
2939 let str_concat_ref = module.declare_func_in_func(str_concat_id, builder.func);
2940 let json_to_str_ref = module.declare_func_in_func(json_to_str_id, builder.func);
2941 let vec_producer_refs: Vec<(VecProducer, ir::FuncRef)> = vec_producer_ids
2942 .iter()
2943 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2944 .collect();
2945 let vec_reducer_refs: Vec<(VecReducer, ir::FuncRef)> = vec_reducer_ids
2946 .iter()
2947 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2948 .collect();
2949 let reg_lane_refs: Vec<(RegLaneRead, ir::FuncRef)> = reg_lane_ids
2950 .iter()
2951 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2952 .collect();
2953 let reg_producer_refs: Vec<(RegProducer, ir::FuncRef)> = reg_producer_ids
2954 .iter()
2955 .map(|(k, id)| (*k, module.declare_func_in_func(*id, builder.func)))
2956 .collect();
2957 let pcg_func_ref = module.declare_func_in_func(pcg_func_id, builder.func);
2958 let pcg_stream_func_ref = module.declare_func_in_func(pcg_stream_func_id, builder.func);
2959 let n_of_func_ref = module.declare_func_in_func(n_of_func_id, builder.func);
2960 let cycle_walk_func_ref = module.declare_func_in_func(cycle_walk_func_id, builder.func);
2961 let perlin_1d_func_ref = module.declare_func_in_func(perlin_1d_func_id, builder.func);
2962 let perlin_2d_func_ref = module.declare_func_in_func(perlin_2d_func_id, builder.func);
2963 let simplex_2d_func_ref = module.declare_func_in_func(simplex_2d_func_id, builder.func);
2964 let fractal_noise_1d_func_ref =
2965 module.declare_func_in_func(fractal_noise_1d_func_id, builder.func);
2966 let fractal_noise_2d_func_ref =
2967 module.declare_func_in_func(fractal_noise_2d_func_id, builder.func);
2968 let thread_id_func_ref = module.declare_func_in_func(thread_id_func_id, builder.func);
2969 let current_epoch_millis_func_ref =
2970 module.declare_func_in_func(current_epoch_millis_func_id, builder.func);
2971 let math_unary_refs: Vec<_> = math_unary_ids
2972 .iter()
2973 .map(|id| module.declare_func_in_func(*id, builder.func))
2974 .collect();
2975 let math_binary_refs: Vec<_> = math_binary_ids
2976 .iter()
2977 .map(|id| module.declare_func_in_func(*id, builder.func))
2978 .collect();
2979 for (step_idx, (jit_op, input_slots, output_slots)) in steps.iter().enumerate() {
2981 let skip_block = if let Some(cp) = clean_ptr {
2983 let skip = builder.create_block();
2984 let cont = builder.create_block();
2985 let offset = builder.ins().iconst(types::I64, step_idx as i64);
2987 let addr = builder.ins().iadd(cp, offset);
2988 let flag = builder.ins().load(types::I8, ir::MemFlags::new(), addr, 0);
2989 let zero = builder.ins().iconst(types::I8, 0);
2990 let is_clean = builder
2991 .ins()
2992 .icmp(ir::condcodes::IntCC::NotEqual, flag, zero);
2993 builder.ins().brif(is_clean, skip, &[], cont, &[]);
2994 builder.switch_to_block(cont);
2995 builder.seal_block(cont);
2996 Some(skip)
2997 } else {
2998 None
2999 };
3000 let tracker_store = tracker.map(|t| {
3004 let idx = builder.ins().iconst(types::I64, step_idx as i64);
3005 let inst = store_slot(&mut builder, buffer_ptr, t, idx);
3006 (inst, builder.func.dfg.num_insts())
3007 });
3008 match jit_op {
3009 JitOp::Identity => {
3010 for (&i, &o) in input_slots.iter().zip(output_slots.iter()) {
3013 let val = load_slot(&mut builder, buffer_ptr, i);
3014 store_slot(&mut builder, buffer_ptr, o, val);
3015 }
3016 }
3017 JitOp::AddConst(c) => {
3018 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3019 let c_val = builder.ins().iconst(types::I64, *c as i64);
3020 let result = builder.ins().iadd(val, c_val);
3021 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3022 }
3023 JitOp::MulConst(c) => {
3024 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3025 let c_val = builder.ins().iconst(types::I64, *c as i64);
3026 let result = builder.ins().imul(val, c_val);
3027 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3028 }
3029 JitOp::DivConst(c) | JitOp::ModConst(c) => {
3030 let is_div = matches!(jit_op, JitOp::DivConst(_));
3033 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3034 if *c == 0 {
3035 let kind = builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3036 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3037 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3038 } else {
3039 let c_val = builder.ins().iconst(types::I64, *c as i64);
3040 let result = if is_div {
3041 builder.ins().udiv(val, c_val)
3042 } else {
3043 builder.ins().urem(val, c_val)
3044 };
3045 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3046 }
3047 }
3048 JitOp::U64DivWire | JitOp::U64ModWire => {
3049 let is_div = matches!(jit_op, JitOp::U64DivWire);
3053 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3054 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3055 let zero = builder.ins().iconst(types::I64, 0);
3056 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3057 let fail_block = builder.create_block();
3058 let ok_block = builder.create_block();
3059 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3060 builder.switch_to_block(fail_block);
3061 builder.seal_block(fail_block);
3062 let kind = builder.ins().iconst(types::I64, if is_div { 0 } else { 1 });
3063 let _ = builder.ins().call(div_zero_fail_ref, &[kind]);
3064 builder.ins().jump(ok_block, &[]);
3065 builder.switch_to_block(ok_block);
3066 builder.seal_block(ok_block);
3067 let result = if is_div {
3068 builder.ins().udiv(a, b)
3069 } else {
3070 builder.ins().urem(a, b)
3071 };
3072 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3073 }
3074 JitOp::ClampConst(min, max) => {
3075 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3076 let min_val = builder.ins().iconst(types::I64, *min as i64);
3077 let max_val = builder.ins().iconst(types::I64, *max as i64);
3078 let clamped_lo = builder.ins().umax(val, min_val);
3079 let clamped = builder.ins().umin(clamped_lo, max_val);
3080 store_slot(&mut builder, buffer_ptr, output_slots[0], clamped);
3081 }
3082 JitOp::Interleave => {
3083 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3084 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3085 let call = builder.ins().call(interleave_func_ref, &[a, b]);
3086 let result = builder.inst_results(call)[0];
3087 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3088 }
3089 JitOp::MixedRadixConst(radixes) => {
3090 let mut remainder = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3092 for (i, &radix) in radixes.iter().enumerate() {
3093 if radix == 0 {
3094 store_slot(&mut builder, buffer_ptr, output_slots[i], remainder);
3096 } else {
3097 let r = builder.ins().iconst(types::I64, radix as i64);
3098 let digit = builder.ins().urem(remainder, r);
3099 store_slot(&mut builder, buffer_ptr, output_slots[i], digit);
3100 remainder = builder.ins().udiv(remainder, r);
3101 }
3102 }
3103 }
3104 JitOp::Hash => {
3105 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3106 let call = builder.ins().call(hash_func_ref, &[val]);
3107 let result = builder.inst_results(call)[0];
3108 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3109 }
3110 JitOp::SplitMix64 => {
3111 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3112 let c_gamma = builder
3113 .ins()
3114 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3115 let x1 = builder.ins().iadd(x0, c_gamma);
3116 let s30 = builder.ins().ushr_imm(x1, 30);
3117 let x2 = builder.ins().bxor(x1, s30);
3118 let c_m1 = builder
3119 .ins()
3120 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3121 let x3 = builder.ins().imul(x2, c_m1);
3122 let s27 = builder.ins().ushr_imm(x3, 27);
3123 let x4 = builder.ins().bxor(x3, s27);
3124 let c_m2 = builder
3125 .ins()
3126 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3127 let x5 = builder.ins().imul(x4, c_m2);
3128 let s31 = builder.ins().ushr_imm(x5, 31);
3129 let result = builder.ins().bxor(x5, s31);
3130 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3131 }
3132 JitOp::FairCoin => {
3133 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3134 let c_gamma = builder
3135 .ins()
3136 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3137 let x1 = builder.ins().iadd(x0, c_gamma);
3138 let s30 = builder.ins().ushr_imm(x1, 30);
3139 let x2 = builder.ins().bxor(x1, s30);
3140 let c_m1 = builder
3141 .ins()
3142 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3143 let x3 = builder.ins().imul(x2, c_m1);
3144 let s27 = builder.ins().ushr_imm(x3, 27);
3145 let x4 = builder.ins().bxor(x3, s27);
3146 let c_m2 = builder
3147 .ins()
3148 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3149 let x5 = builder.ins().imul(x4, c_m2);
3150 let s31 = builder.ins().ushr_imm(x5, 31);
3151 let h = builder.ins().bxor(x5, s31);
3152 let one = builder.ins().iconst(types::I64, 1);
3153 let result = builder.ins().band(h, one);
3154 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3155 }
3156 JitOp::CoinFlipConst(threshold) => {
3157 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3158 let thr = builder.ins().iconst(types::I64, *threshold as i64);
3159 let cmp = builder
3160 .ins()
3161 .icmp(ir::condcodes::IntCC::UnsignedLessThan, x, thr);
3162 let zero = builder.ins().iconst(types::I64, 0);
3163 let one = builder.ins().iconst(types::I64, 1);
3164 let result = builder.ins().select(cmp, one, zero);
3165 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3166 }
3167 JitOp::UnfairCoinConst(p_bits) => {
3168 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3169 let c_gamma = builder
3170 .ins()
3171 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3172 let x1 = builder.ins().iadd(x0, c_gamma);
3173 let s30 = builder.ins().ushr_imm(x1, 30);
3174 let x2 = builder.ins().bxor(x1, s30);
3175 let c_m1 = builder
3176 .ins()
3177 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3178 let x3 = builder.ins().imul(x2, c_m1);
3179 let s27 = builder.ins().ushr_imm(x3, 27);
3180 let x4 = builder.ins().bxor(x3, s27);
3181 let c_m2 = builder
3182 .ins()
3183 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3184 let x5 = builder.ins().imul(x4, c_m2);
3185 let s31 = builder.ins().ushr_imm(x5, 31);
3186 let h = builder.ins().bxor(x5, s31);
3187
3188 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3189 let max_f = builder.ins().f64const(u64::MAX as f64);
3190 let unit = builder.ins().fdiv(fval, max_f);
3191 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3192 let cmp = builder
3193 .ins()
3194 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3195 let zero = builder.ins().iconst(types::I64, 0);
3196 let one = builder.ins().iconst(types::I64, 1);
3197 let result = builder.ins().select(cmp, one, zero);
3198 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3199 }
3200 JitOp::ChanceConst(p_bits) => {
3201 let x0 = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3202 let c_gamma = builder
3203 .ins()
3204 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3205 let x1 = builder.ins().iadd(x0, c_gamma);
3206 let s30 = builder.ins().ushr_imm(x1, 30);
3207 let x2 = builder.ins().bxor(x1, s30);
3208 let c_m1 = builder
3209 .ins()
3210 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3211 let x3 = builder.ins().imul(x2, c_m1);
3212 let s27 = builder.ins().ushr_imm(x3, 27);
3213 let x4 = builder.ins().bxor(x3, s27);
3214 let c_m2 = builder
3215 .ins()
3216 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3217 let x5 = builder.ins().imul(x4, c_m2);
3218 let s31 = builder.ins().ushr_imm(x5, 31);
3219 let h = builder.ins().bxor(x5, s31);
3220
3221 let fval = builder.ins().fcvt_from_uint(types::F64, h);
3222 let max_f = builder.ins().f64const(u64::MAX as f64);
3223 let unit = builder.ins().fdiv(fval, max_f);
3224 let p_f = builder.ins().f64const(f64::from_bits(*p_bits));
3225 let cmp = builder
3226 .ins()
3227 .fcmp(ir::condcodes::FloatCC::LessThan, unit, p_f);
3228 let zero_bits = builder.ins().iconst(types::I64, 0.0_f64.to_bits() as i64);
3229 let one_bits = builder.ins().iconst(types::I64, 1.0_f64.to_bits() as i64);
3230 let result = builder.ins().select(cmp, one_bits, zero_bits);
3231 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3232 }
3233 JitOp::Popcnt => {
3234 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3235 let result = builder.ins().popcnt(val);
3236 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3237 }
3238 JitOp::Clz => {
3239 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3240 let result = builder.ins().clz(val);
3241 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3242 }
3243 JitOp::Ctz => {
3244 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3245 let result = builder.ins().ctz(val);
3246 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3247 }
3248 JitOp::Bswap => {
3249 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3250 let result = builder.ins().bswap(val);
3251 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3252 }
3253 JitOp::ShuffleConst(feedback, size, min) => {
3254 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3255 let fb = builder.ins().iconst(types::I64, *feedback as i64);
3256 let sz = builder.ins().iconst(types::I64, *size as i64);
3257 let mn = builder.ins().iconst(types::I64, *min as i64);
3258 let call = builder.ins().call(shuffle_func_ref, &[val, fb, sz, mn]);
3259 let result = builder.inst_results(call)[0];
3260 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3261 }
3262
3263 JitOp::UnitInterval => {
3265 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3267 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3268 let max_f = builder.ins().f64const(u64::MAX as f64);
3269 let result = builder.ins().fdiv(fval, max_f);
3270 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3271 }
3272 JitOp::F64ToU64 => {
3273 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3274 let result = builder.ins().fcvt_to_uint_sat(types::I64, fval);
3275 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3276 }
3277 JitOp::RoundToU64 => {
3278 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3279 let rounded = round_half_away(&mut builder, fval);
3280 let result = builder.ins().fcvt_to_uint_sat(types::I64, rounded);
3281 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3282 }
3283 JitOp::FloorToU64 => {
3284 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3285 let floored = builder.ins().floor(fval);
3286 let result = builder.ins().fcvt_to_uint_sat(types::I64, floored);
3287 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3288 }
3289 JitOp::CeilToU64 => {
3290 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3291 let ceiled = builder.ins().ceil(fval);
3292 let result = builder.ins().fcvt_to_uint_sat(types::I64, ceiled);
3293 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3294 }
3295 JitOp::ClampF64Const(min_bits, max_bits) => {
3296 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3297 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3298 let fmax = builder.ins().f64const(f64::from_bits(*max_bits));
3299 let clamped = clamp_ir(&mut builder, fval, fmin, fmax);
3300 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], clamped);
3301 }
3302 JitOp::LerpConst(a_bits, b_bits) => {
3303 let t = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3305 let a = builder.ins().f64const(f64::from_bits(*a_bits));
3306 let b = builder.ins().f64const(f64::from_bits(*b_bits));
3307 let diff = builder.ins().fsub(b, a);
3308 let scaled = builder.ins().fmul(t, diff);
3309 let result = builder.ins().fadd(a, scaled);
3310 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3311 }
3312 JitOp::ScaleRangeConst(min_bits, range_bits) => {
3313 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3315 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3316 let max_f = builder.ins().f64const(u64::MAX as f64);
3317 let t = builder.ins().fdiv(fval, max_f);
3318 let fmin = builder.ins().f64const(f64::from_bits(*min_bits));
3319 let frange = builder.ins().f64const(f64::from_bits(*range_bits));
3320 let scaled = builder.ins().fmul(t, frange);
3321 let result = builder.ins().fadd(fmin, scaled);
3322 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3323 }
3324 JitOp::QuantizeConst(step_bits) => {
3325 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3327 let step = builder.ins().f64const(f64::from_bits(*step_bits));
3328 let divided = builder.ins().fdiv(fval, step);
3329 let rounded = round_half_away(&mut builder, divided);
3330 let result = builder.ins().fmul(rounded, step);
3331 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3332 }
3333
3334 JitOp::LutSampleConst(lut_ptr, lut_len) => {
3335 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3337 let ptr_val = builder.ins().iconst(types::I64, *lut_ptr as i64);
3338 let len_val = builder.ins().iconst(types::I64, *lut_len as i64);
3339 let call = builder
3340 .ins()
3341 .call(lut_sample_func_ref, &[input, ptr_val, len_val]);
3342 let result = builder.inst_results(call)[0];
3343 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3344 }
3345 JitOp::DiscretizeConst(range_bits, buckets) => {
3346 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3348 let range = f64::from_bits(*range_bits);
3349 let fzero = builder.ins().f64const(0.0);
3350 let frange_m_eps = builder.ins().f64const(range - f64::EPSILON);
3351 let frange = builder.ins().f64const(range);
3352 let fbuckets = builder.ins().f64const(*buckets as f64);
3353 let clamped = clamp_ir(&mut builder, fval, fzero, frange_m_eps);
3354 let divided = builder.ins().fdiv(clamped, frange);
3355 let scaled = builder.ins().fmul(divided, fbuckets);
3356 let as_u64 = builder.ins().fcvt_to_uint_sat(types::I64, scaled);
3357 let max_bucket = builder.ins().iconst(types::I64, (*buckets - 1) as i64);
3358 let result = builder.ins().umin(as_u64, max_bucket);
3359 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3360 }
3361
3362 JitOp::WeightedPickConst(values_ptr, biases_ptr, primaries_ptr, aliases_ptr, n) => {
3363 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3365 let v_ptr = builder.ins().iconst(types::I64, *values_ptr as i64);
3366 let b_ptr = builder.ins().iconst(types::I64, *biases_ptr as i64);
3367 let p_ptr = builder.ins().iconst(types::I64, *primaries_ptr as i64);
3368 let a_ptr = builder.ins().iconst(types::I64, *aliases_ptr as i64);
3369 let n_val = builder.ins().iconst(types::I64, *n as i64);
3370 let call = builder.ins().call(
3371 weighted_pick_func_ref,
3372 &[input, v_ptr, b_ptr, p_ptr, a_ptr, n_val],
3373 );
3374 let result = builder.inst_results(call)[0];
3375 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3376 }
3377
3378 JitOp::MathUnary(idx) => {
3379 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3380 let func_ref = math_unary_refs[*idx as usize];
3381 let call = builder.ins().call(func_ref, &[input]);
3382 let result = builder.inst_results(call)[0];
3383 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3384 }
3385
3386 JitOp::MathBinary(idx) => {
3387 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3388 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3389 let func_ref = math_binary_refs[*idx as usize];
3390 let call = builder.ins().call(func_ref, &[a, b]);
3391 let result = builder.inst_results(call)[0];
3392 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3393 }
3394
3395 JitOp::ToF64 => {
3396 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3397 let fval = builder.ins().fcvt_from_uint(types::F64, val);
3398 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3399 }
3400
3401 JitOp::RegBinOp(lane, arith) => {
3403 let vt = reg_lane_type(*lane);
3404 let a = load_reg128(&mut builder, buffer_ptr, input_slots[0], vt);
3405 let b = load_reg128(&mut builder, buffer_ptr, input_slots[2], vt);
3406 let is_float = matches!(*lane, 4 | 5);
3407 let r = match (arith, is_float) {
3408 (0, false) => builder.ins().iadd(a, b),
3409 (1, false) => builder.ins().isub(a, b),
3410 (2, false) => builder.ins().imul(a, b),
3411 (0, true) => builder.ins().fadd(a, b),
3412 (1, true) => builder.ins().fsub(a, b),
3413 (2, true) => builder.ins().fmul(a, b),
3414 _ => unreachable!("RegBinOp arith index out of range"),
3415 };
3416 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
3417 }
3418 JitOp::RegCopy => {
3419 let v = load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I64X2);
3420 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3421 }
3422 JitOp::RegSplat(lane) => {
3423 let vt = reg_lane_type(*lane);
3424 let scalar = match *lane {
3425 0 => {
3427 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3428 builder.ins().ireduce(types::I8, v)
3429 }
3430 1 => {
3431 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3432 builder.ins().ireduce(types::I16, v)
3433 }
3434 2 => {
3435 let v = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3436 builder.ins().ireduce(types::I32, v)
3437 }
3438 3 => load_slot(&mut builder, buffer_ptr, input_slots[0]),
3439 4 => {
3441 let f = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3442 builder.ins().fdemote(types::F32, f)
3443 }
3444 5 => load_slot_f64(&mut builder, buffer_ptr, input_slots[0]),
3445 _ => unreachable!("RegSplat lane index out of range"),
3446 };
3447 let v = builder.ins().splat(vt, scalar);
3448 store_reg128(&mut builder, buffer_ptr, output_slots[0], v);
3449 }
3450
3451 JitOp::U64Add2 => {
3453 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3454 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3455 let result = builder.ins().iadd(a, b);
3456 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3457 }
3458 JitOp::U64Sub2 => {
3459 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3460 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3461 let result = builder.ins().isub(a, b);
3462 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3463 }
3464 JitOp::U64Mul2 => {
3465 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3466 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3467 let result = builder.ins().imul(a, b);
3468 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3469 }
3470 JitOp::U64Div2 => {
3471 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3472 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3473 let zero = builder.ins().iconst(types::I64, 0);
3476 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3477 let div_block = builder.create_block();
3478 let merge_block = builder.create_block();
3479 builder.append_block_param(merge_block, types::I64);
3480 builder
3481 .ins()
3482 .brif(is_zero, merge_block, &[zero], div_block, &[]);
3483 builder.switch_to_block(div_block);
3484 builder.seal_block(div_block);
3485 let div_result = builder.ins().udiv(a, b);
3486 builder.ins().jump(merge_block, &[div_result]);
3487 builder.switch_to_block(merge_block);
3488 builder.seal_block(merge_block);
3489 let result = builder.block_params(merge_block)[0];
3490 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3491 }
3492 JitOp::U64Mod2 => {
3493 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3494 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3495 let zero = builder.ins().iconst(types::I64, 0);
3498 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, b, zero);
3499 let rem_block = builder.create_block();
3500 let merge_block = builder.create_block();
3501 builder.append_block_param(merge_block, types::I64);
3502 builder
3503 .ins()
3504 .brif(is_zero, merge_block, &[zero], rem_block, &[]);
3505 builder.switch_to_block(rem_block);
3506 builder.seal_block(rem_block);
3507 let rem_result = builder.ins().urem(a, b);
3508 builder.ins().jump(merge_block, &[rem_result]);
3509 builder.switch_to_block(merge_block);
3510 builder.seal_block(merge_block);
3511 let result = builder.block_params(merge_block)[0];
3512 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3513 }
3514 JitOp::U64And => {
3515 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3516 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3517 let result = builder.ins().band(a, b);
3518 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3519 }
3520 JitOp::U64Or => {
3521 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3522 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3523 let result = builder.ins().bor(a, b);
3524 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3525 }
3526 JitOp::U64Xor => {
3527 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3528 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3529 let result = builder.ins().bxor(a, b);
3530 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3531 }
3532 JitOp::U64Shl => {
3533 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3534 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3535 let result = builder.ins().ishl(a, b);
3536 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3537 }
3538 JitOp::U64Shr => {
3539 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3540 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3541 let result = builder.ins().ushr(a, b);
3542 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3543 }
3544 JitOp::U64Not => {
3545 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3546 let result = builder.ins().bnot(a);
3547 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3548 }
3549
3550 JitOp::F64Add => {
3552 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3553 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3554 let result = builder.ins().fadd(a, b);
3555 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3556 }
3557 JitOp::F64Sub => {
3558 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3559 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3560 let result = builder.ins().fsub(a, b);
3561 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3562 }
3563 JitOp::F64Mul => {
3564 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3565 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3566 let result = builder.ins().fmul(a, b);
3567 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3568 }
3569 JitOp::F64Div => {
3570 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3571 let b = load_slot_f64(&mut builder, buffer_ptr, input_slots[1]);
3572 let zero = builder.ins().f64const(0.0);
3574 let is_zero = builder.ins().fcmp(ir::condcodes::FloatCC::Equal, b, zero);
3575 let div_result = builder.ins().fdiv(a, b);
3576 let result = builder.ins().select(is_zero, zero, div_result);
3577 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3578 }
3579 JitOp::F64Mod => {
3580 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3582 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3583 let call = builder
3584 .ins()
3585 .call(math_binary_refs[F64_MOD_HELPER], &[a, b]);
3586 let result = builder.inst_results(call)[0];
3587 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3588 }
3589
3590 JitOp::IsPositiveCheck { name_ptr, name_len } => {
3591 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3597 let zero = builder.ins().iconst(types::I64, 0);
3598 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, val, zero);
3599 let fail_block = builder.create_block();
3600 let ok_block = builder.create_block();
3601 builder.ins().brif(is_zero, fail_block, &[], ok_block, &[]);
3602
3603 builder.switch_to_block(fail_block);
3604 builder.seal_block(fail_block);
3605 let np = builder.ins().iconst(types::I64, *name_ptr as i64);
3606 let nl = builder.ins().iconst(types::I64, *name_len as i64);
3607 let _ = builder.ins().call(is_positive_fail_ref, &[val, np, nl]);
3608 builder.ins().jump(ok_block, &[]);
3612
3613 builder.switch_to_block(ok_block);
3614 builder.seal_block(ok_block);
3615 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3616 }
3617
3618 JitOp::InRangeCheck(lo, hi) => {
3619 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3623 let lo_v = builder.ins().iconst(types::I64, *lo as i64);
3624 let hi_v = builder.ins().iconst(types::I64, *hi as i64);
3625 let below =
3626 builder
3627 .ins()
3628 .icmp(ir::condcodes::IntCC::UnsignedLessThan, val, lo_v);
3629 let above =
3630 builder
3631 .ins()
3632 .icmp(ir::condcodes::IntCC::UnsignedGreaterThan, val, hi_v);
3633 let out_of_range = builder.ins().bor(below, above);
3634
3635 let fail_block = builder.create_block();
3636 let ok_block = builder.create_block();
3637 builder
3638 .ins()
3639 .brif(out_of_range, fail_block, &[], ok_block, &[]);
3640
3641 builder.switch_to_block(fail_block);
3642 builder.seal_block(fail_block);
3643 let _ = builder.ins().call(in_range_fail_ref, &[val, lo_v, hi_v]);
3644 builder.ins().jump(ok_block, &[]);
3645
3646 builder.switch_to_block(ok_block);
3647 builder.seal_block(ok_block);
3648 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3649 }
3650
3651 JitOp::IsOneOfCheck {
3652 allowed,
3653 set_ptr,
3654 set_len,
3655 } => {
3656 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3662 let mut any_match = builder.ins().iconst(types::I8, 0);
3663 for allow in allowed.iter() {
3664 let c = builder.ins().iconst(types::I64, *allow as i64);
3665 let eq = builder.ins().icmp(ir::condcodes::IntCC::Equal, val, c);
3666 any_match = builder.ins().bor(any_match, eq);
3667 }
3668 let fail_block = builder.create_block();
3669 let ok_block = builder.create_block();
3670 builder
3674 .ins()
3675 .brif(any_match, ok_block, &[], fail_block, &[]);
3676
3677 builder.switch_to_block(fail_block);
3678 builder.seal_block(fail_block);
3679 let sp = builder.ins().iconst(types::I64, *set_ptr as i64);
3680 let sl = builder.ins().iconst(types::I64, *set_len as i64);
3681 let _ = builder.ins().call(is_one_of_fail_ref, &[val, sp, sl]);
3682 builder.ins().jump(ok_block, &[]);
3683
3684 builder.switch_to_block(ok_block);
3685 builder.seal_block(ok_block);
3686 store_slot(&mut builder, buffer_ptr, output_slots[0], val);
3687 }
3688
3689 JitOp::U64Cmp(cc) => {
3690 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3691 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3692 let cmp = builder.ins().icmp(*cc, a, b);
3693 let zero = builder.ins().iconst(types::I64, 0);
3694 let one = builder.ins().iconst(types::I64, 1);
3695 let result = builder.ins().select(cmp, one, zero);
3696 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3697 }
3698 JitOp::F64Cmp(cc) => {
3699 let a = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3700 let b = load_slot_f64(
3701 &mut builder,
3702 buffer_ptr,
3703 if input_slots.len() > 1 {
3704 input_slots[1]
3705 } else {
3706 input_slots[0]
3707 },
3708 );
3709 let cmp = builder.ins().fcmp(*cc, a, b);
3710 let zero = builder.ins().iconst(types::I64, 0);
3711 let one = builder.ins().iconst(types::I64, 1);
3712 let result = builder.ins().select(cmp, one, zero);
3713 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3714 }
3715 JitOp::SelectU64 => {
3716 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3717 let a = load_slot(
3718 &mut builder,
3719 buffer_ptr,
3720 if input_slots.len() > 1 {
3721 input_slots[1]
3722 } else {
3723 input_slots[0]
3724 },
3725 );
3726 let b = load_slot(
3727 &mut builder,
3728 buffer_ptr,
3729 if input_slots.len() > 2 {
3730 input_slots[2]
3731 } else {
3732 input_slots[0]
3733 },
3734 );
3735 let zero = builder.ins().iconst(types::I64, 0);
3736 let is_nonzero = builder
3737 .ins()
3738 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3739 let result = builder.ins().select(is_nonzero, a, b);
3740 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3741 }
3742 JitOp::SelectF64 => {
3743 let cond = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3744 let a = load_slot_f64(
3745 &mut builder,
3746 buffer_ptr,
3747 if input_slots.len() > 1 {
3748 input_slots[1]
3749 } else {
3750 input_slots[0]
3751 },
3752 );
3753 let b = load_slot_f64(
3754 &mut builder,
3755 buffer_ptr,
3756 if input_slots.len() > 2 {
3757 input_slots[2]
3758 } else {
3759 input_slots[0]
3760 },
3761 );
3762 let zero = builder.ins().iconst(types::I64, 0);
3763 let is_nonzero = builder
3764 .ins()
3765 .icmp(ir::condcodes::IntCC::NotEqual, cond, zero);
3766 let result = builder.ins().select(is_nonzero, a, b);
3767 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], result);
3768 }
3769
3770 JitOp::I64ToF64 => {
3771 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3772 let fval = builder.ins().fcvt_from_sint(types::F64, val);
3773 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], fval);
3774 }
3775 JitOp::F64ToI64 => {
3776 let fval = load_slot_f64(&mut builder, buffer_ptr, input_slots[0]);
3777 let ival = builder.ins().fcvt_to_sint(types::I64, fval);
3778 store_slot(&mut builder, buffer_ptr, output_slots[0], ival);
3779 }
3780 JitOp::SignExtendI32 => {
3781 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3782 let i32_val = builder.ins().ireduce(types::I32, val);
3783 let sext_val = builder.ins().sextend(types::I64, i32_val);
3784 store_slot(&mut builder, buffer_ptr, output_slots[0], sext_val);
3785 }
3786 JitOp::SignExtendI16 => {
3787 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3788 let i16_val = builder.ins().ireduce(types::I16, val);
3789 let sext_val = builder.ins().sextend(types::I64, i16_val);
3790 store_slot(&mut builder, buffer_ptr, output_slots[0], sext_val);
3791 }
3792 JitOp::SignExtendI8 => {
3793 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3794 let i8_val = builder.ins().ireduce(types::I8, val);
3795 let sext_val = builder.ins().sextend(types::I64, i8_val);
3796 store_slot(&mut builder, buffer_ptr, output_slots[0], sext_val);
3797 }
3798 JitOp::ZeroExtendU32 => {
3799 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3800 let mask = builder.ins().iconst(types::I64, 0xFFFFFFFFu64 as i64);
3801 let result = builder.ins().band(val, mask);
3802 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3803 }
3804 JitOp::ZeroExtendU16 => {
3805 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3806 let mask = builder.ins().iconst(types::I64, 0xFFFFu64 as i64);
3807 let result = builder.ins().band(val, mask);
3808 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3809 }
3810 JitOp::ZeroExtendU8 => {
3811 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3812 let mask = builder.ins().iconst(types::I64, 0xFFu64 as i64);
3813 let result = builder.ins().band(val, mask);
3814 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3815 }
3816 JitOp::ToBool => {
3817 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3818 let zero = builder.ins().iconst(types::I64, 0);
3819 let one = builder.ins().iconst(types::I64, 1);
3820 let cmp = builder
3821 .ins()
3822 .icmp(ir::condcodes::IntCC::NotEqual, val, zero);
3823 let result = builder.ins().select(cmp, one, zero);
3824 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3825 }
3826 JitOp::ConstU64(v) | JitOp::ConstF64(v) => {
3827 let result = builder.ins().iconst(types::I64, *v as i64);
3828 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
3829 }
3830 JitOp::HashRangeConst(max) => {
3831 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3832 let c_gamma = builder
3833 .ins()
3834 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3835 let x1 = builder.ins().iadd(input, c_gamma);
3836 let s30 = builder.ins().ushr_imm(x1, 30);
3837 let x2 = builder.ins().bxor(x1, s30);
3838 let c_m1 = builder
3839 .ins()
3840 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3841 let x3 = builder.ins().imul(x2, c_m1);
3842 let s27 = builder.ins().ushr_imm(x3, 27);
3843 let x4 = builder.ins().bxor(x3, s27);
3844 let c_m2 = builder
3845 .ins()
3846 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3847 let x5 = builder.ins().imul(x4, c_m2);
3848 let s31 = builder.ins().ushr_imm(x5, 31);
3849 let h = builder.ins().bxor(x5, s31);
3850 if *max == 0 {
3851 let zero = builder.ins().iconst(types::I64, 0);
3852 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
3853 } else {
3854 let m = builder.ins().iconst(types::I64, *max as i64);
3855 let rem = builder.ins().urem(h, m);
3856 store_slot(&mut builder, buffer_ptr, output_slots[0], rem);
3857 }
3858 }
3859 JitOp::HashIntervalConst(min_bits, max_bits) => {
3860 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3861 let c_gamma = builder
3862 .ins()
3863 .iconst(types::I64, 0x9e3779b97f4a7c15u64 as i64);
3864 let x1 = builder.ins().iadd(input, c_gamma);
3865 let s30 = builder.ins().ushr_imm(x1, 30);
3866 let x2 = builder.ins().bxor(x1, s30);
3867 let c_m1 = builder
3868 .ins()
3869 .iconst(types::I64, 0xbf58476d1ce4e5b9u64 as i64);
3870 let x3 = builder.ins().imul(x2, c_m1);
3871 let s27 = builder.ins().ushr_imm(x3, 27);
3872 let x4 = builder.ins().bxor(x3, s27);
3873 let c_m2 = builder
3874 .ins()
3875 .iconst(types::I64, 0x94d049bb133111ebu64 as i64);
3876 let x5 = builder.ins().imul(x4, c_m2);
3877 let s31 = builder.ins().ushr_imm(x5, 31);
3878 let h = builder.ins().bxor(x5, s31);
3879
3880 let h_f = builder.ins().fcvt_from_uint(types::F64, h);
3881 let denom = builder.ins().f64const(u64::MAX as f64);
3882 let unit = builder.ins().fdiv(h_f, denom);
3883 let min_f = f64::from_bits(*min_bits);
3884 let max_f = f64::from_bits(*max_bits);
3885 let span = builder.ins().f64const(max_f - min_f);
3886 let min_val = builder.ins().f64const(min_f);
3887 let scaled = builder.ins().fmul(unit, span);
3888 let res_f = builder.ins().fadd(min_val, scaled);
3889 let res = builder
3890 .ins()
3891 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3892 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3893 }
3894 JitOp::InvLerpConst(a_bits, b_bits) => {
3895 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3896 let in_f = builder
3897 .ins()
3898 .bitcast(types::F64, ir::MemFlags::new(), input);
3899 let a_f = f64::from_bits(*a_bits);
3900 let b_f = f64::from_bits(*b_bits);
3901 let a_val = builder.ins().f64const(a_f);
3902 let inv_span = builder.ins().f64const(1.0 / (b_f - a_f));
3906 let diff = builder.ins().fsub(in_f, a_val);
3907 let t = builder.ins().fmul(diff, inv_span);
3908 let zero = builder.ins().f64const(0.0);
3909 let one = builder.ins().f64const(1.0);
3910 let res_f = clamp_ir(&mut builder, t, zero, one);
3911 let res = builder
3912 .ins()
3913 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3914 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3915 }
3916 JitOp::RemapConst(in_min_bits, in_max_bits, out_min_bits, out_max_bits) => {
3917 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3918 let in_f = builder
3919 .ins()
3920 .bitcast(types::F64, ir::MemFlags::new(), input);
3921 let in_min = f64::from_bits(*in_min_bits);
3922 let in_max = f64::from_bits(*in_max_bits);
3923 let out_min = f64::from_bits(*out_min_bits);
3924 let out_max = f64::from_bits(*out_max_bits);
3925 let in_span_val = builder.ins().f64const(in_max - in_min);
3929 let in_min_val = builder.ins().f64const(in_min);
3930 let out_min_val = builder.ins().f64const(out_min);
3931 let out_span_val = builder.ins().f64const(out_max - out_min);
3932 let diff = builder.ins().fsub(in_f, in_min_val);
3933 let t = builder.ins().fdiv(diff, in_span_val);
3934 let scaled = builder.ins().fmul(t, out_span_val);
3935 let res_f = builder.ins().fadd(out_min_val, scaled);
3936 let res = builder
3937 .ins()
3938 .bitcast(types::I64, ir::MemFlags::new(), res_f);
3939 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3940 }
3941 JitOp::EpochOffsetConst(base) => {
3942 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3943 let b = builder.ins().iconst(types::I64, *base as i64);
3944 let res = builder.ins().iadd(val, b);
3945 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3946 }
3947 JitOp::EpochScaleConst(factor) => {
3948 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3949 let f = builder.ins().iconst(types::I64, *factor as i64);
3950 let res = builder.ins().imul(val, f);
3951 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3952 }
3953 JitOp::ThreadId => {
3954 let call = builder.ins().call(thread_id_func_ref, &[]);
3955 let res = builder.inst_results(call)[0];
3956 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3957 }
3958 JitOp::CurrentEpochMillis => {
3959 let call = builder.ins().call(current_epoch_millis_func_ref, &[]);
3960 let res = builder.inst_results(call)[0];
3961 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3962 }
3963 JitOp::Perlin1dConst(perm_ptr, freq_bits) => {
3964 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3965 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
3966 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
3967 let call = builder.ins().call(perlin_1d_func_ref, &[input, p, fb]);
3968 let res = builder.inst_results(call)[0];
3969 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3970 }
3971 JitOp::Perlin2dConst(perm_ptr, freq_bits) => {
3972 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3973 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3974 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
3975 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
3976 let call = builder.ins().call(perlin_2d_func_ref, &[x, y, p, fb]);
3977 let res = builder.inst_results(call)[0];
3978 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3979 }
3980 JitOp::Simplex2dConst(perm_ptr, freq_bits) => {
3981 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3982 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
3983 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
3984 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
3985 let call = builder.ins().call(simplex_2d_func_ref, &[x, y, p, fb]);
3986 let res = builder.inst_results(call)[0];
3987 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3988 }
3989 JitOp::FractalNoise1dConst(perm_ptr, freq_bits, octaves) => {
3990 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
3991 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
3992 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
3993 let oct = builder.ins().iconst(types::I64, *octaves as i64);
3994 let call = builder
3995 .ins()
3996 .call(fractal_noise_1d_func_ref, &[input, p, fb, oct]);
3997 let res = builder.inst_results(call)[0];
3998 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
3999 }
4000 JitOp::FractalNoise2dConst(perm_ptr, freq_bits, octaves) => {
4001 let x = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4002 let y = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4003 let p = builder.ins().iconst(types::I64, *perm_ptr as i64);
4004 let fb = builder.ins().iconst(types::I64, *freq_bits as i64);
4005 let oct = builder.ins().iconst(types::I64, *octaves as i64);
4006 let call = builder
4007 .ins()
4008 .call(fractal_noise_2d_func_ref, &[x, y, p, fb, oct]);
4009 let res = builder.inst_results(call)[0];
4010 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4011 }
4012 JitOp::CycleWalkConst(range, seed, inc) => {
4013 let pos = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4014 let r = builder.ins().iconst(types::I64, *range as i64);
4015 let s = builder.ins().iconst(types::I64, *seed as i64);
4016 let i = builder.ins().iconst(types::I64, *inc as i64);
4017 let call = builder.ins().call(cycle_walk_func_ref, &[pos, r, s, i]);
4018 let res = builder.inst_results(call)[0];
4019 store_slot(&mut builder, buffer_ptr, output_slots[0], res);
4020 }
4021
4022 JitOp::VariadicSum => {
4023 if input_slots.is_empty() {
4024 let zero = builder.ins().iconst(types::I64, 0);
4025 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4026 } else {
4027 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4028 for &slot in &input_slots[1..] {
4029 let v = load_slot(&mut builder, buffer_ptr, slot);
4030 acc = builder.ins().iadd(acc, v);
4031 }
4032 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4033 }
4034 }
4035 JitOp::VariadicProduct => {
4036 if input_slots.is_empty() {
4037 let one = builder.ins().iconst(types::I64, 1);
4038 store_slot(&mut builder, buffer_ptr, output_slots[0], one);
4039 } else {
4040 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4041 for &slot in &input_slots[1..] {
4042 let v = load_slot(&mut builder, buffer_ptr, slot);
4043 acc = builder.ins().imul(acc, v);
4044 }
4045 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4046 }
4047 }
4048 JitOp::VariadicMin => {
4049 if input_slots.is_empty() {
4050 let zero = builder.ins().iconst(types::I64, 0);
4051 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4052 } else {
4053 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4054 for &slot in &input_slots[1..] {
4055 let v = load_slot(&mut builder, buffer_ptr, slot);
4056 let cmp =
4057 builder
4058 .ins()
4059 .icmp(ir::condcodes::IntCC::UnsignedLessThan, v, acc);
4060 acc = builder.ins().select(cmp, v, acc);
4061 }
4062 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4063 }
4064 }
4065 JitOp::VariadicMax => {
4066 if input_slots.is_empty() {
4067 let zero = builder.ins().iconst(types::I64, 0);
4068 store_slot(&mut builder, buffer_ptr, output_slots[0], zero);
4069 } else {
4070 let mut acc = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4071 for &slot in &input_slots[1..] {
4072 let v = load_slot(&mut builder, buffer_ptr, slot);
4073 let cmp = builder.ins().icmp(
4074 ir::condcodes::IntCC::UnsignedGreaterThan,
4075 v,
4076 acc,
4077 );
4078 acc = builder.ins().select(cmp, v, acc);
4079 }
4080 store_slot(&mut builder, buffer_ptr, output_slots[0], acc);
4081 }
4082 }
4083
4084 JitOp::CeilToMultiple => {
4085 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4086 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4087 let zero = builder.ins().iconst(types::I64, 0);
4088 let one = builder.ins().iconst(types::I64, 1);
4089 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4090 let calc_block = builder.create_block();
4091 let merge_block = builder.create_block();
4092 builder.append_block_param(merge_block, types::I64);
4093 builder
4094 .ins()
4095 .brif(is_zero, merge_block, &[val], calc_block, &[]);
4096 builder.switch_to_block(calc_block);
4097 builder.seal_block(calc_block);
4098 let div = div_ceil(&mut builder, val, m, one);
4101 let high = builder.ins().umulhi(div, m);
4102 let low = builder.ins().imul(div, m);
4103 let zero_hi = builder.ins().iconst(types::I64, 0);
4104 let overflows =
4105 builder
4106 .ins()
4107 .icmp(ir::condcodes::IntCC::NotEqual, high, zero_hi);
4108 let max = builder.ins().iconst(types::I64, -1);
4109 let mul = builder.ins().select(overflows, max, low);
4110 builder.ins().jump(merge_block, &[mul]);
4111 builder.switch_to_block(merge_block);
4112 builder.seal_block(merge_block);
4113 let result = builder.block_params(merge_block)[0];
4114 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4115 }
4116 JitOp::CheckedAdd => {
4117 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4118 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4119 let sum = builder.ins().iadd(a, b);
4120 let is_overflow =
4121 builder
4122 .ins()
4123 .icmp(ir::condcodes::IntCC::UnsignedLessThan, sum, a);
4124 let zero = builder.ins().iconst(types::I64, 0);
4125 let result = builder.ins().select(is_overflow, zero, sum);
4126 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4127 }
4128 JitOp::CheckedSub => {
4129 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4130 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4131 let is_lt = builder
4132 .ins()
4133 .icmp(ir::condcodes::IntCC::UnsignedLessThan, a, b);
4134 let diff = builder.ins().isub(a, b);
4135 let zero = builder.ins().iconst(types::I64, 0);
4136 let result = builder.ins().select(is_lt, zero, diff);
4137 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4138 }
4139 JitOp::CheckedMul => {
4140 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4141 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4142 let prod = builder.ins().imul(a, b);
4143 let zero = builder.ins().iconst(types::I64, 0);
4144 let a_is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, a, zero);
4145 let div_block = builder.create_block();
4146 let merge_block = builder.create_block();
4147 builder.append_block_param(merge_block, types::I64);
4148 builder
4149 .ins()
4150 .brif(a_is_zero, merge_block, &[zero], div_block, &[]);
4151 builder.switch_to_block(div_block);
4152 builder.seal_block(div_block);
4153 let div = builder.ins().udiv(prod, a);
4154 let ok = builder.ins().icmp(ir::condcodes::IntCC::Equal, div, b);
4155 let mul_res = builder.ins().select(ok, prod, zero);
4156 builder.ins().jump(merge_block, &[mul_res]);
4157 builder.switch_to_block(merge_block);
4158 builder.seal_block(merge_block);
4159 let result = builder.block_params(merge_block)[0];
4160 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4161 }
4162 JitOp::MultiplesAtLeast => {
4163 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4164 let m = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4165 let zero = builder.ins().iconst(types::I64, 0);
4166 let one = builder.ins().iconst(types::I64, 1);
4167 let is_zero = builder.ins().icmp(ir::condcodes::IntCC::Equal, m, zero);
4168 let calc_block = builder.create_block();
4169 let merge_block = builder.create_block();
4170 builder.append_block_param(merge_block, types::I64);
4171 builder
4172 .ins()
4173 .brif(is_zero, merge_block, &[zero], calc_block, &[]);
4174 builder.switch_to_block(calc_block);
4175 builder.seal_block(calc_block);
4176 let div = div_ceil(&mut builder, val, m, one);
4177 builder.ins().jump(merge_block, &[div]);
4178 builder.switch_to_block(merge_block);
4179 builder.seal_block(merge_block);
4180 let result = builder.block_params(merge_block)[0];
4181 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4182 }
4183
4184 JitOp::BlendConst(mix_bits) => {
4185 let a = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4190 let b = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4191 let fa = builder.ins().bitcast(types::F64, ir::MemFlags::new(), a);
4192 let fb = builder.ins().bitcast(types::F64, ir::MemFlags::new(), b);
4193 let mix_f64 = f64::from_bits(*mix_bits);
4194 let mix_val = builder.ins().f64const(mix_f64);
4195 let one = builder.ins().f64const(1.0);
4196 let one_minus_mix = builder.ins().fsub(one, mix_val);
4197 let a_part = builder.ins().fmul(fa, one_minus_mix);
4198 let b_part = builder.ins().fmul(fb, mix_val);
4199 let sum = builder.ins().fadd(a_part, b_part);
4200 let result = builder.ins().bitcast(types::I64, ir::MemFlags::new(), sum);
4201 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4202 }
4203 JitOp::LfsrStepConst(feedback) => {
4204 let val = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4205 let feedback = builder.ins().iconst(types::I64, *feedback as i64);
4206 let one = builder.ins().iconst(types::I64, 1);
4207 let zero = builder.ins().iconst(types::I64, 0);
4208 let shifted = builder.ins().ushr(val, one);
4209 let lsb = builder.ins().band(val, one);
4210 let is_odd = builder
4211 .ins()
4212 .icmp(ir::condcodes::IntCC::NotEqual, lsb, zero);
4213 let fb_mask = builder.ins().select(is_odd, feedback, zero);
4214 let result = builder.ins().bxor(shifted, fb_mask);
4215 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4216 }
4217 JitOp::PcgConst(seed, stream) => {
4218 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4219 let s = builder.ins().iconst(types::I64, *seed as i64);
4220 let st = builder.ins().iconst(types::I64, *stream as i64);
4221 let call = builder.ins().call(pcg_func_ref, &[input, s, st]);
4222 let result = builder.inst_results(call)[0];
4223 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4224 }
4225 JitOp::PcgStreamConst(seed) => {
4226 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4227 let st = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4228 let s = builder.ins().iconst(types::I64, *seed as i64);
4229 let call = builder.ins().call(pcg_stream_func_ref, &[input, st, s]);
4230 let result = builder.inst_results(call)[0];
4231 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4232 }
4233 JitOp::NOfConst(n, m) => {
4234 let input = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4235 let n_val = builder.ins().iconst(types::I64, *n as i64);
4236 let m_val = builder.ins().iconst(types::I64, *m as i64);
4237 let call = builder.ins().call(n_of_func_ref, &[input, n_val, m_val]);
4238 let result = builder.inst_results(call)[0];
4239 store_slot(&mut builder, buffer_ptr, output_slots[0], result);
4240 }
4241
4242 JitOp::SlotCall { kit, scratch_base } => {
4243 let n_in = input_slots.len();
4249 let n_out = output_slots.len();
4250 let frame = |builder: &mut FunctionBuilder, n: usize| {
4251 builder.create_sized_stack_slot(ir::StackSlotData::new(
4252 ir::StackSlotKind::ExplicitSlot,
4253 (n.max(1) * 8) as u32,
4254 3,
4255 ))
4256 };
4257 let in_frame = frame(&mut builder, n_in);
4258 let out_frame = frame(&mut builder, n_out);
4259 for (k, &s) in input_slots.iter().enumerate() {
4260 let v = load_slot(&mut builder, buffer_ptr, s);
4261 builder.ins().stack_store(v, in_frame, (k * 8) as i32);
4262 }
4263 let kit_ptr = builder
4264 .ins()
4265 .iconst(types::I64, std::sync::Arc::as_ptr(&kit.0) as usize as i64);
4266 let in_ptr = builder.ins().stack_addr(types::I64, in_frame, 0);
4267 let n_in_v = builder.ins().iconst(types::I64, n_in as i64);
4268 let out_ptr = builder.ins().stack_addr(types::I64, out_frame, 0);
4269 let n_out_v = builder.ins().iconst(types::I64, n_out as i64);
4270 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4271 let n_sc_v = builder.ins().iconst(types::I64, kit.0.scratch.len() as i64);
4272 builder.ins().call(
4273 slot_call_ref,
4274 &[
4275 kit_ptr,
4276 in_ptr,
4277 n_in_v,
4278 out_ptr,
4279 n_out_v,
4280 scratch_ptr,
4281 base_v,
4282 n_sc_v,
4283 ],
4284 );
4285 for (k, &s) in output_slots.iter().enumerate() {
4286 let v = builder
4287 .ins()
4288 .stack_load(types::I64, out_frame, (k * 8) as i32);
4289 store_slot(&mut builder, buffer_ptr, s, v);
4290 }
4291 }
4292
4293 JitOp::U64ToStr { scratch_base }
4294 | JitOp::I64ToStr { scratch_base }
4295 | JitOp::F64ToStr { scratch_base } => {
4296 let func = match jit_op {
4299 JitOp::U64ToStr { .. } => u64_to_str_ref,
4300 JitOp::I64ToStr { .. } => i64_to_str_ref,
4301 _ => f64_to_str_ref,
4302 };
4303 let value = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4304 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4305 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4306 builder
4307 .ins()
4308 .call(func, &[scratch_ptr, base_v, buffer_ptr, out_v, value]);
4309 }
4310 JitOp::JsonToStr { scratch_base } => {
4311 let ptr = load_slot(&mut builder, buffer_ptr, input_slots[0]);
4312 let len = load_slot(&mut builder, buffer_ptr, input_slots[1]);
4313 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4314 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4315 builder.ins().call(
4316 json_to_str_ref,
4317 &[scratch_ptr, base_v, buffer_ptr, out_v, ptr, len],
4318 );
4319 }
4320 JitOp::StrConcat { scratch_base } => {
4321 let n_words = input_slots.len();
4324 let frame = builder.create_sized_stack_slot(ir::StackSlotData::new(
4325 ir::StackSlotKind::ExplicitSlot,
4326 (n_words.max(1) * 8) as u32,
4327 3,
4328 ));
4329 for (k, &s) in input_slots.iter().enumerate() {
4330 let v = load_slot(&mut builder, buffer_ptr, s);
4331 builder.ins().stack_store(v, frame, (k * 8) as i32);
4332 }
4333 let pairs_ptr = builder.ins().stack_addr(types::I64, frame, 0);
4334 let n_v = builder.ins().iconst(types::I64, (n_words / 2) as i64);
4335 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4336 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4337 builder.ins().call(
4338 str_concat_ref,
4339 &[scratch_ptr, base_v, buffer_ptr, out_v, pairs_ptr, n_v],
4340 );
4341 }
4342
4343 JitOp::VecProduce { kind, scratch_base } => {
4344 let func = func_of(&vec_producer_refs, *kind);
4348 let base_v = builder.ins().iconst(types::I64, *scratch_base as i64);
4349 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4350 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4351 let mut args = vec![scratch_ptr, base_v, buffer_ptr, out_v];
4352 args.extend(words);
4353 builder.ins().call(func, &args);
4354 }
4355 JitOp::VecReduce(kind) => {
4356 let func = func_of(&vec_reducer_refs, *kind);
4357 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4358 let call = builder.ins().call(func, &words);
4359 let bits = builder.inst_results(call)[0];
4360 store_slot(&mut builder, buffer_ptr, output_slots[0], bits);
4361 }
4362 JitOp::RegLane(kind) => {
4363 let func = func_of(®_lane_refs, *kind);
4364 let words = load_words(&mut builder, buffer_ptr, input_slots, 3);
4365 let call = builder.ins().call(func, &words);
4366 let word = builder.inst_results(call)[0];
4367 store_slot(&mut builder, buffer_ptr, output_slots[0], word);
4368 }
4369 JitOp::RegProduce(kind) => {
4370 let func = func_of(®_producer_refs, *kind);
4371 let out_v = builder.ins().iconst(types::I64, output_slots[0] as i64);
4372 let words = load_words(&mut builder, buffer_ptr, input_slots, 4);
4373 let mut args = vec![buffer_ptr, out_v];
4374 args.extend(words);
4375 builder.ins().call(func, &args);
4376 }
4377 JitOp::RegDotF32 => {
4378 let a = load_reg128(&mut builder, buffer_ptr, input_slots[0], types::F32X4);
4382 let b = load_reg128(&mut builder, buffer_ptr, input_slots[2], types::F32X4);
4383 let p = builder.ins().fmul(a, b);
4384 let p0 = builder.ins().extractlane(p, 0);
4385 let p1 = builder.ins().extractlane(p, 1);
4386 let p2 = builder.ins().extractlane(p, 2);
4387 let p3 = builder.ins().extractlane(p, 3);
4388 let s01 = builder.ins().fadd(p0, p1);
4389 let s23 = builder.ins().fadd(p2, p3);
4390 let s = builder.ins().fadd(s01, s23);
4391 let wide = builder.ins().fpromote(types::F64, s);
4392 store_slot_f64(&mut builder, buffer_ptr, output_slots[0], wide);
4393 }
4394 JitOp::RegShuffleConst(mask) => {
4395 let x = load_reg128(&mut builder, buffer_ptr, input_slots[0], types::I8X16);
4399 let imm = builder
4400 .func
4401 .dfg
4402 .immediates
4403 .push(ir::ConstantData::from(&mask[..]));
4404 let r = builder.ins().shuffle(x, x, imm);
4405 store_reg128(&mut builder, buffer_ptr, output_slots[0], r);
4406 }
4407
4408 JitOp::Fallback => {
4409 }
4412 }
4413 if let Some((inst, mark)) = tracker_store {
4414 let calls = (mark..builder.func.dfg.num_insts()).any(|i| {
4415 builder.func.dfg.insts[ir::Inst::from_u32(i as u32)]
4416 .opcode()
4417 .is_call()
4418 });
4419 if !calls {
4420 builder.func.layout.remove_inst(inst);
4421 }
4422 }
4423
4424 if let (Some(cp), Some(skip)) = (clean_ptr, skip_block) {
4426 let offset = builder.ins().iconst(types::I64, step_idx as i64);
4427 let addr = builder.ins().iadd(cp, offset);
4428 let one = builder.ins().iconst(types::I8, 1);
4429 builder.ins().store(ir::MemFlags::new(), one, addr, 0);
4430 builder.ins().jump(skip, &[]);
4431 builder.switch_to_block(skip);
4432 builder.seal_block(skip);
4433 }
4434 }
4435
4436 builder.ins().return_(&[]);
4437 builder.finalize();
4438 }
4439 let fallible = ctx.func.layout.blocks().any(|block| {
4442 ctx.func
4443 .layout
4444 .block_insts(block)
4445 .any(|inst| ctx.func.dfg.insts[inst].opcode().is_call())
4446 });
4447
4448 module
4449 .define_function(func_id, &mut ctx)
4450 .map_err(|e| format!("define function: {e}"))?;
4451 module.clear_context(&mut ctx);
4452 module
4453 .finalize_definitions()
4454 .map_err(|e| format!("finalize: {e}"))?;
4455
4456 let code_ptr = module.get_finalized_function(func_id);
4457 let kits: Vec<SlotKitRef> = steps
4459 .iter()
4460 .filter_map(|(op, _, _)| op.slot_kit().cloned())
4461 .collect();
4462 let code = super::kernels::JitCode::new(module, kits, fallible);
4463
4464 if provenance {
4465 let prov_fn: NativeProvFn = unsafe { mem::transmute(code_ptr) };
4466 let dummy_raw: NativeFn = unsafe { mem::transmute(code_ptr) };
4467 Ok((dummy_raw, prov_fn, code))
4468 } else {
4469 let raw_fn: NativeFn = unsafe { mem::transmute(code_ptr) };
4470 let dummy_prov: NativeProvFn = unsafe { mem::transmute(code_ptr) };
4471 Ok((raw_fn, dummy_prov, code))
4472 }
4473}
4474
4475fn load_slot(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4479 let offset = (slot * 8) as i32;
4480 builder
4481 .ins()
4482 .load(types::I64, ir::MemFlags::trusted(), buffer_ptr, offset)
4483}
4484
4485fn store_slot(
4487 builder: &mut FunctionBuilder,
4488 buffer_ptr: ir::Value,
4489 slot: usize,
4490 value: ir::Value,
4491) -> ir::Inst {
4492 let offset = (slot * 8) as i32;
4493 builder
4494 .ins()
4495 .store(ir::MemFlags::trusted(), value, buffer_ptr, offset)
4496}
4497
4498fn reg_lane_type(lane: u8) -> ir::Type {
4501 match lane {
4502 0 => types::I8X16,
4503 1 => types::I16X8,
4504 2 => types::I32X4,
4505 3 => types::I64X2,
4506 4 => types::F32X4,
4507 5 => types::F64X2,
4508 _ => unreachable!("register lane index out of range"),
4509 }
4510}
4511
4512fn load_reg128(
4517 builder: &mut FunctionBuilder,
4518 buffer_ptr: ir::Value,
4519 first_slot: usize,
4520 vt: ir::Type,
4521) -> ir::Value {
4522 let offset = (first_slot * 8) as i32;
4523 builder
4524 .ins()
4525 .load(vt, ir::MemFlags::new(), buffer_ptr, offset)
4526}
4527
4528fn store_reg128(
4530 builder: &mut FunctionBuilder,
4531 buffer_ptr: ir::Value,
4532 first_slot: usize,
4533 value: ir::Value,
4534) {
4535 let offset = (first_slot * 8) as i32;
4536 builder
4537 .ins()
4538 .store(ir::MemFlags::new(), value, buffer_ptr, offset);
4539}
4540
4541fn round_half_away(builder: &mut FunctionBuilder, x: ir::Value) -> ir::Value {
4546 let t = builder.ins().trunc(x);
4547 let frac = builder.ins().fsub(x, t);
4548 let mag = builder.ins().fabs(frac);
4549 let half = builder.ins().f64const(0.5);
4550 let reaches = builder
4551 .ins()
4552 .fcmp(ir::condcodes::FloatCC::GreaterThanOrEqual, mag, half);
4553 let one = builder.ins().f64const(1.0);
4554 let step = builder.ins().fcopysign(one, x);
4555 let up = builder.ins().fadd(t, step);
4556 builder.ins().select(reaches, up, t)
4557}
4558
4559fn clamp_ir(
4564 builder: &mut FunctionBuilder,
4565 x: ir::Value,
4566 lo: ir::Value,
4567 hi: ir::Value,
4568) -> ir::Value {
4569 let below = builder.ins().fcmp(ir::condcodes::FloatCC::LessThan, x, lo);
4570 let above = builder
4571 .ins()
4572 .fcmp(ir::condcodes::FloatCC::GreaterThan, x, hi);
4573 let capped = builder.ins().select(above, hi, x);
4574 builder.ins().select(below, lo, capped)
4575}
4576
4577fn div_ceil(
4580 builder: &mut FunctionBuilder,
4581 val: ir::Value,
4582 m: ir::Value,
4583 one: ir::Value,
4584) -> ir::Value {
4585 let q = builder.ins().udiv(val, m);
4586 let r = builder.ins().urem(val, m);
4587 let zero = builder.ins().iconst(types::I64, 0);
4588 let inexact = builder.ins().icmp(ir::condcodes::IntCC::NotEqual, r, zero);
4589 let q1 = builder.ins().iadd(q, one);
4590 builder.ins().select(inexact, q1, q)
4591}
4592
4593fn func_of<K: PartialEq + Copy>(refs: &[(K, ir::FuncRef)], key: K) -> ir::FuncRef {
4595 refs.iter()
4596 .find(|(k, _)| *k == key)
4597 .map(|(_, r)| *r)
4598 .expect("every helper of the group is declared")
4599}
4600
4601fn load_words(
4604 builder: &mut FunctionBuilder,
4605 buffer_ptr: ir::Value,
4606 input_slots: &[usize],
4607 n: usize,
4608) -> Vec<ir::Value> {
4609 (0..n)
4610 .map(|k| match input_slots.get(k) {
4611 Some(&s) => load_slot(builder, buffer_ptr, s),
4612 None => builder.ins().iconst(types::I64, 0),
4613 })
4614 .collect()
4615}
4616
4617fn load_slot_f64(builder: &mut FunctionBuilder, buffer_ptr: ir::Value, slot: usize) -> ir::Value {
4619 let i64_val = load_slot(builder, buffer_ptr, slot);
4620 builder
4621 .ins()
4622 .bitcast(types::F64, ir::MemFlags::new(), i64_val)
4623}
4624
4625fn store_slot_f64(
4627 builder: &mut FunctionBuilder,
4628 buffer_ptr: ir::Value,
4629 slot: usize,
4630 value: ir::Value,
4631) {
4632 let i64_val = builder
4633 .ins()
4634 .bitcast(types::I64, ir::MemFlags::new(), value);
4635 store_slot(builder, buffer_ptr, slot, i64_val);
4636}
4637
4638#[cfg(test)]
4641mod tests {
4642 use super::*;
4643
4644 #[test]
4645 fn jit_identity() {
4646 let steps = vec![(JitOp::Identity, vec![0], vec![1])];
4647 let mut output_map = HashMap::new();
4648 output_map.insert("out".into(), 1);
4649 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4650 kernel.eval(&[42]);
4651 assert_eq!(kernel.get("out"), 42);
4652 }
4653
4654 #[test]
4655 fn jit_add_const() {
4656 let steps = vec![(JitOp::AddConst(100), vec![0], vec![1])];
4657 let mut output_map = HashMap::new();
4658 output_map.insert("out".into(), 1);
4659 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4660 kernel.eval(&[5]);
4661 assert_eq!(kernel.get("out"), 105);
4662 }
4663
4664 #[test]
4665 fn jit_mul_const() {
4666 let steps = vec![(JitOp::MulConst(7), vec![0], vec![1])];
4667 let mut output_map = HashMap::new();
4668 output_map.insert("out".into(), 1);
4669 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4670 kernel.eval(&[6]);
4671 assert_eq!(kernel.get("out"), 42);
4672 }
4673
4674 #[test]
4675 fn jit_mod_const() {
4676 let steps = vec![(JitOp::ModConst(100), vec![0], vec![1])];
4677 let mut output_map = HashMap::new();
4678 output_map.insert("out".into(), 1);
4679 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4680 kernel.eval(&[542]);
4681 assert_eq!(kernel.get("out"), 42);
4682 }
4683
4684 #[test]
4685 fn jit_hash() {
4686 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
4687 let mut output_map = HashMap::new();
4688 output_map.insert("out".into(), 1);
4689 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4690
4691 kernel.eval(&[42]);
4692 let v1 = kernel.get("out");
4693
4694 let expected = xxhash_rust::xxh3::xxh3_64(&42u64.to_le_bytes());
4696 assert_eq!(v1, expected);
4697 }
4698
4699 #[test]
4700 fn jit_hash_deterministic() {
4701 let steps = vec![(JitOp::Hash, vec![0], vec![1])];
4702 let mut output_map = HashMap::new();
4703 output_map.insert("out".into(), 1);
4704 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4705
4706 kernel.eval(&[42]);
4707 let v1 = kernel.get("out");
4708 kernel.eval(&[42]);
4709 let v2 = kernel.get("out");
4710 assert_eq!(v1, v2);
4711 }
4712
4713 #[test]
4714 fn jit_chain_hash_mod() {
4715 let steps = vec![
4717 (JitOp::Hash, vec![0], vec![1]), (JitOp::ModConst(1_000_000), vec![1], vec![2]), ];
4720 let mut output_map = HashMap::new();
4721 output_map.insert("user_id".into(), 2);
4722 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
4723
4724 kernel.eval(&[42]);
4725 let uid = kernel.get("user_id");
4726 assert!(uid < 1_000_000, "got {uid}");
4727 }
4728
4729 #[test]
4730 fn jit_clamp_const() {
4731 let steps = vec![(JitOp::ClampConst(10, 50), vec![0], vec![1])];
4732 let mut output_map = HashMap::new();
4733 output_map.insert("out".into(), 1);
4734 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4735
4736 kernel.eval(&[5]);
4737 assert_eq!(kernel.get("out"), 10); kernel.eval(&[30]);
4740 assert_eq!(kernel.get("out"), 30); kernel.eval(&[100]);
4743 assert_eq!(kernel.get("out"), 50); }
4745
4746 #[test]
4747 fn jit_interleave() {
4748 let steps = vec![(JitOp::Interleave, vec![0, 1], vec![2])];
4749 let mut output_map = HashMap::new();
4750 output_map.insert("out".into(), 2);
4751 let mut kernel = compile_jit_raw(2, 3, steps, output_map, Vec::new()).unwrap();
4752
4753 kernel.eval(&[0b101, 0b010]);
4754 assert_eq!(kernel.get("out"), 0b01_10_01);
4756 }
4757
4758 #[test]
4759 fn jit_mixed_radix() {
4760 let steps = vec![(
4762 JitOp::MixedRadixConst(vec![100, 1000, 0]),
4763 vec![0],
4764 vec![1, 2, 3],
4765 )];
4766 let mut output_map = HashMap::new();
4767 output_map.insert("d0".into(), 1);
4768 output_map.insert("d1".into(), 2);
4769 output_map.insert("d2".into(), 3);
4770 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
4771
4772 kernel.eval(&[4_201_337]);
4774 assert_eq!(kernel.get("d0"), 37);
4775 assert_eq!(kernel.get("d1"), 13);
4776 assert_eq!(kernel.get("d2"), 42);
4777 }
4778
4779 #[test]
4780 fn jit_unit_interval() {
4781 let steps = vec![(JitOp::UnitInterval, vec![0], vec![1])];
4782 let mut output_map = HashMap::new();
4783 output_map.insert("out".into(), 1);
4784 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4785
4786 kernel.eval(&[0]);
4787 let v = f64::from_bits(kernel.get("out"));
4788 assert!((v - 0.0).abs() < 1e-10);
4789
4790 kernel.eval(&[u64::MAX]);
4791 let v = f64::from_bits(kernel.get("out"));
4792 assert!((v - 1.0).abs() < 1e-10);
4793 }
4794
4795 #[test]
4796 fn jit_f64_to_u64() {
4797 let steps = vec![(JitOp::F64ToU64, vec![0], vec![1])];
4799 let mut output_map = HashMap::new();
4800 output_map.insert("out".into(), 1);
4801 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4802
4803 kernel.eval(&[3.7f64.to_bits()]);
4804 assert_eq!(kernel.get("out"), 3); }
4806
4807 #[test]
4808 fn jit_round_to_u64() {
4809 let steps = vec![(JitOp::RoundToU64, vec![0], vec![1])];
4810 let mut output_map = HashMap::new();
4811 output_map.insert("out".into(), 1);
4812 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4813
4814 kernel.eval(&[3.7f64.to_bits()]);
4815 assert_eq!(kernel.get("out"), 4);
4816
4817 kernel.eval(&[3.2f64.to_bits()]);
4818 assert_eq!(kernel.get("out"), 3);
4819 }
4820
4821 #[test]
4822 fn jit_clamp_f64() {
4823 let steps = vec![(
4824 JitOp::ClampF64Const(0.0f64.to_bits(), 1.0f64.to_bits()),
4825 vec![0],
4826 vec![1],
4827 )];
4828 let mut output_map = HashMap::new();
4829 output_map.insert("out".into(), 1);
4830 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4831
4832 kernel.eval(&[(-0.5f64).to_bits()]);
4833 assert_eq!(f64::from_bits(kernel.get("out")), 0.0);
4834
4835 kernel.eval(&[0.5f64.to_bits()]);
4836 assert_eq!(f64::from_bits(kernel.get("out")), 0.5);
4837
4838 kernel.eval(&[1.5f64.to_bits()]);
4839 assert_eq!(f64::from_bits(kernel.get("out")), 1.0);
4840 }
4841
4842 #[test]
4843 fn jit_lerp() {
4844 let steps = vec![(
4845 JitOp::LerpConst(10.0f64.to_bits(), 20.0f64.to_bits()),
4846 vec![0],
4847 vec![1],
4848 )];
4849 let mut output_map = HashMap::new();
4850 output_map.insert("out".into(), 1);
4851 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4852
4853 kernel.eval(&[0.0f64.to_bits()]);
4854 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
4855
4856 kernel.eval(&[1.0f64.to_bits()]);
4857 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
4858
4859 kernel.eval(&[0.5f64.to_bits()]);
4860 assert_eq!(f64::from_bits(kernel.get("out")), 15.0);
4861 }
4862
4863 #[test]
4864 fn jit_scale_range() {
4865 let steps = vec![(
4866 JitOp::ScaleRangeConst(10.0f64.to_bits(), 10.0f64.to_bits()),
4867 vec![0],
4868 vec![1],
4869 )];
4870 let mut output_map = HashMap::new();
4871 output_map.insert("out".into(), 1);
4872 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4873
4874 kernel.eval(&[0]);
4875 let v = f64::from_bits(kernel.get("out"));
4876 assert!((v - 10.0).abs() < 0.001);
4877
4878 kernel.eval(&[u64::MAX]);
4879 let v = f64::from_bits(kernel.get("out"));
4880 assert!((v - 20.0).abs() < 0.001);
4881 }
4882
4883 #[test]
4884 fn jit_quantize() {
4885 let steps = vec![(JitOp::QuantizeConst(10.0f64.to_bits()), vec![0], vec![1])];
4886 let mut output_map = HashMap::new();
4887 output_map.insert("out".into(), 1);
4888 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4889
4890 kernel.eval(&[13.0f64.to_bits()]);
4891 assert_eq!(f64::from_bits(kernel.get("out")), 10.0);
4892
4893 kernel.eval(&[17.0f64.to_bits()]);
4894 assert_eq!(f64::from_bits(kernel.get("out")), 20.0);
4895 }
4896
4897 #[test]
4898 fn jit_discretize() {
4899 let steps = vec![(
4900 JitOp::DiscretizeConst(100.0f64.to_bits(), 10),
4901 vec![0],
4902 vec![1],
4903 )];
4904 let mut output_map = HashMap::new();
4905 output_map.insert("out".into(), 1);
4906 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4907
4908 kernel.eval(&[0.0f64.to_bits()]);
4909 assert_eq!(kernel.get("out"), 0);
4910
4911 kernel.eval(&[55.0f64.to_bits()]);
4912 assert_eq!(kernel.get("out"), 5);
4913
4914 kernel.eval(&[99.0f64.to_bits()]);
4915 assert_eq!(kernel.get("out"), 9);
4916
4917 kernel.eval(&[200.0f64.to_bits()]);
4919 assert_eq!(kernel.get("out"), 9);
4920 }
4921
4922 #[test]
4923 fn jit_chain_unit_interval_lerp() {
4924 let steps = vec![
4926 (JitOp::UnitInterval, vec![0], vec![1]),
4927 (
4928 JitOp::LerpConst(100.0f64.to_bits(), 200.0f64.to_bits()),
4929 vec![1],
4930 vec![2],
4931 ),
4932 ];
4933 let mut output_map = HashMap::new();
4934 output_map.insert("out".into(), 2);
4935 let mut kernel = compile_jit_raw(1, 3, steps, output_map, Vec::new()).unwrap();
4936
4937 kernel.eval(&[0]);
4938 let v = f64::from_bits(kernel.get("out"));
4939 assert!((v - 100.0).abs() < 0.001);
4940
4941 kernel.eval(&[u64::MAX]);
4942 let v = f64::from_bits(kernel.get("out"));
4943 assert!((v - 200.0).abs() < 0.001);
4944 }
4945
4946 #[test]
4947 fn jit_multi_step_chain() {
4948 let steps = vec![
4950 (JitOp::AddConst(10), vec![0], vec![1]),
4951 (JitOp::MulConst(3), vec![1], vec![2]),
4952 (JitOp::ModConst(100), vec![2], vec![3]),
4953 ];
4954 let mut output_map = HashMap::new();
4955 output_map.insert("out".into(), 3);
4956 let mut kernel = compile_jit_raw(1, 4, steps, output_map, Vec::new()).unwrap();
4957
4958 kernel.eval(&[5]);
4959 assert_eq!(kernel.get("out"), 45);
4961 }
4962
4963 #[test]
4966 fn jit_is_positive_check_passes_positive() {
4967 let steps = vec![(
4968 JitOp::IsPositiveCheck {
4969 name_ptr: 0,
4970 name_len: 0,
4971 },
4972 vec![0],
4973 vec![1],
4974 )];
4975 let mut output_map = HashMap::new();
4976 output_map.insert("out".into(), 1);
4977 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4978 kernel.eval(&[42]);
4979 assert_eq!(kernel.get("out"), 42);
4980 kernel.eval(&[u64::MAX]);
4983 assert_eq!(kernel.get("out"), u64::MAX);
4984 }
4985
4986 #[test]
4987 fn jit_in_range_check_passes_interior() {
4988 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
4989 let mut output_map = HashMap::new();
4990 output_map.insert("out".into(), 1);
4991 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
4992 kernel.eval(&[50]);
4993 assert_eq!(kernel.get("out"), 50);
4994 kernel.eval(&[10]);
4996 assert_eq!(kernel.get("out"), 10);
4997 kernel.eval(&[100]);
4998 assert_eq!(kernel.get("out"), 100);
4999 }
5000
5001 #[test]
5011 fn jit_is_one_of_check_passes_allowed_values() {
5012 let steps = vec![(
5013 JitOp::IsOneOfCheck {
5014 allowed: vec![1, 2, 3, 5, 8],
5015 set_ptr: 0,
5016 set_len: 0,
5017 },
5018 vec![0],
5019 vec![1],
5020 )];
5021 let mut output_map = HashMap::new();
5022 output_map.insert("out".into(), 1);
5023 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5024 for v in [1u64, 2, 3, 5, 8] {
5026 kernel.eval(&[v]);
5027 assert_eq!(kernel.get("out"), v);
5028 }
5029 }
5030
5031 #[test]
5032 fn jit_is_one_of_check_accepts_single_element_allow_list() {
5033 let steps = vec![(
5036 JitOp::IsOneOfCheck {
5037 allowed: vec![42],
5038 set_ptr: 0,
5039 set_len: 0,
5040 },
5041 vec![0],
5042 vec![1],
5043 )];
5044 let mut output_map = HashMap::new();
5045 output_map.insert("out".into(), 1);
5046 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5047 kernel.eval(&[42]);
5048 assert_eq!(kernel.get("out"), 42);
5049 }
5050
5051 fn extract_panic_msg(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
5061 payload
5062 .downcast_ref::<String>()
5063 .cloned()
5064 .or_else(|| payload.downcast_ref::<&str>().map(|s| s.to_string()))
5065 .unwrap_or_else(|| "(non-string panic)".into())
5066 }
5067
5068 #[test]
5069 fn jit_is_positive_violation_is_catchable() {
5070 let steps = vec![(
5071 JitOp::IsPositiveCheck {
5072 name_ptr: 0,
5073 name_len: 0,
5074 },
5075 vec![0],
5076 vec![1],
5077 )];
5078 let mut output_map = HashMap::new();
5079 output_map.insert("out".into(), 1);
5080 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5081 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5082 .expect_err("JIT violation should panic");
5083 assert!(extract_panic_msg(err).contains("must be > 0"));
5084 }
5085
5086 #[test]
5087 fn jit_in_range_violation_is_catchable() {
5088 let steps = vec![(JitOp::InRangeCheck(10, 100), vec![0], vec![1])];
5089 let mut output_map = HashMap::new();
5090 output_map.insert("out".into(), 1);
5091 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5092
5093 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[5])))
5094 .expect_err("below-range should panic");
5095 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5096
5097 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[500])))
5098 .expect_err("above-range should panic");
5099 assert!(extract_panic_msg(err).contains("outside [10, 100]"));
5100 }
5101
5102 #[test]
5103 fn jit_is_one_of_violation_is_catchable() {
5104 let steps = vec![(
5105 JitOp::IsOneOfCheck {
5106 allowed: vec![1, 3, 5],
5107 set_ptr: 0,
5108 set_len: 0,
5109 },
5110 vec![0],
5111 vec![1],
5112 )];
5113 let mut output_map = HashMap::new();
5114 output_map.insert("out".into(), 1);
5115 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5116 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[2])))
5117 .expect_err("disallowed value should panic");
5118 assert!(extract_panic_msg(err).contains("not in allowed set"));
5119 }
5120
5121 #[test]
5122 fn invoke_with_catch_restores_slot_after_foreign_panic() {
5123 let caught = std::panic::catch_unwind(|| {
5131 invoke_with_catch(|| panic!("foreign panic"));
5132 });
5133 assert!(caught.is_err(), "foreign panic should propagate out");
5134
5135 let steps = vec![(
5137 JitOp::IsPositiveCheck {
5138 name_ptr: 0,
5139 name_len: 0,
5140 },
5141 vec![0],
5142 vec![1],
5143 )];
5144 let mut output_map = HashMap::new();
5145 output_map.insert("out".into(), 1);
5146 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5147 let err = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5148 .expect_err("JIT violation should panic cleanly after foreign panic");
5149 assert!(extract_panic_msg(err).contains("must be > 0"));
5150
5151 kernel.eval(&[42]);
5153 assert_eq!(kernel.get("out"), 42);
5154 }
5155
5156 #[test]
5157 fn jit_kernel_survives_multiple_violations() {
5158 let steps = vec![(
5162 JitOp::IsPositiveCheck {
5163 name_ptr: 0,
5164 name_len: 0,
5165 },
5166 vec![0],
5167 vec![1],
5168 )];
5169 let mut output_map = HashMap::new();
5170 output_map.insert("out".into(), 1);
5171 let mut kernel = compile_jit_raw(1, 2, steps, output_map, Vec::new()).unwrap();
5172
5173 for _ in 0..3 {
5174 let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| kernel.eval(&[0])))
5175 .expect_err("violation should still panic");
5176 }
5177 kernel.eval(&[42]);
5179 assert_eq!(kernel.get("out"), 42);
5180 }
5181}