Skip to main content

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}