use crate::ffi::jit_kinds::*;
use crate::ffi::value_ffi::*;
pub const TIMESTAMPS_PTR_OFFSET: i32 = 24;
pub const COLUMN_PTRS_OFFSET: i32 = 32;
pub const COLUMN_COUNT_OFFSET: i32 = 40;
pub const ROW_COUNT_OFFSET: i32 = 48;
pub const CURRENT_ROW_OFFSET: i32 = 56;
pub const LOCALS_OFFSET: i32 = 64;
pub const STACK_OFFSET: i32 = 2112; pub const STACK_KINDS_OFFSET: i32 = 6208; pub const STACK_PTR_OFFSET: i32 = 6720;
pub const GC_SAFEPOINT_FLAG_PTR_OFFSET: i32 = 6840;
pub const RETURN_TAG_NANBOXED: u8 = 0;
pub const RETURN_TAG_F64: u8 = 1;
pub const RETURN_TAG_I64: u8 = 2;
pub const RETURN_TAG_I32: u8 = 3;
pub const RETURN_TAG_BOOL: u8 = 4;
pub const RETURN_TAG_UNIT: u8 = 5;
pub const JIT_SIGNAL_DIVISION_BY_ZERO: i32 = -2;
pub const RETURN_TYPE_TAG_OFFSET: usize = std::mem::offset_of!(JITContext, return_type_tag);
pub const PENDING_CALL_ERROR_OFFSET: usize =
std::mem::offset_of!(JITContext, pending_call_error);
pub const SIGNAL_TRAMPOLINE_ERROR: i32 = -3;
pub const JIT_SIGNAL_INDEX_OUT_OF_BOUNDS: i32 = -4;
const _: () = {
assert!(
std::mem::offset_of!(JITContext, timestamps_ptr) == TIMESTAMPS_PTR_OFFSET as usize,
"TIMESTAMPS_PTR_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, column_ptrs) == COLUMN_PTRS_OFFSET as usize,
"COLUMN_PTRS_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, column_count) == COLUMN_COUNT_OFFSET as usize,
"COLUMN_COUNT_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, row_count) == ROW_COUNT_OFFSET as usize,
"ROW_COUNT_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, current_row) == CURRENT_ROW_OFFSET as usize,
"CURRENT_ROW_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, locals) == LOCALS_OFFSET as usize,
"LOCALS_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, stack) == STACK_OFFSET as usize,
"STACK_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, stack_kinds) == STACK_KINDS_OFFSET as usize,
"STACK_KINDS_OFFSET does not match JITContext layout (ADR-006 §2.7.7 \
parallel-kind track must follow `stack` in lockstep)"
);
assert!(
std::mem::offset_of!(JITContext, stack_ptr) == STACK_PTR_OFFSET as usize,
"STACK_PTR_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, gc_safepoint_flag_ptr) == GC_SAFEPOINT_FLAG_PTR_OFFSET as usize,
"GC_SAFEPOINT_FLAG_PTR_OFFSET does not match JITContext layout"
);
assert!(
std::mem::offset_of!(JITContext, pending_call_error) == PENDING_CALL_ERROR_OFFSET,
"PENDING_CALL_ERROR_OFFSET does not match JITContext layout"
);
};
pub type JittedStrategyFn = unsafe extern "C" fn(*mut JITContext) -> i32;
pub type JittedFn = unsafe extern "C" fn(*mut f64, *const f64, usize) -> f64;
pub type OsrEntryFn = unsafe extern "C" fn(*mut JITContext) -> i32;
pub type SimulationKernelFn = unsafe extern "C" fn(
cursor_index: usize,
series_ptrs: *const *const f64,
state_ptr: *mut u8,
) -> i32;
pub type CorrelatedKernelFn = unsafe extern "C" fn(
cursor_index: usize,
series_ptrs: *const *const f64,
table_count: usize,
state_ptr: *mut u8,
) -> i32;
#[derive(Debug, Clone)]
pub struct SimulationKernelConfig {
pub column_map: Vec<(String, usize)>,
pub table_map: Vec<(String, usize)>,
pub state_field_offsets: Vec<(String, usize)>,
pub state_schema_id: u32,
pub column_count: usize,
pub table_count: usize,
}
impl SimulationKernelConfig {
pub fn new(state_schema_id: u32, column_count: usize) -> Self {
Self {
column_map: Vec::new(),
table_map: Vec::new(),
state_field_offsets: Vec::new(),
state_schema_id,
column_count,
table_count: 0,
}
}
pub fn new_multi_table(state_schema_id: u32, table_count: usize) -> Self {
Self {
column_map: Vec::new(),
table_map: Vec::new(),
state_field_offsets: Vec::new(),
state_schema_id,
column_count: 0,
table_count,
}
}
pub fn map_column(mut self, field_name: &str, column_index: usize) -> Self {
self.column_map.push((field_name.to_string(), column_index));
self
}
pub fn map_series(mut self, series_name: &str, series_index: usize) -> Self {
self.table_map.push((series_name.to_string(), series_index));
self
}
pub fn map_state_field(mut self, field_name: &str, offset: usize) -> Self {
self.state_field_offsets
.push((field_name.to_string(), offset));
self
}
pub fn get_column_index(&self, field_name: &str) -> Option<usize> {
self.column_map
.iter()
.find(|(name, _)| name == field_name)
.map(|(_, idx)| *idx)
}
pub fn get_series_index(&self, series_name: &str) -> Option<usize> {
self.table_map
.iter()
.find(|(name, _)| name == series_name)
.map(|(_, idx)| *idx)
}
pub fn get_state_offset(&self, field_name: &str) -> Option<usize> {
self.state_field_offsets
.iter()
.find(|(name, _)| name == field_name)
.map(|(_, offset)| *offset)
}
pub fn is_multi_table(&self) -> bool {
self.table_count > 0 || !self.table_map.is_empty()
}
}
#[repr(C)]
pub struct JITClosure {
pub function_id: u16,
pub captures_count: u16,
pub captures_ptr: *const u64, }
impl JITClosure {
pub fn new(function_id: u16, captures: &[u64]) -> Box<Self> {
let captures_box: Box<[u64]> = captures.to_vec().into_boxed_slice();
let captures_ptr = Box::into_raw(captures_box) as *const u64;
Box::new(JITClosure {
function_id,
captures_count: captures.len() as u16,
captures_ptr,
})
}
#[inline]
pub unsafe fn get_capture(&self, index: usize) -> u64 {
debug_assert!(index < self.captures_count as usize);
unsafe { *self.captures_ptr.add(index) }
}
pub unsafe fn drop_captures(&mut self) {
if !self.captures_ptr.is_null() && self.captures_count > 0 {
let count = self.captures_count as usize;
let _ = unsafe {
Box::from_raw(std::slice::from_raw_parts_mut(
self.captures_ptr as *mut u64,
count,
))
};
self.captures_ptr = std::ptr::null();
}
}
}
impl Drop for JITClosure {
fn drop(&mut self) {
unsafe { self.drop_captures() };
}
}
#[repr(C)]
pub struct JITDuration {
pub value: f64,
pub unit: u8, }
impl JITDuration {
pub fn new(value: f64, unit: u8) -> Box<Self> {
Box::new(JITDuration { value, unit })
}
pub fn box_duration(duration: Box<JITDuration>) -> u64 {
use crate::ffi::jit_kinds::jit_box;
use crate::ffi::value_ffi::HK_DURATION;
jit_box(HK_DURATION, *duration)
}
}
#[repr(C)]
pub struct JITRange {
pub start: u64, pub end: u64, }
impl JITRange {
pub fn new(start: u64, end: u64) -> Box<Self> {
Box::new(JITRange { start, end })
}
pub fn box_range(range: Box<JITRange>) -> u64 {
use crate::ffi::jit_kinds::jit_box;
use crate::ffi::value_ffi::HK_RANGE;
jit_box(HK_RANGE, *range)
}
}
#[repr(C)]
pub struct JITDataReference {
pub timestamp: i64,
pub symbol: *const String, pub timeframe_value: u32, pub timeframe_unit: u8, pub has_timezone: bool,
pub timezone: *const String, }
impl JITDataReference {
pub fn box_data_ref(data_ref: Box<JITDataReference>) -> u64 {
use crate::ffi::jit_kinds::jit_box;
use crate::ffi::value_ffi::HK_DATA_REFERENCE;
jit_box(HK_DATA_REFERENCE, *data_ref)
}
}
#[repr(C)]
#[derive(Debug, Clone)]
pub struct JITContext {
pub in_position: bool,
pub position_side: i8, pub entry_price: u64, pub unrealized_pnl_pct: u64,
pub timestamps_ptr: *const i64,
pub column_ptrs: *const *const f64,
pub column_count: usize,
pub row_count: usize,
pub current_row: usize,
pub locals: [u64; 256],
pub stack: [u64; 512],
pub stack_kinds: [u8; 512],
pub stack_ptr: usize,
pub heap_ptr: *mut std::ffi::c_void,
pub function_table: *const JittedStrategyFn,
pub function_table_len: usize,
pub exec_context_ptr: *mut std::ffi::c_void,
pub function_names_ptr: *const String,
pub function_names_len: usize,
pub event_queue_ptr: *mut std::ffi::c_void,
pub suspension_state: u32,
pub iterations_since_yield: u64,
pub yield_threshold: u64,
pub alert_pipeline_ptr: *mut std::ffi::c_void,
pub simulation_mode: u32,
pub simulation_state_ptr: *mut u8,
pub simulation_state_size: usize,
pub gc_safepoint_flag_ptr: *const u8,
pub gc_heap_ptr: *mut std::ffi::c_void,
pub foreign_bridge_ptr: *const std::ffi::c_void,
pub return_type_tag: u8,
pub pending_call_error: u8,
}
impl Default for JITContext {
fn default() -> Self {
Self {
in_position: false,
position_side: 0,
entry_price: box_number(0.0),
unrealized_pnl_pct: box_number(0.0),
timestamps_ptr: std::ptr::null(),
column_ptrs: std::ptr::null(),
column_count: 0,
row_count: 0,
current_row: 0,
locals: [TAG_NULL; 256],
stack: [TAG_NULL; 512],
stack_kinds: [crate::ffi::stack_kind_code::SENTINEL; 512],
stack_ptr: 0,
heap_ptr: std::ptr::null_mut(),
function_table: std::ptr::null(),
function_table_len: 0,
exec_context_ptr: std::ptr::null_mut(),
function_names_ptr: std::ptr::null(),
function_names_len: 0,
event_queue_ptr: std::ptr::null_mut(),
suspension_state: 0,
iterations_since_yield: 0,
yield_threshold: 0, alert_pipeline_ptr: std::ptr::null_mut(),
simulation_mode: 0,
simulation_state_ptr: std::ptr::null_mut(),
simulation_state_size: 0,
gc_safepoint_flag_ptr: std::ptr::null(),
gc_heap_ptr: std::ptr::null_mut(),
foreign_bridge_ptr: std::ptr::null(),
return_type_tag: 0,
pending_call_error: 0,
}
}
}
impl JITContext {
pub fn get_column_value(&self, column_index: usize, offset: i32) -> f64 {
if self.column_ptrs.is_null() || column_index >= self.column_count {
return 0.0;
}
let row_idx = (self.current_row as i32 + offset) as usize;
if row_idx < self.row_count {
unsafe {
let col_ptr = *self.column_ptrs.add(column_index);
if !col_ptr.is_null() {
*col_ptr.add(row_idx)
} else {
0.0
}
}
} else {
0.0
}
}
#[inline]
pub fn set_current_row(&mut self, index: usize) {
self.current_row = index;
}
#[inline]
pub fn update_current_row(&mut self, index: usize) {
self.current_row = index;
}
#[inline]
pub fn is_simulation_mode(&self) -> bool {
self.simulation_mode > 0
}
pub fn setup_simulation(
&mut self,
state_ptr: *mut u8,
state_size: usize,
column_ptrs: *const *const f64,
column_count: usize,
row_count: usize,
timestamps: *const i64,
) {
self.simulation_mode = 1; self.simulation_state_ptr = state_ptr;
self.simulation_state_size = state_size;
self.column_ptrs = column_ptrs;
self.column_count = column_count;
self.row_count = row_count;
self.current_row = 0;
self.timestamps_ptr = timestamps;
}
#[inline]
pub unsafe fn get_state_field_f64(&self, offset: usize) -> f64 {
if self.simulation_state_ptr.is_null() {
return 0.0;
}
let field_ptr = unsafe { self.simulation_state_ptr.add(8 + offset) } as *const u64;
let bits = unsafe { *field_ptr };
unbox_number(bits)
}
#[inline]
pub unsafe fn set_state_field_f64(&mut self, offset: usize, value: f64) {
if self.simulation_state_ptr.is_null() {
return;
}
let field_ptr = unsafe { self.simulation_state_ptr.add(8 + offset) } as *mut u64;
unsafe { *field_ptr = box_number(value) };
}
pub fn clear_simulation(&mut self) {
self.simulation_mode = 0;
self.simulation_state_ptr = std::ptr::null_mut();
self.simulation_state_size = 0;
}
}
pub struct JITDataFrame {
pub columns: Vec<Vec<f64>>,
pub column_ptrs: Vec<*const f64>,
pub timestamps: Vec<i64>,
pub row_count: usize,
}
impl JITDataFrame {
pub fn new() -> Self {
Self {
columns: Vec::new(),
column_ptrs: Vec::new(),
timestamps: Vec::new(),
row_count: 0,
}
}
pub fn from_execution_context(
ctx: &shape_runtime::context::ExecutionContext,
schema: &shape_vm::bytecode::DataFrameSchema,
) -> Self {
let mut data = Self::new();
let _ = (ctx, schema); for _ in 0..schema.column_names.len() {
data.columns.push(Vec::new());
data.column_ptrs.push(std::ptr::null());
}
data
}
pub fn populate_context(&self, ctx: &mut JITContext) {
if !self.column_ptrs.is_empty() {
ctx.column_ptrs = self.column_ptrs.as_ptr();
ctx.column_count = self.column_ptrs.len();
}
ctx.row_count = self.row_count;
if !self.timestamps.is_empty() {
ctx.timestamps_ptr = self.timestamps.as_ptr();
}
}
pub fn len(&self) -> usize {
self.row_count
}
pub fn is_empty(&self) -> bool {
self.row_count == 0
}
pub fn column_count(&self) -> usize {
self.columns.len()
}
pub fn from_datatable(dt: &shape_value::DataTable) -> Self {
use arrow_array::cast::AsArray;
use arrow_schema::{DataType, TimeUnit};
let batch = dt.inner();
let schema = batch.schema();
let num_rows = batch.num_rows();
let mut columns = Vec::new();
let mut timestamps = Vec::new();
for (i, field) in schema.fields().iter().enumerate() {
match field.data_type() {
DataType::Float64 => {
let arr = batch
.column(i)
.as_primitive::<arrow_array::types::Float64Type>();
let col: Vec<f64> = (0..num_rows).map(|r| arr.value(r)).collect();
columns.push(col);
}
DataType::Timestamp(TimeUnit::Microsecond, _) => {
let arr = batch
.column(i)
.as_primitive::<arrow_array::types::TimestampMicrosecondType>();
timestamps = (0..num_rows).map(|r| arr.value(r)).collect();
}
DataType::Int64 => {
let arr = batch
.column(i)
.as_primitive::<arrow_array::types::Int64Type>();
let col: Vec<f64> = (0..num_rows).map(|r| arr.value(r) as f64).collect();
columns.push(col);
}
_ => {
}
}
}
let column_ptrs: Vec<*const f64> = columns.iter().map(|c| c.as_ptr()).collect();
Self {
columns,
column_ptrs,
timestamps,
row_count: num_rows,
}
}
}
impl Default for JITDataFrame {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct JITConfig {
pub opt_level: u8,
pub debug_symbols: bool,
pub jit_threshold: usize,
}
impl Default for JITConfig {
fn default() -> Self {
Self {
opt_level: 3,
debug_symbols: false,
jit_threshold: 100,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_closure_dynamic_captures_0() {
let closure = JITClosure::new(42, &[]);
assert_eq!(closure.function_id, 42);
assert_eq!(closure.captures_count, 0);
let mut closure = closure;
unsafe { closure.drop_captures() };
}
#[test]
fn test_closure_dynamic_captures_5() {
let captures = [
box_number(1.0),
box_number(2.0),
box_number(3.0),
TAG_BOOL_TRUE,
TAG_NULL,
];
let closure = JITClosure::new(7, &captures);
assert_eq!(closure.function_id, 7);
assert_eq!(closure.captures_count, 5);
unsafe {
assert_eq!(unbox_number(closure.get_capture(0)), 1.0);
assert_eq!(unbox_number(closure.get_capture(1)), 2.0);
assert_eq!(unbox_number(closure.get_capture(2)), 3.0);
assert_eq!(closure.get_capture(3), TAG_BOOL_TRUE);
assert_eq!(closure.get_capture(4), TAG_NULL);
}
}
#[test]
fn test_closure_dynamic_captures_20() {
let captures: Vec<u64> = (0..20).map(|i| box_number(i as f64)).collect();
let closure = JITClosure::new(99, &captures);
assert_eq!(closure.captures_count, 20);
unsafe {
for i in 0..20 {
assert_eq!(unbox_number(closure.get_capture(i)), i as f64);
}
}
}
#[test]
fn test_closure_dynamic_captures_64() {
let captures: Vec<u64> = (0..64).map(|i| box_number(i as f64 * 10.0)).collect();
let closure = JITClosure::new(1, &captures);
assert_eq!(closure.captures_count, 64);
unsafe {
for i in 0..64 {
assert_eq!(unbox_number(closure.get_capture(i)), i as f64 * 10.0);
}
}
}
#[test]
fn test_closure_captures_drop() {
let captures: Vec<u64> = (0..32).map(|i| box_number(i as f64)).collect();
let mut closure = JITClosure::new(5, &captures);
assert_eq!(closure.captures_count, 32);
unsafe {
assert_eq!(unbox_number(closure.get_capture(0)), 0.0);
assert_eq!(unbox_number(closure.get_capture(31)), 31.0);
}
unsafe { closure.drop_captures() };
assert!(closure.captures_ptr.is_null());
assert_eq!(closure.captures_count, 32); }
#[test]
fn test_closure_jit_box_roundtrip_via_heap_kind_prefix() {
let captures = [box_number(42.0), TAG_BOOL_FALSE];
let closure = JITClosure::new(10, &captures);
let bits = jit_box(HK_CLOSURE, *closure);
assert_ne!(bits, 0, "allocation pointer is non-null");
assert_eq!(
unsafe { crate::ffi::jit_kinds::read_heap_kind(bits) },
HK_CLOSURE,
"heap-kind prefix at offset 0 discriminates the JIT allocation"
);
let recovered = unsafe { jit_unbox::<JITClosure>(bits) };
assert_eq!(recovered.function_id, 10);
assert_eq!(recovered.captures_count, 2);
unsafe {
assert_eq!(unbox_number(recovered.get_capture(0)), 42.0);
assert_eq!(recovered.get_capture(1), TAG_BOOL_FALSE);
}
unsafe { crate::ffi::jit_kinds::jit_drop::<JITClosure>(bits) };
}
#[test]
fn test_closure_drop_impl_frees_captures_via_jit_drop() {
let captures: Vec<u64> = (0..24).map(|i| box_number(i as f64)).collect();
let closure = JITClosure::new(3, &captures);
let bits = jit_box(HK_CLOSURE, *closure);
let recovered = unsafe { jit_unbox::<JITClosure>(bits) };
assert_eq!(recovered.captures_count, 24);
unsafe {
assert_eq!(unbox_number(recovered.get_capture(23)), 23.0);
}
unsafe { jit_drop::<JITClosure>(bits) };
}
#[test]
fn test_closure_implicit_drop_on_box() {
let captures: Vec<u64> = (0..10).map(|i| box_number(i as f64)).collect();
let closure = JITClosure::new(1, &captures);
drop(closure);
}
}