shape_jit/context.rs
1//! JIT Context and Related Data Structures
2//!
3//! Contains the runtime context and data structures used by JIT-compiled code.
4
5use crate::ffi::jit_kinds::*;
6use crate::ffi::value_ffi::*;
7
8// ============================================================================
9// JITContext Field Offsets for Direct Memory Access
10// ============================================================================
11//
12// These must match the #[repr(C)] struct layout of JITContext
13// Regenerate with: rustc --edition 2024 scripts/jit_offsets.rs && ./jit_offsets
14
15// Generic column access - columns are accessed via column_ptrs array indexed by column_map
16// Timestamps pointer for time-based access
17pub const TIMESTAMPS_PTR_OFFSET: i32 = 24;
18
19// DataFrame access offsets
20pub const COLUMN_PTRS_OFFSET: i32 = 32;
21pub const COLUMN_COUNT_OFFSET: i32 = 40;
22pub const ROW_COUNT_OFFSET: i32 = 48;
23pub const CURRENT_ROW_OFFSET: i32 = 56;
24
25// Locals and stack offsets
26pub const LOCALS_OFFSET: i32 = 64;
27pub const STACK_OFFSET: i32 = 2112; // 64 + (256 * 8)
28// Parallel `NativeKind` byte track on `JITContext.stack` per ADR-006 §2.7.7 /
29// Q9 (JIT-side analog of the VM `Vec<u64> + Vec<NativeKind>` lockstep
30// invariant). Every push into `stack[i]` writes the slot's kind code into
31// `stack_kinds[i]` in lockstep; every pop reads both. The kind sources from
32// the producing call signature at JIT-compile time (no tag-bit decode,
33// no `is_heap()` probe) per the §2.7.5 stamp-at-compile-time discipline.
34// 8-byte raw payload per stack slot is preserved; the kind track is a
35// parallel 1-byte-per-slot side array — same shape as the VM
36// `crates/shape-vm/src/executor/vm_impl/stack.rs::VmStack` design.
37pub const STACK_KINDS_OFFSET: i32 = 6208; // 2112 + (512 * 8)
38pub const STACK_PTR_OFFSET: i32 = 6720; // 6208 + 512
39
40// GC safepoint flag pointer offset (for inline safepoint check)
41pub const GC_SAFEPOINT_FLAG_PTR_OFFSET: i32 = 6840;
42
43// ============================================================================
44// Return Type Tags for stack[0]
45// ============================================================================
46//
47// Used by `return_type_tag` field in JITContext to tell the executor how to
48// interpret the raw bits in stack[0] without NaN-box decoding.
49
50/// Legacy NaN-boxed value (default). Executor uses existing unboxing path.
51pub const RETURN_TAG_NANBOXED: u8 = 0;
52/// Raw f64 bits. Executor reads as `f64::from_bits(stack[0])`.
53pub const RETURN_TAG_F64: u8 = 1;
54/// Raw i64 bits. Executor reads as `stack[0] as i64`.
55pub const RETURN_TAG_I64: u8 = 2;
56/// Raw i32 bits (zero-extended to u64). Executor reads as `stack[0] as i32`.
57pub const RETURN_TAG_I32: u8 = 3;
58/// Raw bool (0 or 1). Executor reads as `stack[0] != 0`.
59pub const RETURN_TAG_BOOL: u8 = 4;
60/// Unit / null / no-value return (W11-jit-new-array). Stamped when the
61/// program's terminal expression is a side-effecting call that
62/// produces no value (e.g. `print(x)` at top level). Executor maps
63/// this to `WireValue::Null`.
64pub const RETURN_TAG_UNIT: u8 = 5;
65
66// ============================================================================
67// JIT Runtime-Error Signal Codes
68// ============================================================================
69//
70// A JIT-compiled function returns an `i32` signal. The executor treats any
71// negative signal as an error (`crates/shape-jit/src/executor.rs`); a
72// successful run returns `0`. The codes below carve out specific negative
73// values so a recoverable Shape-level runtime error (one the bytecode VM
74// handles cleanly) surfaces with the SAME diagnostic in `--mode jit` as in
75// `--mode vm`, instead of an `ud2`/`sdiv`-trap that crashes the process.
76//
77// r5c-2-gz-cp2-jit-div: integer division/modulo by zero is the first such
78// error. The codegen emits a guarded branch — on a zero divisor it does an
79// immediate `return_` of `JIT_SIGNAL_DIVISION_BY_ZERO` (the W12 fall-through
80// shape: a clean diagnostic, not a trap). The executor maps the code back to
81// the VM's `VMError::DivisionByZero` text ("Division by zero").
82
83/// Signal returned by a JIT-compiled function when an integer division or
84/// modulo had a zero divisor. The executor maps this to the same
85/// `Division by zero` diagnostic the bytecode VM emits for
86/// `VMError::DivisionByZero`. Mirrors the i64/i32/narrow/u64 div/mod arms in
87/// `crates/shape-jit/src/mir_compiler/rvalues.rs` + `v2_int.rs`.
88pub const JIT_SIGNAL_DIVISION_BY_ZERO: i32 = -2;
89
90/// Byte offset of `return_type_tag` in JITContext (for Cranelift codegen).
91pub const RETURN_TYPE_TAG_OFFSET: usize = std::mem::offset_of!(JITContext, return_type_tag);
92
93/// Byte offset of `pending_call_error` in JITContext (for Cranelift codegen).
94///
95/// `r5c-2-bz-b-jit-err-surface`: a VM-trampoline FFI call (`jit_call_method`)
96/// whose VM-side handler returns `Err` writes `1` here, stores the error
97/// message in the `ffi::control::JIT_RUNTIME_ERROR` thread-local, and returns
98/// a value-shaped placeholder. The MIR emitter loads this byte immediately
99/// after every trampoline FFI call and, when set, branches to a deopt block
100/// that returns a negative `i32` signal — abandoning the JIT frame BEFORE the
101/// placeholder result can reach a heap-kinded `write_place` / refcount-retain
102/// site. The `JITExecutor` surfaces the stored message verbatim. This is the
103/// W12 architectural fall-through (a genuine JIT failure abandons the JIT
104/// frame; the VM produces the clean error) — no runtime tag-bit decode, no
105/// value-shaped error sentinel flowing into the typed slot ABI.
106pub const PENDING_CALL_ERROR_OFFSET: usize =
107 std::mem::offset_of!(JITContext, pending_call_error);
108
109/// Negative `i32` deopt signal returned by a JIT-compiled function when a
110/// VM-trampoline FFI call surfaced an `Err`. Distinct from `-1` (generic JIT
111/// execution error) so `JITExecutor` can route to the stored VM error message.
112pub const SIGNAL_TRAMPOLINE_ERROR: i32 = -3;
113
114/// Signal returned by a JIT-compiled function when a typed-array element
115/// access (`v2_array_get` / `v2_array_set`) had an out-of-bounds index. The
116/// executor maps this to the same `Index out of bounds` diagnostic the
117/// bytecode VM emits for `VMError::IndexOutOfBounds`. The prior JIT codegen
118/// silently produced the element-type zero (read) or skipped the store
119/// (write) — a VM/JIT divergence and a memory-unsafe access surfacing as a
120/// fabricated value. This emits a guarded early `return_` instead: a clean
121/// diagnostic, not silent-wrong-output. Mirrors the `JIT_SIGNAL_DIVISION_BY_ZERO`
122/// fall-through shape.
123pub const JIT_SIGNAL_INDEX_OUT_OF_BOUNDS: i32 = -4;
124
125// ============================================================================
126// Compile-time layout verification for JITContext
127// ============================================================================
128//
129// These assertions ensure the hardcoded byte offsets above remain in sync with
130// the actual #[repr(C)] struct layout. A mismatch will produce a compile error.
131const _: () = {
132 assert!(
133 std::mem::offset_of!(JITContext, timestamps_ptr) == TIMESTAMPS_PTR_OFFSET as usize,
134 "TIMESTAMPS_PTR_OFFSET does not match JITContext layout"
135 );
136 assert!(
137 std::mem::offset_of!(JITContext, column_ptrs) == COLUMN_PTRS_OFFSET as usize,
138 "COLUMN_PTRS_OFFSET does not match JITContext layout"
139 );
140 assert!(
141 std::mem::offset_of!(JITContext, column_count) == COLUMN_COUNT_OFFSET as usize,
142 "COLUMN_COUNT_OFFSET does not match JITContext layout"
143 );
144 assert!(
145 std::mem::offset_of!(JITContext, row_count) == ROW_COUNT_OFFSET as usize,
146 "ROW_COUNT_OFFSET does not match JITContext layout"
147 );
148 assert!(
149 std::mem::offset_of!(JITContext, current_row) == CURRENT_ROW_OFFSET as usize,
150 "CURRENT_ROW_OFFSET does not match JITContext layout"
151 );
152 assert!(
153 std::mem::offset_of!(JITContext, locals) == LOCALS_OFFSET as usize,
154 "LOCALS_OFFSET does not match JITContext layout"
155 );
156 assert!(
157 std::mem::offset_of!(JITContext, stack) == STACK_OFFSET as usize,
158 "STACK_OFFSET does not match JITContext layout"
159 );
160 assert!(
161 std::mem::offset_of!(JITContext, stack_kinds) == STACK_KINDS_OFFSET as usize,
162 "STACK_KINDS_OFFSET does not match JITContext layout (ADR-006 §2.7.7 \
163 parallel-kind track must follow `stack` in lockstep)"
164 );
165 assert!(
166 std::mem::offset_of!(JITContext, stack_ptr) == STACK_PTR_OFFSET as usize,
167 "STACK_PTR_OFFSET does not match JITContext layout"
168 );
169 assert!(
170 std::mem::offset_of!(JITContext, gc_safepoint_flag_ptr) == GC_SAFEPOINT_FLAG_PTR_OFFSET as usize,
171 "GC_SAFEPOINT_FLAG_PTR_OFFSET does not match JITContext layout"
172 );
173 assert!(
174 std::mem::offset_of!(JITContext, pending_call_error) == PENDING_CALL_ERROR_OFFSET,
175 "PENDING_CALL_ERROR_OFFSET does not match JITContext layout"
176 );
177};
178
179// ============================================================================
180// Type Aliases
181// ============================================================================
182
183/// Function pointer type for JIT-compiled strategy functions
184pub type JittedStrategyFn = unsafe extern "C" fn(*mut JITContext) -> i32;
185
186/// Legacy function signature for simple numeric computations
187pub type JittedFn = unsafe extern "C" fn(*mut f64, *const f64, usize) -> f64;
188
189/// OSR entry function signature.
190///
191/// This has the same binary signature as `JittedStrategyFn` -- the difference
192/// is semantic: for OSR entry, the caller pre-fills `JITContext.locals` from
193/// the interpreter's live frame before invocation, and reads modified locals
194/// back on return.
195///
196/// # Arguments
197/// * `ctx_ptr` - Pointer to a `JITContext` with locals pre-filled from the
198/// interpreter frame (marshaled using the `OsrEntryPoint.local_kinds`).
199///
200/// # Returns
201/// * `0` - Success: execution completed. Modified locals are in
202/// `JITContext.locals`. The VM reads them back and continues
203/// at `OsrEntryPoint.exit_ip`.
204/// * `i32::MIN+1` - Deopt requested: a type guard failed mid-loop. The VM
205/// reads locals from `JITContext.locals` and resumes at
206/// the `DeoptInfo.resume_ip` for the failing guard.
207/// * Other negative - Error.
208pub type OsrEntryFn = unsafe extern "C" fn(*mut JITContext) -> i32;
209
210// ============================================================================
211// Simulation Kernel ABI (Zero-Allocation Hot Path)
212// ============================================================================
213
214/// Function pointer type for simulation kernel functions (single series).
215///
216/// This is the "fused step" ABI that enables >10M ticks/sec by:
217/// - Bypassing JITContext setup overhead
218/// - Using direct pointer arithmetic for data access
219/// - Avoiding all allocations in the hot loop
220///
221/// # Arguments
222/// * `cursor_index` - Current position in the series (0-based)
223/// * `series_ptrs` - Pointer to array of column pointers (*const *const f64)
224/// * `state_ptr` - Pointer to TypedObject state (*mut u8)
225///
226/// # Returns
227/// * 0 = continue execution
228/// * 1 = signal generated (written to state)
229/// * negative = error
230///
231/// # Safety
232/// The caller must ensure:
233/// - `cursor_index` is within bounds
234/// - `series_ptrs` points to valid column pointer array
235/// - `state_ptr` points to valid TypedObject with correct schema
236pub type SimulationKernelFn = unsafe extern "C" fn(
237 cursor_index: usize,
238 series_ptrs: *const *const f64,
239 state_ptr: *mut u8,
240) -> i32;
241
242/// Function pointer type for correlated (multi-series) kernel functions.
243///
244/// This extends the simulation kernel ABI to support multiple aligned time series.
245/// Each series is accessed via compile-time resolved indices.
246///
247/// # Arguments
248/// * `cursor_index` - Current position in all series (0-based, must be aligned)
249/// * `series_ptrs` - Pointer to array of series data pointers (*const *const f64)
250/// Each pointer is a single f64 array (one series's data)
251/// * `table_count` - Number of series (for bounds checking, known at compile time)
252/// * `state_ptr` - Pointer to TypedObject state (*mut u8)
253///
254/// # Memory Layout
255/// ```text
256/// series_ptrs[0] -> [spy_close[0], spy_close[1], ..., spy_close[n-1]]
257/// series_ptrs[1] -> [vix_close[0], vix_close[1], ..., vix_close[n-1]]
258/// ...
259/// ```
260///
261/// # JIT Access Pattern
262/// ```asm
263/// ; context.spy (series index 0)
264/// mov rax, [series_ptrs + 0*8] ; series pointer
265/// mov xmm0, [rax + cursor_index*8] ; value at cursor
266/// ```
267///
268/// # Returns
269/// * 0 = continue execution
270/// * 1 = signal generated (written to state)
271/// * negative = error
272///
273/// # Safety
274/// The caller must ensure:
275/// - `cursor_index` is within bounds for ALL series
276/// - `series_ptrs` points to valid array of `table_count` data pointers
277/// - All series have the same length (aligned timestamps)
278/// - `state_ptr` points to valid TypedObject with correct schema
279pub type CorrelatedKernelFn = unsafe extern "C" fn(
280 cursor_index: usize,
281 series_ptrs: *const *const f64,
282 table_count: usize,
283 state_ptr: *mut u8,
284) -> i32;
285
286/// Configuration for compiling a simulation kernel.
287///
288/// This provides the field offset mappings needed to generate
289/// direct memory access code for the kernel ABI.
290///
291/// Supports two modes:
292/// - **Single-series**: `column_map` maps field names (close, volume) to column indices
293/// - **Multi-series**: `table_map` maps series names (spy, vix) to series indices
294#[derive(Debug, Clone)]
295pub struct SimulationKernelConfig {
296 /// Column index mappings: (field_name, column_index)
297 /// e.g., [("close", 3), ("volume", 4)]
298 /// Used in single-series mode for accessing columns within one series
299 pub column_map: Vec<(String, usize)>,
300
301 /// Series index mappings: (series_name, series_index)
302 /// e.g., [("spy", 0), ("vix", 1), ("temperature", 2)]
303 /// Used in multi-series mode for accessing multiple correlated series
304 ///
305 /// CRITICAL for JIT: Resolved at compile time, NOT runtime.
306 /// `context.spy` → `series_ptrs[0][cursor_idx]`
307 pub table_map: Vec<(String, usize)>,
308
309 /// State field offsets: (field_name, byte_offset)
310 /// e.g., [("cash", 0), ("position", 8), ("entry_price", 16)]
311 pub state_field_offsets: Vec<(String, usize)>,
312
313 /// Schema ID for the state TypedObject
314 pub state_schema_id: u32,
315
316 /// Total number of columns in the data (single-series mode)
317 pub column_count: usize,
318
319 /// Total number of series (multi-series mode)
320 pub table_count: usize,
321}
322
323impl SimulationKernelConfig {
324 /// Create a new kernel config for single-series mode.
325 pub fn new(state_schema_id: u32, column_count: usize) -> Self {
326 Self {
327 column_map: Vec::new(),
328 table_map: Vec::new(),
329 state_field_offsets: Vec::new(),
330 state_schema_id,
331 column_count,
332 table_count: 0,
333 }
334 }
335
336 /// Create a new kernel config for multi-series (correlated) mode.
337 ///
338 /// Use this when simulating across multiple aligned time series
339 /// (e.g., SPY vs VIX, temperature vs pressure).
340 pub fn new_multi_table(state_schema_id: u32, table_count: usize) -> Self {
341 Self {
342 column_map: Vec::new(),
343 table_map: Vec::new(),
344 state_field_offsets: Vec::new(),
345 state_schema_id,
346 column_count: 0,
347 table_count,
348 }
349 }
350
351 /// Map a data field name to a column index (single-series mode).
352 pub fn map_column(mut self, field_name: &str, column_index: usize) -> Self {
353 self.column_map.push((field_name.to_string(), column_index));
354 self
355 }
356
357 /// Map a series name to a series index (multi-series mode).
358 ///
359 /// CRITICAL: This mapping is resolved at compile time.
360 /// `context.spy` in Shape → `series_ptrs[0][cursor_idx]` in generated code.
361 pub fn map_series(mut self, series_name: &str, series_index: usize) -> Self {
362 self.table_map.push((series_name.to_string(), series_index));
363 self
364 }
365
366 /// Map a state field name to a byte offset.
367 pub fn map_state_field(mut self, field_name: &str, offset: usize) -> Self {
368 self.state_field_offsets
369 .push((field_name.to_string(), offset));
370 self
371 }
372
373 /// Get column index for a field name (single-series mode).
374 pub fn get_column_index(&self, field_name: &str) -> Option<usize> {
375 self.column_map
376 .iter()
377 .find(|(name, _)| name == field_name)
378 .map(|(_, idx)| *idx)
379 }
380
381 /// Get series index for a series name (multi-series mode).
382 ///
383 /// This is used by the JIT compiler at compile time.
384 pub fn get_series_index(&self, series_name: &str) -> Option<usize> {
385 self.table_map
386 .iter()
387 .find(|(name, _)| name == series_name)
388 .map(|(_, idx)| *idx)
389 }
390
391 /// Get state field offset for a field name.
392 pub fn get_state_offset(&self, field_name: &str) -> Option<usize> {
393 self.state_field_offsets
394 .iter()
395 .find(|(name, _)| name == field_name)
396 .map(|(_, offset)| *offset)
397 }
398
399 /// Check if this config is for multi-series mode.
400 pub fn is_multi_table(&self) -> bool {
401 self.table_count > 0 || !self.table_map.is_empty()
402 }
403}
404
405// ============================================================================
406// JIT Data Structures
407// ============================================================================
408
409/// JIT-compatible closure structure
410/// Holds function_id and a pointer to a heap-allocated array of captured values.
411/// Supports unlimited captures (no fixed-size limit).
412#[repr(C)]
413pub struct JITClosure {
414 pub function_id: u16,
415 pub captures_count: u16,
416 pub captures_ptr: *const u64, // Pointer to heap-allocated capture array (NaN-boxed)
417}
418
419impl JITClosure {
420 /// Create a new JITClosure with dynamically allocated captures.
421 ///
422 /// The captures slice is copied into a heap-allocated `Box<[u64]>` that is
423 /// leaked into a raw pointer. Call `drop_captures()` to reclaim the memory.
424 pub fn new(function_id: u16, captures: &[u64]) -> Box<Self> {
425 let captures_box: Box<[u64]> = captures.to_vec().into_boxed_slice();
426 let captures_ptr = Box::into_raw(captures_box) as *const u64;
427 Box::new(JITClosure {
428 function_id,
429 captures_count: captures.len() as u16,
430 captures_ptr,
431 })
432 }
433
434 /// Safely read a capture value by index.
435 ///
436 /// # Safety
437 /// The captures_ptr must be valid and index must be < captures_count.
438 #[inline]
439 pub unsafe fn get_capture(&self, index: usize) -> u64 {
440 debug_assert!(index < self.captures_count as usize);
441 unsafe { *self.captures_ptr.add(index) }
442 }
443
444 /// Free the heap-allocated captures array.
445 ///
446 /// Idempotent: safe to call multiple times (no-op after first call).
447 ///
448 /// # Safety
449 /// The captures_ptr must point to a valid allocation created by `new()`,
450 /// or be null (no-op).
451 pub unsafe fn drop_captures(&mut self) {
452 if !self.captures_ptr.is_null() && self.captures_count > 0 {
453 let count = self.captures_count as usize;
454 let _ = unsafe {
455 Box::from_raw(std::slice::from_raw_parts_mut(
456 self.captures_ptr as *mut u64,
457 count,
458 ))
459 };
460 self.captures_ptr = std::ptr::null();
461 }
462 }
463}
464
465impl Drop for JITClosure {
466 fn drop(&mut self) {
467 // SAFETY: drop_captures is idempotent — if captures_ptr is already null
468 // (e.g. from an explicit drop_captures() call), this is a no-op.
469 unsafe { self.drop_captures() };
470 }
471}
472
473/// JIT-compatible duration structure
474#[repr(C)]
475pub struct JITDuration {
476 pub value: f64,
477 pub unit: u8, // 0=seconds, 1=minutes, 2=hours, 3=days, 4=weeks, 5=bars
478}
479
480impl JITDuration {
481 pub fn new(value: f64, unit: u8) -> Box<Self> {
482 Box::new(JITDuration { value, unit })
483 }
484
485 pub fn box_duration(duration: Box<JITDuration>) -> u64 {
486 use crate::ffi::jit_kinds::jit_box;
487 use crate::ffi::value_ffi::HK_DURATION;
488 jit_box(HK_DURATION, *duration)
489 }
490}
491
492/// JIT-compatible range structure
493/// Represents a range with start and end values (both NaN-boxed)
494#[repr(C)]
495pub struct JITRange {
496 pub start: u64, // NaN-boxed start value
497 pub end: u64, // NaN-boxed end value
498}
499
500impl JITRange {
501 pub fn new(start: u64, end: u64) -> Box<Self> {
502 Box::new(JITRange { start, end })
503 }
504
505 pub fn box_range(range: Box<JITRange>) -> u64 {
506 use crate::ffi::jit_kinds::jit_box;
507 use crate::ffi::value_ffi::HK_RANGE;
508 jit_box(HK_RANGE, *range)
509 }
510}
511
512/// JIT-compatible data reference structure
513/// Represents a reference to a specific data row in time
514#[repr(C)]
515pub struct JITDataReference {
516 pub timestamp: i64,
517 pub symbol: *const String, // Pointer to symbol string
518 pub timeframe_value: u32, // Timeframe value
519 pub timeframe_unit: u8, // 0=Second, 1=Minute, 2=Hour, 3=Day, 4=Week, 5=Month, 6=Bar
520 pub has_timezone: bool,
521 pub timezone: *const String, // Pointer to timezone string (may be null)
522}
523
524impl JITDataReference {
525 pub fn box_data_ref(data_ref: Box<JITDataReference>) -> u64 {
526 use crate::ffi::jit_kinds::jit_box;
527 use crate::ffi::value_ffi::HK_DATA_REFERENCE;
528 jit_box(HK_DATA_REFERENCE, *data_ref)
529 }
530}
531
532// ============================================================================
533// JITContext - Main Execution Context
534// ============================================================================
535
536/// JIT execution context passed to compiled functions
537/// This struct must be C-compatible (#[repr(C)]) for FFI
538///
539/// Uses NaN-boxing for full type support
540#[repr(C)]
541#[derive(Debug, Clone)]
542pub struct JITContext {
543 // Position state
544 pub in_position: bool,
545 pub position_side: i8, // 0=None, 1=Long, -1=Short
546 pub entry_price: u64, // NaN-boxed f64
547 pub unrealized_pnl_pct: u64, // NaN-boxed f64
548
549 // Timestamps pointer for time-based data access
550 pub timestamps_ptr: *const i64,
551
552 // ========== Generic DataFrame Access (industry-agnostic) ==========
553 /// Array of column pointers (SIMD-aligned f64 arrays)
554 /// Column order matches DataFrameSchema.column_names
555 pub column_ptrs: *const *const f64,
556 /// Number of columns in the DataFrame
557 pub column_count: usize,
558 /// Number of rows in the DataFrame
559 pub row_count: usize,
560 /// Current row index (for backtest iteration)
561 pub current_row: usize,
562
563 // Local variables (NaN-boxed values)
564 pub locals: [u64; 256],
565
566 // NaN-boxed stack for JIT execution
567 pub stack: [u64; 512],
568 // Parallel `NativeKind` byte track per ADR-006 §2.7.7 / Q9 — the JIT-side
569 // analog of the VM `Vec<u64> + Vec<NativeKind>` lockstep invariant.
570 // `stack_kinds[i]` carries the §2.7.5 stamp-at-compile-time kind code
571 // for the slot at `stack[i]`, written by the MIR emitter at every push
572 // site in lockstep with the data write and read at every pop site (no
573 // tag-bit decode, no `is_heap()` probe). See
574 // `crates/shape-jit/src/ffi/stack_kind_code.rs` for the encoding and
575 // `crates/shape-vm/src/executor/vm_impl/stack.rs` for the VM mirror.
576 pub stack_kinds: [u8; 512],
577 pub stack_ptr: usize,
578
579 // Heap object storage (owned by VM, JIT just holds pointers)
580 pub heap_ptr: *mut std::ffi::c_void,
581
582 // Function table for Call opcode (pointer to array of function pointers)
583 pub function_table: *const JittedStrategyFn,
584 pub function_table_len: usize,
585
586 // ExecutionContext pointer for fallback to interpreter
587 pub exec_context_ptr: *mut std::ffi::c_void,
588
589 // Function names for closure-to-Value conversion
590 // Points to contiguous String array from BytecodeProgram.functions
591 pub function_names_ptr: *const String,
592 pub function_names_len: usize,
593
594 // ========== Async Execution Support ==========
595 /// Pointer to event queue (for FFI calls to poll/push events)
596 /// Points to a SharedEventQueue behind the scenes
597 pub event_queue_ptr: *mut std::ffi::c_void,
598
599 /// Suspension state: 0 = running, 1 = yielded, 2 = suspended
600 pub suspension_state: u32,
601
602 /// Iterations since last yield (for cooperative scheduling)
603 pub iterations_since_yield: u64,
604
605 /// Yield threshold - yield after this many iterations
606 /// 0 = never yield automatically
607 pub yield_threshold: u64,
608
609 /// Alert pipeline pointer (for FFI calls to emit alerts)
610 /// Points to AlertRouter behind the scenes
611 pub alert_pipeline_ptr: *mut std::ffi::c_void,
612
613 // ========== Simulation Mode Support ==========
614 /// Simulation mode: 0 = disabled, 1 = DenseKernel, 2 = HybridKernel
615 pub simulation_mode: u32,
616
617 /// Pointer to simulation state (TypedObject for DenseKernel)
618 /// JIT code accesses state fields via direct memory offset
619 pub simulation_state_ptr: *mut u8,
620
621 /// Size of simulation state data (for deallocation)
622 pub simulation_state_size: usize,
623
624 // ========== GC Integration ==========
625 /// Pointer to GC safepoint flag (AtomicBool raw pointer).
626 /// Null when GC is not enabled. The JIT safepoint function reads this
627 /// to determine if a GC cycle is requested.
628 pub gc_safepoint_flag_ptr: *const u8,
629
630 /// Pointer to GcHeap for allocation fast path.
631 /// Null when GC is not enabled.
632 pub gc_heap_ptr: *mut std::ffi::c_void,
633
634 /// Opaque pointer to JIT foreign-call bridge state.
635 /// Null when no foreign functions are linked for this execution.
636 pub foreign_bridge_ptr: *const std::ffi::c_void,
637
638 /// v2: type tag for the return value in stack[0].
639 /// 0 = NaN-boxed (legacy), 1 = raw f64, 2 = raw i64, 3 = raw i32, 4 = raw bool
640 pub return_type_tag: u8,
641
642 /// `r5c-2-bz-b-jit-err-surface`: VM-trampoline error flag.
643 ///
644 /// `0` = no pending error (default). `1` = a VM-trampoline FFI call
645 /// (`jit_call_method`) hit an `Err` from the VM-side handler. The FFI body
646 /// sets this and stores the error message in the
647 /// `ffi::control::JIT_RUNTIME_ERROR` thread-local. The MIR emitter loads
648 /// this byte right after every trampoline FFI call and deopts (returns
649 /// `SIGNAL_TRAMPOLINE_ERROR`) when set — so the FFI's placeholder return
650 /// value never reaches a heap-kinded refcount-retain site.
651 pub pending_call_error: u8,
652}
653
654impl Default for JITContext {
655 fn default() -> Self {
656 Self {
657 in_position: false,
658 position_side: 0,
659 entry_price: box_number(0.0),
660 unrealized_pnl_pct: box_number(0.0),
661 // Timestamps pointer
662 timestamps_ptr: std::ptr::null(),
663 // Generic DataFrame access
664 column_ptrs: std::ptr::null(),
665 column_count: 0,
666 row_count: 0,
667 current_row: 0,
668 // Local variables and stack
669 locals: [TAG_NULL; 256],
670 stack: [TAG_NULL; 512],
671 // ADR-006 §2.7.7: parallel-kind track initialized to the
672 // SENTINEL kind code (`stack_kind_code::SENTINEL`). Live slots
673 // overwrite this with the producing-site kind in lockstep with
674 // the data write; the sentinel surfaces a kind-source gap if a
675 // pop reads an unwritten slot (forbidden #9 / W-series Bool-
676 // default rationalization).
677 stack_kinds: [crate::ffi::stack_kind_code::SENTINEL; 512],
678 stack_ptr: 0,
679 heap_ptr: std::ptr::null_mut(),
680 function_table: std::ptr::null(),
681 function_table_len: 0,
682 exec_context_ptr: std::ptr::null_mut(),
683 function_names_ptr: std::ptr::null(),
684 function_names_len: 0,
685 // Async execution support
686 event_queue_ptr: std::ptr::null_mut(),
687 suspension_state: 0,
688 iterations_since_yield: 0,
689 yield_threshold: 0, // 0 = no automatic yielding
690 alert_pipeline_ptr: std::ptr::null_mut(),
691 // Simulation mode support
692 simulation_mode: 0,
693 simulation_state_ptr: std::ptr::null_mut(),
694 simulation_state_size: 0,
695 // GC integration
696 gc_safepoint_flag_ptr: std::ptr::null(),
697 gc_heap_ptr: std::ptr::null_mut(),
698 foreign_bridge_ptr: std::ptr::null(),
699 return_type_tag: 0,
700 pending_call_error: 0,
701 }
702 }
703}
704
705impl JITContext {
706 /// Get column value at offset from current row
707 /// column_index is the column index in the DataFrame schema
708 pub fn get_column_value(&self, column_index: usize, offset: i32) -> f64 {
709 if self.column_ptrs.is_null() || column_index >= self.column_count {
710 return 0.0;
711 }
712 let row_idx = (self.current_row as i32 + offset) as usize;
713 if row_idx < self.row_count {
714 unsafe {
715 let col_ptr = *self.column_ptrs.add(column_index);
716 if !col_ptr.is_null() {
717 *col_ptr.add(row_idx)
718 } else {
719 0.0
720 }
721 }
722 } else {
723 0.0
724 }
725 }
726
727 /// Update current row index for DataFrame iteration
728 #[inline]
729 pub fn set_current_row(&mut self, index: usize) {
730 self.current_row = index;
731 }
732
733 /// Update current row for backtest iteration (alias for backward compatibility)
734 #[inline]
735 pub fn update_current_row(&mut self, index: usize) {
736 self.current_row = index;
737 }
738
739 // ========================================================================
740 // Simulation Mode Methods
741 // ========================================================================
742
743 /// Check if in simulation mode
744 #[inline]
745 pub fn is_simulation_mode(&self) -> bool {
746 self.simulation_mode > 0
747 }
748
749 /// Set up context for DenseKernel simulation.
750 ///
751 /// # Arguments
752 /// * `state_ptr` - Pointer to TypedObject state
753 /// * `state_size` - Size of state data
754 /// * `column_ptrs` - Pointers to data columns
755 /// * `column_count` - Number of columns
756 /// * `row_count` - Number of rows
757 /// * `timestamps` - Pointer to timestamp array
758 pub fn setup_simulation(
759 &mut self,
760 state_ptr: *mut u8,
761 state_size: usize,
762 column_ptrs: *const *const f64,
763 column_count: usize,
764 row_count: usize,
765 timestamps: *const i64,
766 ) {
767 self.simulation_mode = 1; // DenseKernel mode
768 self.simulation_state_ptr = state_ptr;
769 self.simulation_state_size = state_size;
770 self.column_ptrs = column_ptrs;
771 self.column_count = column_count;
772 self.row_count = row_count;
773 self.current_row = 0;
774 self.timestamps_ptr = timestamps;
775 }
776
777 /// Get simulation state field as f64.
778 ///
779 /// # Safety
780 /// Caller must ensure offset is valid for the state TypedObject.
781 #[inline]
782 pub unsafe fn get_state_field_f64(&self, offset: usize) -> f64 {
783 if self.simulation_state_ptr.is_null() {
784 return 0.0;
785 }
786 let field_ptr = unsafe { self.simulation_state_ptr.add(8 + offset) } as *const u64;
787 let bits = unsafe { *field_ptr };
788 unbox_number(bits)
789 }
790
791 /// Set simulation state field as f64.
792 ///
793 /// # Safety
794 /// Caller must ensure offset is valid for the state TypedObject.
795 #[inline]
796 pub unsafe fn set_state_field_f64(&mut self, offset: usize, value: f64) {
797 if self.simulation_state_ptr.is_null() {
798 return;
799 }
800 let field_ptr = unsafe { self.simulation_state_ptr.add(8 + offset) } as *mut u64;
801 unsafe { *field_ptr = box_number(value) };
802 }
803
804 /// Clear simulation mode.
805 pub fn clear_simulation(&mut self) {
806 self.simulation_mode = 0;
807 self.simulation_state_ptr = std::ptr::null_mut();
808 self.simulation_state_size = 0;
809 }
810}
811
812// ============================================================================
813// JITDataFrame - Generic DataFrame for JIT (industry-agnostic)
814// ============================================================================
815
816/// Generic DataFrame storage for JIT execution.
817/// Stores data as an array of columns, matching the generic column_ptrs
818/// design in JITContext.
819///
820/// Column order MUST match the DataFrameSchema used during compilation.
821pub struct JITDataFrame {
822 /// Column data arrays (each Vec is one column)
823 /// Columns are ordered by index as defined in DataFrameSchema
824 pub columns: Vec<Vec<f64>>,
825 /// Pointers to column data (for JITContext.column_ptrs)
826 pub column_ptrs: Vec<*const f64>,
827 /// Timestamps (always present, column 0 equivalent)
828 pub timestamps: Vec<i64>,
829 /// Number of rows
830 pub row_count: usize,
831}
832
833impl JITDataFrame {
834 /// Create an empty JITDataFrame
835 pub fn new() -> Self {
836 Self {
837 columns: Vec::new(),
838 column_ptrs: Vec::new(),
839 timestamps: Vec::new(),
840 row_count: 0,
841 }
842 }
843
844 /// Create from ExecutionContext using a schema mapping.
845 /// The schema determines which columns to extract and their order.
846 pub fn from_execution_context(
847 ctx: &shape_runtime::context::ExecutionContext,
848 schema: &shape_vm::bytecode::DataFrameSchema,
849 ) -> Self {
850 let mut data = Self::new();
851
852 // NOTE: Series caching not yet implemented in ExecutionContext
853 // For now, initialize empty columns for each schema column
854 // TODO: Implement series caching when available
855 let _ = (ctx, schema); // Suppress unused warnings
856 for _ in 0..schema.column_names.len() {
857 data.columns.push(Vec::new());
858 data.column_ptrs.push(std::ptr::null());
859 }
860
861 data
862 }
863
864 /// Populate a JITContext with generic DataFrame pointers.
865 /// This sets column_ptrs, column_count, row_count, and timestamps_ptr.
866 pub fn populate_context(&self, ctx: &mut JITContext) {
867 if !self.column_ptrs.is_empty() {
868 ctx.column_ptrs = self.column_ptrs.as_ptr();
869 ctx.column_count = self.column_ptrs.len();
870 }
871 ctx.row_count = self.row_count;
872
873 if !self.timestamps.is_empty() {
874 ctx.timestamps_ptr = self.timestamps.as_ptr();
875 }
876 }
877
878 /// Get the number of rows
879 pub fn len(&self) -> usize {
880 self.row_count
881 }
882
883 /// Check if empty
884 pub fn is_empty(&self) -> bool {
885 self.row_count == 0
886 }
887
888 /// Get number of columns
889 pub fn column_count(&self) -> usize {
890 self.columns.len()
891 }
892
893 /// Create from a DataTable by extracting f64 columns and an optional timestamp column.
894 ///
895 /// All f64 columns are copied into SIMD-aligned buffers. If a column named
896 /// "timestamp" (or typed as Timestamp) exists, it is extracted as i64.
897 pub fn from_datatable(dt: &shape_value::DataTable) -> Self {
898 use arrow_array::cast::AsArray;
899 use arrow_schema::{DataType, TimeUnit};
900
901 let batch = dt.inner();
902 let schema = batch.schema();
903 let num_rows = batch.num_rows();
904 let mut columns = Vec::new();
905 let mut timestamps = Vec::new();
906
907 for (i, field) in schema.fields().iter().enumerate() {
908 match field.data_type() {
909 DataType::Float64 => {
910 let arr = batch
911 .column(i)
912 .as_primitive::<arrow_array::types::Float64Type>();
913 let col: Vec<f64> = (0..num_rows).map(|r| arr.value(r)).collect();
914 columns.push(col);
915 }
916 DataType::Timestamp(TimeUnit::Microsecond, _) => {
917 let arr = batch
918 .column(i)
919 .as_primitive::<arrow_array::types::TimestampMicrosecondType>();
920 timestamps = (0..num_rows).map(|r| arr.value(r)).collect();
921 }
922 DataType::Int64 => {
923 // Convert i64 to f64 for JIT column access
924 let arr = batch
925 .column(i)
926 .as_primitive::<arrow_array::types::Int64Type>();
927 let col: Vec<f64> = (0..num_rows).map(|r| arr.value(r) as f64).collect();
928 columns.push(col);
929 }
930 _ => {
931 // Skip non-numeric columns (strings, bools, etc.)
932 }
933 }
934 }
935
936 let column_ptrs: Vec<*const f64> = columns.iter().map(|c| c.as_ptr()).collect();
937
938 Self {
939 columns,
940 column_ptrs,
941 timestamps,
942 row_count: num_rows,
943 }
944 }
945}
946
947impl Default for JITDataFrame {
948 fn default() -> Self {
949 Self::new()
950 }
951}
952
953// ============================================================================
954// JITConfig - Compilation Configuration
955// ============================================================================
956
957/// JIT compilation configuration
958#[derive(Debug, Clone)]
959pub struct JITConfig {
960 /// Optimization level (0-3)
961 pub opt_level: u8,
962 /// Enable debug symbols
963 pub debug_symbols: bool,
964 /// Minimum execution count before JIT compilation
965 pub jit_threshold: usize,
966}
967
968impl Default for JITConfig {
969 fn default() -> Self {
970 Self {
971 opt_level: 3,
972 debug_symbols: false,
973 jit_threshold: 100,
974 }
975 }
976}
977
978#[cfg(test)]
979mod tests {
980 use super::*;
981
982 #[test]
983 fn test_closure_dynamic_captures_0() {
984 // Zero captures — captures_ptr should be a valid (empty) allocation
985 let closure = JITClosure::new(42, &[]);
986 assert_eq!(closure.function_id, 42);
987 assert_eq!(closure.captures_count, 0);
988 // Drop is safe even with 0 captures
989 let mut closure = closure;
990 unsafe { closure.drop_captures() };
991 }
992
993 #[test]
994 fn test_closure_dynamic_captures_5() {
995 // Typical case: 5 captures
996 let captures = [
997 box_number(1.0),
998 box_number(2.0),
999 box_number(3.0),
1000 TAG_BOOL_TRUE,
1001 TAG_NULL,
1002 ];
1003 let closure = JITClosure::new(7, &captures);
1004 assert_eq!(closure.function_id, 7);
1005 assert_eq!(closure.captures_count, 5);
1006
1007 unsafe {
1008 assert_eq!(unbox_number(closure.get_capture(0)), 1.0);
1009 assert_eq!(unbox_number(closure.get_capture(1)), 2.0);
1010 assert_eq!(unbox_number(closure.get_capture(2)), 3.0);
1011 assert_eq!(closure.get_capture(3), TAG_BOOL_TRUE);
1012 assert_eq!(closure.get_capture(4), TAG_NULL);
1013 }
1014 }
1015
1016 #[test]
1017 fn test_closure_dynamic_captures_20() {
1018 // Exceeds old 16-capture limit
1019 let captures: Vec<u64> = (0..20).map(|i| box_number(i as f64)).collect();
1020 let closure = JITClosure::new(99, &captures);
1021 assert_eq!(closure.captures_count, 20);
1022
1023 unsafe {
1024 for i in 0..20 {
1025 assert_eq!(unbox_number(closure.get_capture(i)), i as f64);
1026 }
1027 }
1028 }
1029
1030 #[test]
1031 fn test_closure_dynamic_captures_64() {
1032 // Stress test: 64 captures
1033 let captures: Vec<u64> = (0..64).map(|i| box_number(i as f64 * 10.0)).collect();
1034 let closure = JITClosure::new(1, &captures);
1035 assert_eq!(closure.captures_count, 64);
1036
1037 unsafe {
1038 for i in 0..64 {
1039 assert_eq!(unbox_number(closure.get_capture(i)), i as f64 * 10.0);
1040 }
1041 }
1042 }
1043
1044 #[test]
1045 fn test_closure_captures_drop() {
1046 // Verify memory is properly freed (no leak under Miri/ASAN)
1047 let captures: Vec<u64> = (0..32).map(|i| box_number(i as f64)).collect();
1048 let mut closure = JITClosure::new(5, &captures);
1049 assert_eq!(closure.captures_count, 32);
1050
1051 // Verify captures are valid before drop
1052 unsafe {
1053 assert_eq!(unbox_number(closure.get_capture(0)), 0.0);
1054 assert_eq!(unbox_number(closure.get_capture(31)), 31.0);
1055 }
1056
1057 // Drop captures
1058 unsafe { closure.drop_captures() };
1059 assert!(closure.captures_ptr.is_null());
1060 assert_eq!(closure.captures_count, 32); // count unchanged, ptr nulled
1061 }
1062
1063 /// `jit_box(HK_CLOSURE, JITClosure)` round-trip — strict-typed rewrite
1064 /// (W12-deleted-valuewordshape-tests-rewrite, 2026-05-12).
1065 ///
1066 /// Pre-rewrite the test asserted `is_heap_kind(bits, HK_CLOSURE) ==
1067 /// true` after `jit_box(HK_CLOSURE, *closure)`. Under ADR-006 §2.7.5
1068 /// JIT-side producers return raw `Box::into_raw(...) as u64`; the
1069 /// `is_heap_kind` consumer first checks `is_heap(bits)` which requires
1070 /// `is_tagged(bits)` (negative-NaN tag bits). Raw pointers don't have
1071 /// those tag bits, so `is_heap_kind` returns false — the consumer is
1072 /// in the production-code migration gap surfaced separately.
1073 ///
1074 /// The strict-typed discriminator reads the `kind: u16` prefix at
1075 /// offset 0 of the JIT allocation via `read_heap_kind(bits)`. Per
1076 /// §2.7.5 "*not* tag-bit dispatch — it reads a field from a
1077 /// heap-resident struct that the producing call placed there."
1078 ///
1079 /// Same round-trip semantics expressed through the strict-typed
1080 /// predicate.
1081 #[test]
1082 fn test_closure_jit_box_roundtrip_via_heap_kind_prefix() {
1083 // Verify JITClosure survives jit_box/jit_unbox roundtrip
1084 let captures = [box_number(42.0), TAG_BOOL_FALSE];
1085 let closure = JITClosure::new(10, &captures);
1086 let bits = jit_box(HK_CLOSURE, *closure);
1087
1088 // Construction-side contract: `jit_box` writes `HK_CLOSURE` into
1089 // the `kind: u16` prefix at offset 0 of the `JitAlloc<JITClosure>`
1090 // allocation. The strict-typed §2.7.5 discriminator reads the
1091 // prefix directly — no tag-bit dispatch.
1092 assert_ne!(bits, 0, "allocation pointer is non-null");
1093 assert_eq!(
1094 unsafe { crate::ffi::jit_kinds::read_heap_kind(bits) },
1095 HK_CLOSURE,
1096 "heap-kind prefix at offset 0 discriminates the JIT allocation"
1097 );
1098
1099 // Round-trip via direct `jit_unbox`: reads the `data` field of
1100 // the `JitAlloc<JITClosure>` without gating on tag bits, recovering
1101 // the JITClosure the producer stored.
1102 let recovered = unsafe { jit_unbox::<JITClosure>(bits) };
1103 assert_eq!(recovered.function_id, 10);
1104 assert_eq!(recovered.captures_count, 2);
1105 unsafe {
1106 assert_eq!(unbox_number(recovered.get_capture(0)), 42.0);
1107 assert_eq!(recovered.get_capture(1), TAG_BOOL_FALSE);
1108 }
1109
1110 // Clean up the JitAlloc allocation directly. The deleted
1111 // ValueWord-shape clean-up went through the JIT-emitted `jit_drop`
1112 // call; here we go straight to `jit_drop::<JITClosure>` which is
1113 // the §2.7.5 direct-path cleanup.
1114 unsafe { crate::ffi::jit_kinds::jit_drop::<JITClosure>(bits) };
1115 }
1116
1117 #[test]
1118 fn test_closure_drop_impl_frees_captures_via_jit_drop() {
1119 // Verify the Drop impl on JITClosure frees the captures array
1120 // when the owning JitAlloc is freed via jit_drop.
1121 // Under Miri/ASAN this would catch a leak if Drop didn't work.
1122 let captures: Vec<u64> = (0..24).map(|i| box_number(i as f64)).collect();
1123 let closure = JITClosure::new(3, &captures);
1124 let bits = jit_box(HK_CLOSURE, *closure);
1125
1126 // Read captures to confirm they're valid
1127 let recovered = unsafe { jit_unbox::<JITClosure>(bits) };
1128 assert_eq!(recovered.captures_count, 24);
1129 unsafe {
1130 assert_eq!(unbox_number(recovered.get_capture(23)), 23.0);
1131 }
1132
1133 // jit_drop frees JitAlloc<JITClosure>, which calls Drop::drop on
1134 // JITClosure, which frees the captures array.
1135 unsafe { jit_drop::<JITClosure>(bits) };
1136 }
1137
1138 #[test]
1139 fn test_closure_implicit_drop_on_box() {
1140 // Verify that simply dropping a Box<JITClosure> frees the captures.
1141 // (This tests the Drop impl without jit_box involvement.)
1142 let captures: Vec<u64> = (0..10).map(|i| box_number(i as f64)).collect();
1143 let closure = JITClosure::new(1, &captures);
1144 // closure is Box<JITClosure>, dropping it should free captures via Drop
1145 drop(closure);
1146 // No leak under Miri/ASAN
1147 }
1148}