mod blocks;
pub mod bounds_elision;
mod conversions;
mod ownership;
mod places;
mod rvalues;
mod statements;
mod terminators;
pub(crate) mod types;
pub(crate) mod v2_array;
pub(crate) mod v2_field;
pub(crate) mod v2_int;
pub(crate) mod v2_refcount;
pub(crate) mod v2_string;
pub(crate) mod v2_typed_map;
#[cfg(all(test, feature = "deep-tests"))]
mod integration_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod v2_array_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod closure_dispatch_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod short_circuit_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod ref_param_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod field_ref_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod array_builder_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod groupby_surface_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod typedarray_ptr_regression_tests;
#[cfg(all(test, feature = "deep-tests"))]
mod jit_array_param_regression_tests;
use cranelift::codegen::ir::{FuncRef, StackSlot};
use cranelift::prelude::*;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use crate::ffi_refs::FFIFuncRefs;
use shape_value::v2::closure_layout::ClosureLayout;
use shape_value::v2::struct_layout::FieldKind;
use shape_value::v2::ConcreteType;
use shape_vm::bytecode::MirFunctionData;
use shape_vm::mir::types::*;
use shape_vm::type_tracking::NativeKind;
#[derive(Debug, Clone)]
pub(crate) struct StackClosureCallInfo {
pub(crate) function_id: u16,
pub(crate) capture_offsets: Vec<i32>,
pub(crate) capture_types: Vec<cranelift::prelude::Type>,
}
pub struct MirToIR<'a, 'b> {
pub(crate) builder: &'a mut FunctionBuilder<'b>,
pub(crate) ctx_ptr: Value,
pub(crate) ffi: FFIFuncRefs,
pub(crate) entry_block: Block,
pub(crate) block_map: HashMap<BasicBlockId, Block>,
pub(crate) locals: HashMap<SlotId, Variable>,
pub(crate) local_types: Vec<LocalTypeInfo>,
pub(crate) slot_kinds: Vec<Option<NativeKind>>,
pub(crate) concrete_types: Vec<ConcreteType>,
pub(crate) next_var: usize,
pub(crate) mir: &'a MirFunction,
pub(crate) mir_data: &'a MirFunctionData,
pub(crate) strings: &'a [String],
pub(crate) function_indices: &'a HashMap<String, u16>,
pub(crate) user_func_refs: HashMap<u16, FuncRef>,
pub(crate) user_func_arities: HashMap<u16, u16>,
pub(crate) ref_stack_slots: HashMap<SlotId, (StackSlot, Type)>,
pub(crate) field_byte_offsets: HashMap<String, u16>,
pub(crate) field_native_kinds: HashMap<String, NativeKind>,
pub(crate) field_array_elem_kinds: HashMap<String, NativeKind>,
pub(crate) non_escaping_closure_slots: HashSet<SlotId>,
pub(crate) stack_closure_slots: HashMap<SlotId, StackSlot>,
pub(crate) stack_closure_call_info:
HashMap<SlotId, StackClosureCallInfo>,
pub(crate) ref_param_slots: HashSet<SlotId>,
pub(crate) closure_function_layouts: HashMap<u16, Arc<ClosureLayout>>,
pub(crate) owned_mutable_capture_slots: HashMap<SlotId, FieldKind>,
pub(crate) shared_capture_slots: HashMap<SlotId, FieldKind>,
pub(crate) shared_local_slots: HashSet<SlotId>,
pub(crate) closure_placeholder_fids: Vec<u16>,
pub(crate) next_closure_placeholder_idx: std::cell::Cell<usize>,
pub(crate) bounds_elision: bounds_elision::BoundsElisionPlan,
pub(crate) monomorphized_method_call_sites:
HashMap<(shape_ast::ast::span::Span, Option<usize>), usize>,
pub(crate) caller_function_id: Option<usize>,
pub(crate) operator_trait_dispatch_sites:
HashMap<shape_ast::ast::span::Span, (String, u16)>,
}
pub struct MirPreflightResult {
pub can_compile: bool,
pub blockers: Vec<String>,
}
pub fn preflight(mir_data: &MirFunctionData) -> MirPreflightResult {
let mut blockers = Vec::new();
for block in &mir_data.mir.blocks {
for stmt in &block.statements {
match &stmt.kind {
StatementKind::Assign(place, rvalue) => {
if !is_simple_place(place) {
blockers.push(format!(
"complex place in assignment at {:?}",
stmt.span
));
}
match rvalue {
Rvalue::TypePatternTest { type_annotation, .. } => {
blockers.push(format!(
"TypePatternTest (W15.2-LANG-5): \
`Pattern::Typed` codegen pending, \
annotation = {:?} at {:?}",
type_annotation, stmt.span
));
}
Rvalue::EnumDiscriminantTest {
enum_name,
variant_name,
..
} => {
blockers.push(format!(
"EnumDiscriminantTest (W15.2-LANG-1): \
user-defined `Pattern::Constructor` codegen \
pending, enum = {:?}, variant = {:?} at {:?}",
enum_name, variant_name, stmt.span
));
}
Rvalue::EnumPayload { variant, .. } => {
blockers.push(format!(
"EnumPayload (R8 W9 G.2 Step 2 Bucket 2): \
`Pattern::Constructor` payload binder \
(`Ok(_)` / `Err(_)` / `Some(_)`) codegen \
has receiver-recovery soundness gap at the \
user-fn return-kind boundary per ADR-006 \
\u{a7}2.7.17; whole-program deopt via W12 \
`[jit-fallback]` routes to the bytecode \
interpreter (which compiles EnumPayload via \
the kind-aware opcode dispatch). \
variant = {:?} at {:?}. Tracked v0.4 per \
`docs/v0.3-close-summary.md` \u{a7}5.16 \
JIT-lowering followup workstream.",
variant, stmt.span
));
}
_ => {}
}
}
StatementKind::Drop(place) => {
if !is_simple_place(place) {
blockers.push(format!("complex place in drop at {:?}", stmt.span));
}
}
StatementKind::TaskBoundary(_, _) => {
}
StatementKind::ClosureCapture { function_id, .. } => {
if function_id.is_none() {
blockers.push("ClosureCapture missing function_id".to_string());
}
}
_ => {}
}
}
match &block.terminator.kind {
TerminatorKind::Goto(_)
| TerminatorKind::SwitchBool { .. }
| TerminatorKind::Return
| TerminatorKind::Unreachable => {}
TerminatorKind::Call { .. } => {
}
}
}
MirPreflightResult {
can_compile: blockers.is_empty(),
blockers,
}
}
fn is_simple_place(place: &Place) -> bool {
match place {
Place::Local(_) => true,
Place::Field(inner, _) | Place::Index(inner, _) => is_simple_place(inner),
Place::Deref(inner) => is_simple_place(inner),
}
}
impl<'a, 'b> MirToIR<'a, 'b> {
pub fn new(
builder: &'a mut FunctionBuilder<'b>,
ctx_ptr: Value,
ffi: FFIFuncRefs,
mir_data: &'a MirFunctionData,
slot_kinds: Vec<Option<NativeKind>>,
strings: &'a [String],
entry_block: Block,
function_indices: &'a HashMap<String, u16>,
user_func_refs: HashMap<u16, FuncRef>,
user_func_arities: HashMap<u16, u16>,
) -> Self {
Self::new_with_concrete_types(
builder,
ctx_ptr,
ffi,
mir_data,
slot_kinds,
Vec::new(),
strings,
entry_block,
function_indices,
user_func_refs,
user_func_arities,
)
}
pub fn new_with_concrete_types(
builder: &'a mut FunctionBuilder<'b>,
ctx_ptr: Value,
ffi: FFIFuncRefs,
mir_data: &'a MirFunctionData,
slot_kinds: Vec<Option<NativeKind>>,
concrete_types: Vec<ConcreteType>,
strings: &'a [String],
entry_block: Block,
function_indices: &'a HashMap<String, u16>,
user_func_refs: HashMap<u16, FuncRef>,
user_func_arities: HashMap<u16, u16>,
) -> Self {
Self::new_with_closure_layouts(
builder,
ctx_ptr,
ffi,
mir_data,
slot_kinds,
concrete_types,
strings,
entry_block,
function_indices,
user_func_refs,
user_func_arities,
HashMap::new(),
)
}
pub fn new_with_closure_layouts(
builder: &'a mut FunctionBuilder<'b>,
ctx_ptr: Value,
ffi: FFIFuncRefs,
mir_data: &'a MirFunctionData,
slot_kinds: Vec<Option<NativeKind>>,
concrete_types: Vec<ConcreteType>,
strings: &'a [String],
entry_block: Block,
function_indices: &'a HashMap<String, u16>,
user_func_refs: HashMap<u16, FuncRef>,
user_func_arities: HashMap<u16, u16>,
closure_function_layouts: HashMap<u16, Arc<ClosureLayout>>,
) -> Self {
let local_types = mir_data.mir.local_types.clone();
let _ = slot_kinds;
let concrete_seed: Vec<Option<NativeKind>> = concrete_types
.iter()
.map(|ct| types::native_kind_from_concrete_type(ct))
.collect();
let slot_kinds = types::infer_slot_kinds_with_concrete(
&mir_data.mir,
&concrete_seed,
&concrete_types,
);
let non_escaping_closure_slots =
mir_data.storage_plan.non_escaping_closure_slots.clone();
use shape_vm::type_tracking::{BindingStorageClass, EscapeStatus};
let param_slot_set: HashSet<SlotId> =
mir_data.mir.param_slots.iter().copied().collect();
let mut shared_local_slots: HashSet<SlotId> = HashSet::new();
for (slot, class) in &mir_data.storage_plan.slot_classes {
if !matches!(class, BindingStorageClass::SharedCow) {
continue;
}
if param_slot_set.contains(slot) {
continue;
}
let is_captured = mir_data
.storage_plan
.slot_semantics
.get(slot)
.map(|sem| matches!(sem.escape_status, EscapeStatus::Captured))
.unwrap_or(false);
if !is_captured {
continue;
}
shared_local_slots.insert(*slot);
}
use shape_value::v2::closure_layout::CaptureKind;
use shape_vm::mir::types::{Operand as MirOperand, Place as MirPlace, StatementKind};
for block in &mir_data.mir.blocks {
for stmt in &block.statements {
let StatementKind::ClosureCapture {
operands,
function_id,
..
} = &stmt.kind
else {
continue;
};
let Some(fid) = *function_id else {
continue;
};
let Some(layout) = closure_function_layouts.get(&fid) else {
continue;
};
for (i, op) in operands.iter().enumerate() {
if i >= layout.capture_count() {
break;
}
if !matches!(layout.capture_storage_kind(i), CaptureKind::Shared) {
continue;
}
let root = match op {
MirOperand::Copy(p)
| MirOperand::Move(p)
| MirOperand::MoveExplicit(p) => match p {
MirPlace::Local(s) => Some(*s),
_ => None,
},
MirOperand::Constant(_) => None,
};
if let Some(slot) = root {
if param_slot_set.contains(&slot) {
continue;
}
shared_local_slots.insert(slot);
}
}
}
}
let closure_placeholder_fids =
scan_closure_placeholder_fids(&mir_data.mir, function_indices);
let field_native_kinds =
types::infer_field_native_kinds(&mir_data.mir, &slot_kinds);
let ref_param_slots: HashSet<SlotId> = mir_data
.mir
.param_slots
.iter()
.zip(mir_data.mir.param_reference_kinds.iter())
.filter_map(|(slot, kind)| kind.as_ref().map(|_| *slot))
.collect();
Self {
builder,
ctx_ptr,
ffi,
entry_block,
block_map: HashMap::new(),
locals: HashMap::new(),
local_types,
slot_kinds,
concrete_types,
next_var: 0,
mir: &mir_data.mir,
mir_data,
strings,
function_indices,
user_func_refs,
user_func_arities,
ref_stack_slots: HashMap::new(),
field_byte_offsets: HashMap::new(),
field_native_kinds,
field_array_elem_kinds: HashMap::new(),
non_escaping_closure_slots,
stack_closure_slots: HashMap::new(),
stack_closure_call_info: HashMap::new(),
closure_function_layouts,
owned_mutable_capture_slots: HashMap::new(),
shared_capture_slots: HashMap::new(),
shared_local_slots,
closure_placeholder_fids,
next_closure_placeholder_idx: std::cell::Cell::new(0),
bounds_elision: bounds_elision::BoundsElisionPlan::default(),
monomorphized_method_call_sites: HashMap::new(),
caller_function_id: None,
operator_trait_dispatch_sites: HashMap::new(),
ref_param_slots,
}
}
pub fn set_operator_trait_dispatch_sites(
&mut self,
sites: HashMap<shape_ast::ast::span::Span, (String, u16)>,
) {
self.operator_trait_dispatch_sites = sites;
}
pub fn set_monomorph_routing_context(
&mut self,
sites: HashMap<(shape_ast::ast::span::Span, Option<usize>), usize>,
caller_function_id: Option<usize>,
) {
self.monomorphized_method_call_sites = sites;
self.caller_function_id = caller_function_id;
}
pub fn set_bounds_elision_plan(&mut self, plan: bounds_elision::BoundsElisionPlan) {
self.bounds_elision = plan;
}
pub fn populate_field_byte_offsets_from_schemas(
&mut self,
registry: &shape_runtime::type_schema::TypeSchemaRegistry,
) {
use shape_runtime::type_schema::FieldType;
let referenced_names: std::collections::HashSet<&str> = self
.mir
.field_name_table
.values()
.map(|s| s.as_str())
.collect();
fn field_type_to_native_kind(ft: &FieldType) -> Option<shape_value::NativeKind> {
use shape_value::NativeKind;
match ft {
FieldType::F64 => Some(NativeKind::Float64),
FieldType::I64 => Some(NativeKind::Int64),
FieldType::Bool => Some(NativeKind::Bool),
FieldType::String => Some(NativeKind::String),
FieldType::I8
| FieldType::U8
| FieldType::I16
| FieldType::U16
| FieldType::I32
| FieldType::U32
| FieldType::U64 => Some(NativeKind::Int64),
FieldType::Timestamp => Some(NativeKind::Int64),
FieldType::Object(_)
| FieldType::Array(_)
| FieldType::Option(_)
| FieldType::Decimal
| FieldType::Any
| FieldType::HashMap { .. }
| FieldType::Set(_) => None,
}
}
fn array_elem_to_native_kind(ft: &FieldType) -> Option<shape_value::NativeKind> {
use shape_value::NativeKind;
let FieldType::Array(elem) = ft else {
return None;
};
match elem.as_ref() {
FieldType::F64 => Some(NativeKind::Float64),
FieldType::I64 | FieldType::Timestamp => Some(NativeKind::Int64),
FieldType::Bool => Some(NativeKind::Bool),
FieldType::I8 | FieldType::U8 => Some(NativeKind::Int8),
FieldType::I16 | FieldType::U16 => Some(NativeKind::Int16),
FieldType::I32 | FieldType::U32 => Some(NativeKind::Int32),
FieldType::U64 => Some(NativeKind::UInt64),
FieldType::String
| FieldType::Decimal
| FieldType::Object(_)
| FieldType::Array(_)
| FieldType::Option(_)
| FieldType::Any
| FieldType::HashMap { .. }
| FieldType::Set(_) => None,
}
}
for type_name in registry.type_names().collect::<Vec<_>>() {
let Some(schema) = registry.get(type_name) else {
continue;
};
for (i, field) in schema.fields.iter().enumerate() {
if !referenced_names.contains(field.name.as_str()) {
continue;
}
let byte_off = (i as u16) * 8;
self.field_byte_offsets
.entry(field.name.clone())
.or_insert(byte_off);
if let Some(kind) = field_type_to_native_kind(&field.field_type) {
self.field_native_kinds
.entry(field.name.clone())
.or_insert(kind);
}
if let Some(elem_kind) = array_elem_to_native_kind(&field.field_type) {
self.field_array_elem_kinds
.entry(field.name.clone())
.or_insert(elem_kind);
}
}
}
}
pub fn register_owned_mutable_capture_slots(
&mut self,
captures_count: u16,
layout: &ClosureLayout,
) {
use shape_value::v2::closure_layout::CaptureKind;
let captures_count = captures_count as usize;
if captures_count == 0 {
return;
}
let len = captures_count.min(layout.capture_kinds.len());
for (i, ¶m_slot) in self
.mir
.param_slots
.iter()
.take(len)
.enumerate()
{
let capture_kind = layout.capture_storage_kind(i);
let is_cell_capture = matches!(
capture_kind,
CaptureKind::OwnedMutable | CaptureKind::Shared
);
if !is_cell_capture {
continue;
}
let inner_kind = layout.capture_inner_kind(i);
match capture_kind {
CaptureKind::OwnedMutable => {
self.owned_mutable_capture_slots
.insert(param_slot, inner_kind);
}
CaptureKind::Shared => {
self.shared_capture_slots.insert(param_slot, inner_kind);
}
CaptureKind::Immutable => unreachable!(),
}
if let Some(concrete) = layout.capture_types.get(i) {
if let Some(kind) = types::elem_slot_kind_for_concrete(concrete) {
let idx = param_slot.0 as usize;
if idx < self.slot_kinds.len() && self.slot_kinds[idx].is_none() {
self.slot_kinds[idx] = Some(kind);
}
}
}
}
}
pub fn compile(&mut self) -> Result<(), String> {
self.create_blocks();
self.declare_locals();
self.initialize_shared_local_slots();
self.compile_body()
}
pub fn compile_body(&mut self) -> Result<(), String> {
if tracing::enabled!(target: "shape_jit::mir", tracing::Level::TRACE) {
for (bi, block) in self.mir.blocks.iter().enumerate() {
tracing::trace!(
target: "shape_jit::mir",
bb = bi,
stmts = block.statements.len(),
term = ?block.terminator.kind,
"mir-trace block",
);
for (si, stmt) in block.statements.iter().enumerate() {
tracing::trace!(
target: "shape_jit::mir",
bb = bi,
s = si,
stmt = ?stmt.kind,
"mir-trace statement",
);
}
}
}
for block_idx in 0..self.mir.blocks.len() {
let block = &self.mir.blocks[block_idx];
let cl_block = self.block_map[&block.id];
if block_idx != 0 {
self.builder.switch_to_block(cl_block);
}
for stmt in &block.statements {
self.compile_statement(stmt)?;
}
self.compile_terminator(&block.terminator)?;
}
self.builder.seal_all_blocks();
Ok(())
}
pub(crate) fn reload_referenced_locals(&mut self) {
let refs: Vec<_> = self
.ref_stack_slots
.iter()
.map(|(&slot_id, &(stack_slot, cl_ty))| (slot_id, stack_slot, cl_ty))
.collect();
for (slot_id, stack_slot, cl_ty) in refs {
let reloaded = self.builder.ins().stack_load(cl_ty, stack_slot, 0);
if let Some(&var) = self.locals.get(&slot_id) {
self.builder.def_var(var, reloaded);
}
}
}
}
pub(crate) mod v2_call_abi;
fn scan_closure_placeholder_fids(
mir: &shape_vm::mir::types::MirFunction,
function_indices: &std::collections::HashMap<String, u16>,
) -> Vec<u16> {
use shape_vm::mir::types::{MirConstant, Operand, Rvalue, StatementKind};
let mut result: Vec<u16> = Vec::new();
let mut closure_idx: u32 = 0;
let mut has_capture = false;
for block in &mir.blocks {
for stmt in &block.statements {
let is_placeholder = matches!(
&stmt.kind,
StatementKind::Assign(
_,
Rvalue::Use(Operand::Constant(MirConstant::ClosurePlaceholder))
)
);
if is_placeholder {
if has_capture {
result.push(u16::MAX);
has_capture = false;
} else {
let name = format!("__closure_{}", closure_idx);
let fid = function_indices.get(&name).copied().unwrap_or(u16::MAX);
result.push(fid);
closure_idx = closure_idx.saturating_add(1);
}
continue;
}
if let StatementKind::ClosureCapture {
function_id: Some(_),
..
} = &stmt.kind
{
closure_idx = closure_idx.saturating_add(1);
has_capture = true;
}
}
}
result
}