Int(Value),

Float(Value),

Bool(Value),

None,

Null,

Tuple(Vec<Self>),

FuncRef(FuncRef),

fn from_type_and_value(ty: JitType, val: Value) -> Self {

match ty {

JitType::Int => Self::Int(val),

JitType::Float => Self::Float(val),

JitType::Bool => Self::Bool(val),

}

}

fn to_jit_type(&self) -> Option<JitType> {

match self {

Self::Int(_) => Some(JitType::Int),

Self::Float(_) => Some(JitType::Float),

Self::Bool(_) => Some(JitType::Bool),

Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None,

}

}

fn into_value(self) -> Option<Value> {

match self {

Self::Int(val) | Self::Float(val) | Self::Bool(val) => Some(val),

Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None,

}

}

builder: &'a mut FunctionBuilder<'b>,

stack: Vec<JitValue>,

variables: Box<[Option<Local>]>,

label_to_block: HashMap<Label, Block>,

pub(crate) sig: JitSig,

pub(crate) fn new(

builder: &'a mut FunctionBuilder<'b>,

num_variables: usize,

arg_types: &[JitType],

ret_type: Option<JitType>,

entry_block: Block,

) -> Self {

let mut compiler = Self {

builder,

stack: Vec::new(),

variables: vec![None; num_variables].into_boxed_slice(),

label_to_block: HashMap::new(),

sig: JitSig {

args: arg_types.to_vec(),

ret: ret_type,

},

};

let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();

for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() {

compiler

.store_variable(

(i as u32).into(),

JitValue::from_type_and_value(ty.clone(), val),

)

.unwrap();

}

compiler

}

fn pop_multiple(&mut self, count: usize) -> Vec<JitValue> {

let stack_len = self.stack.len();

self.stack.drain(stack_len - count..).collect()

}

fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {

#[expect(clippy::mut_mut, reason = "This seems like a false positive")]

let builder = &mut self.builder;

let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;

let local = self.variables[idx].get_or_insert_with(|| {

let var = builder.declare_var(ty.to_cranelift());

Local {

var,

ty: ty.clone(),

}

});

if ty != local.ty {

Err(JitCompileError::NotSupported)

} else {

self.builder.def_var(local.var, val.into_value().unwrap());

Ok(())

}

}

fn boolean_val(&mut self, val: JitValue) -> Result<Value, JitCompileError> {

match val {

JitValue::Float(val) => {

let zero = self.builder.ins().f64const(0.0);

let val = self.builder.ins().fcmp(FloatCC::NotEqual, val, zero);

Ok(val)

}

JitValue::Int(val) => {

let zero = self.builder.ins().iconst(types::I64, 0);

let val = self.builder.ins().icmp(IntCC::NotEqual, val, zero);

Ok(val)

}

JitValue::Bool(val) => Ok(val),

JitValue::None => Ok(self.builder.ins().iconst(types::I8, 0)),

JitValue::Null | JitValue::Tuple(_) | JitValue::FuncRef(_) => {

Err(JitCompileError::NotSupported)

}

}

}

fn get_or_create_block(&mut self, label: Label) -> Block {

#[expect(clippy::mut_mut, reason = "This seems like a false positive")]

let builder = &mut self.builder;

*self

.label_to_block

.entry(label)

.or_insert_with(|| builder.create_block())

}

fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result<Label, JitCompileError> {

let after = offset

.checked_add(1)

.and_then(|i| i.checked_add(caches))

.ok_or(JitCompileError::BadBytecode)?;

let target = after

.checked_add(u32::from(arg))

.ok_or(JitCompileError::BadBytecode)?;

Ok(Label::from_u32(target))

}

fn jump_target_backward(

offset: u32,

caches: u32,

arg: OpArg,

) -> Result<Label, JitCompileError> {

let after = offset

.checked_add(1)

.and_then(|i| i.checked_add(caches))

.ok_or(JitCompileError::BadBytecode)?;

let target = after

.checked_sub(u32::from(arg))

.ok_or(JitCompileError::BadBytecode)?;

Ok(Label::from_u32(target))

}

fn instruction_target(

offset: u32,

instruction: Instruction,

arg: OpArg,

) -> Result<Option<Label>, JitCompileError> {

let caches = instruction.cache_entries() as u32;

let target = match instruction {

Instruction::JumpForward { .. } => {

Some(Self::jump_target_forward(offset, caches, arg)?)

}

Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => {

Some(Self::jump_target_backward(offset, caches, arg)?)

}

Instruction::PopJumpIfFalse { .. }

| Instruction::PopJumpIfTrue { .. }

| Instruction::PopJumpIfNone { .. }

| Instruction::PopJumpIfNotNone { .. }

| Instruction::ForIter { .. }

| Instruction::Send { .. } => Some(Self::jump_target_forward(offset, caches, arg)?),

_ => None,

};

Ok(target)

}

pub(crate) fn compile<C: bytecode::Constant>(

&mut self,

func_ref: FuncRef,

bytecode: &CodeObject<C>,

) -> Result<(), JitCompileError> {

// JIT should consume a stable instruction stream: de-specialized opcodes

// with zeroed CACHE entries, not runtime-mutated quickened code.

let clean_instructions: bytecode::CodeUnits = bytecode

.instructions

.original_bytes()

.as_slice()

.try_into()

.map_err(|_| JitCompileError::BadBytecode)?;

let mut label_targets = BTreeSet::new();

let mut target_arg_state = OpArgState::default();

for (offset, &raw_instr) in clean_instructions.iter().enumerate() {

let (instruction, arg) = target_arg_state.get(raw_instr);

if let Some(target) = Self::instruction_target(offset as u32, instruction, arg)? {

label_targets.insert(target);

}

}

let mut arg_state = OpArgState::default();

// Track whether we have "returned" in the current block

let mut in_unreachable_code = false;

for (offset, &raw_instr) in clean_instructions.iter().enumerate() {

let label = Label::from_u32(offset as u32);

let (instruction, arg) = arg_state.get(raw_instr);

// If this is a label that some earlier jump can target,

// treat it as the start of a new reachable block:

if label_targets.contains(&label) {

// Create or get the block for this label:

let target_block = self.get_or_create_block(label);

// If the current block isn't terminated, add a fallthrough jump

if let Some(cur) = self.builder.current_block()

&& cur != target_block

{

// Check if the block needs a terminator by examining the last instruction

let needs_terminator = match self.builder.func.layout.last_inst(cur) {

None => true, // Empty block needs terminator

Some(inst) => {

// Check if the last instruction is a terminator

!self.builder.func.dfg.insts[inst].opcode().is_terminator()

}

};

if needs_terminator {

self.builder.ins().jump(target_block, &[]);

}

}

// Switch to the target block

if self.builder.current_block() != Some(target_block) {

self.builder.switch_to_block(target_block);

}

// We are definitely reachable again at this label

in_unreachable_code = false;

}

// If we're in unreachable code, skip this instruction unless the label re-entered above.

if in_unreachable_code {

continue;

}

// Actually compile this instruction:

self.add_instruction(func_ref, bytecode, offset as u32, instruction, arg)?;

// If that was an unconditional branch or return, mark future instructions unreachable

match instruction {

Instruction::ReturnValue

| Instruction::JumpBackward { .. }

| Instruction::JumpBackwardNoInterrupt { .. }

| Instruction::JumpForward { .. } => {

in_unreachable_code = true;

}

_ => {}

}

}

// After processing, if the current block is unterminated, insert a trap

if let Some(cur) = self.builder.current_block() {

let needs_terminator = match self.builder.func.layout.last_inst(cur) {

None => true,

Some(inst) => !self.builder.func.dfg.insts[inst].opcode().is_terminator(),

};

if needs_terminator {

self.builder.ins().trap(TrapCode::user(0).unwrap());

}

}

Ok(())

}

fn prepare_const<C: bytecode::Constant>(

&mut self,

constant: BorrowedConstant<'_, C>,

) -> Result<JitValue, JitCompileError> {

let value = match constant {

BorrowedConstant::Integer { value } => {

let val = self.builder.ins().iconst(

types::I64,

value.to_i64().ok_or(JitCompileError::NotSupported)?,

);

JitValue::Int(val)

}

BorrowedConstant::Float { value } => {

let val = self.builder.ins().f64const(value);

JitValue::Float(val)

}

BorrowedConstant::Boolean { value } => {

let val = self.builder.ins().iconst(types::I8, value as i64);

JitValue::Bool(val)

}

BorrowedConstant::None => JitValue::None,

_ => return Err(JitCompileError::NotSupported),

};

Ok(value)

}

fn return_value(&mut self, val: JitValue) -> Result<(), JitCompileError> {

if let Some(ref ty) = self.sig.ret {

// If the signature has a return type, enforce it

if val.to_jit_type().as_ref() != Some(ty) {

return Err(JitCompileError::NotSupported);

}

} else {

// First time we see a return, define it in the signature

let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?;

self.sig.ret = Some(ty.clone());

self.builder

.func

.signature

.returns

.push(AbiParam::new(ty.to_cranelift()));

}

// If this is e.g. an Int, Float, or Bool we have a Cranelift `Value`.

// If we have JitValue::None or .Tuple(...) but can't handle that, error out (or handle differently).

let cr_val = val.into_value().ok_or(JitCompileError::NotSupported)?;

self.builder.ins().return_(&[cr_val]);

Ok(())

}

pub(crate) fn add_instruction<C: bytecode::Constant>(

&mut self,

func_ref: FuncRef,

bytecode: &CodeObject<C>,

offset: u32,

instruction: Instruction,

arg: OpArg,

) -> Result<(), JitCompileError> {

match instruction {

Instruction::BinaryOp { op } => {

let op = op.get(arg);

// the rhs is popped off first

let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let a_type = a.to_jit_type();

let b_type = b.to_jit_type();

let val = match (op, a, b) {

(

BinaryOperator::Add | BinaryOperator::InplaceAdd,

JitValue::Int(a),

JitValue::Int(b),

) => {

let (out, carry) = self.builder.ins().sadd_overflow(a, b);

self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW);

JitValue::Int(out)

}

(

BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.compile_sub(a, b)),

(

BinaryOperator::FloorDivide | BinaryOperator::InplaceFloorDivide,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().sdiv(a, b)),

(

BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,

JitValue::Int(a),

JitValue::Int(b),

) => {

// Check if b == 0, If so trap with a division by zero error

self.builder

.ins()

.trapz(b, TrapCode::INTEGER_DIVISION_BY_ZERO);

// Else convert to float and divide

let a_float = self.builder.ins().fcvt_from_sint(types::F64, a);

let b_float = self.builder.ins().fcvt_from_sint(types::F64, b);

JitValue::Float(self.builder.ins().fdiv(a_float, b_float))

}

(

BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().imul(a, b)),

(

BinaryOperator::Remainder | BinaryOperator::InplaceRemainder,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().srem(a, b)),

(

BinaryOperator::Power | BinaryOperator::InplacePower,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.compile_ipow(a, b)),

(

BinaryOperator::Lshift | BinaryOperator::Rshift,

JitValue::Int(a),

JitValue::Int(b),

) => {

// Shifts throw an exception if we have a negative shift count

// Remove all bits except the sign bit, and trap if its 1 (i.e. negative).

let sign = self.builder.ins().ushr_imm(b, 63);

self.builder.ins().trapnz(

sign,

TrapCode::user(CustomTrapCode::NegativeShiftCount as u8).unwrap(),

);

let out =

if matches!(op, BinaryOperator::Lshift | BinaryOperator::InplaceLshift)

{

self.builder.ins().ishl(a, b)

} else {

self.builder.ins().sshr(a, b)

};

JitValue::Int(out)

}

(

BinaryOperator::And | BinaryOperator::InplaceAnd,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().band(a, b)),

(

BinaryOperator::Or | BinaryOperator::InplaceOr,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().bor(a, b)),

(

BinaryOperator::Xor | BinaryOperator::InplaceXor,

JitValue::Int(a),

JitValue::Int(b),

) => JitValue::Int(self.builder.ins().bxor(a, b)),

// Floats

(

BinaryOperator::Add | BinaryOperator::InplaceAdd,

JitValue::Float(a),

JitValue::Float(b),

) => JitValue::Float(self.builder.ins().fadd(a, b)),

(

BinaryOperator::Subtract | BinaryOperator::InplaceSubtract,

JitValue::Float(a),

JitValue::Float(b),

) => JitValue::Float(self.builder.ins().fsub(a, b)),

(

BinaryOperator::Multiply | BinaryOperator::InplaceMultiply,

JitValue::Float(a),

JitValue::Float(b),

) => JitValue::Float(self.builder.ins().fmul(a, b)),

(

BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide,

JitValue::Float(a),

JitValue::Float(b),

) => JitValue::Float(self.builder.ins().fdiv(a, b)),

(

BinaryOperator::Power | BinaryOperator::InplacePower,

JitValue::Float(a),

JitValue::Float(b),

) => JitValue::Float(self.compile_fpow(a, b)),

// Floats and Integers

(_, JitValue::Int(a), JitValue::Float(b))

| (_, JitValue::Float(a), JitValue::Int(b)) => {

let operand_one = match a_type.unwrap() {

JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, a),

_ => a,

};

let operand_two = match b_type.unwrap() {

JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, b),

_ => b,

};

match op {

BinaryOperator::Add | BinaryOperator::InplaceAdd => {

JitValue::Float(self.builder.ins().fadd(operand_one, operand_two))

}

BinaryOperator::Subtract | BinaryOperator::InplaceSubtract => {

JitValue::Float(self.builder.ins().fsub(operand_one, operand_two))

}

BinaryOperator::Multiply | BinaryOperator::InplaceMultiply => {

JitValue::Float(self.builder.ins().fmul(operand_one, operand_two))

}

BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide => {

JitValue::Float(self.builder.ins().fdiv(operand_one, operand_two))

}

BinaryOperator::Power | BinaryOperator::InplacePower => {

JitValue::Float(self.compile_fpow(operand_one, operand_two))

}

_ => return Err(JitCompileError::NotSupported),

}

}

_ => return Err(JitCompileError::NotSupported),

};

self.stack.push(val);

Ok(())

}

Instruction::BuildTuple { count } => {

let elements = self.pop_multiple(count.get(arg) as usize);

self.stack.push(JitValue::Tuple(elements));

Ok(())

}

Instruction::Call { argc } => {

let nargs = argc.get(arg);

let mut args = Vec::new();

for _ in 0..nargs {

let arg = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

args.push(arg.into_value().unwrap());

}

// Pop self_or_null (should be Null for JIT-compiled recursive calls)

let self_or_null = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

if !matches!(self_or_null, JitValue::Null) {

return Err(JitCompileError::NotSupported);

}

match self.stack.pop().ok_or(JitCompileError::BadBytecode)? {

JitValue::FuncRef(reference) => {

let call = self.builder.ins().call(reference, &args);

let returns = self.builder.inst_results(call);

self.stack.push(JitValue::Int(returns[0]));

Ok(())

}

_ => Err(JitCompileError::BadBytecode),

}

}

Instruction::PushNull => {

self.stack.push(JitValue::Null);

Ok(())

}

Instruction::CallIntrinsic1 { func } => {

match func.get(arg) {

IntrinsicFunction1::UnaryPositive => {

match self.stack.pop().ok_or(JitCompileError::BadBytecode)? {

JitValue::Int(val) => {

// Nothing to do

self.stack.push(JitValue::Int(val));

Ok(())

}

_ => Err(JitCompileError::NotSupported),

}

}

_ => Err(JitCompileError::NotSupported),

}

}

Instruction::CompareOp { opname } => {

let op = opname.get(arg);

// the rhs is popped off first

let b = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let a = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let a_type: Option<JitType> = a.to_jit_type();

let b_type: Option<JitType> = b.to_jit_type();

match (a, b) {

(

JitValue::Int(a) | JitValue::Bool(a),

JitValue::Int(b) | JitValue::Bool(b),

) => {

let operand_one = match a_type.unwrap() {

JitType::Bool => self.builder.ins().uextend(types::I64, a),

_ => a,

};

let operand_two = match b_type.unwrap() {

JitType::Bool => self.builder.ins().uextend(types::I64, b),

_ => b,

};

let cond = match op {

ComparisonOperator::Equal => IntCC::Equal,

ComparisonOperator::NotEqual => IntCC::NotEqual,

ComparisonOperator::Less => IntCC::SignedLessThan,

ComparisonOperator::LessOrEqual => IntCC::SignedLessThanOrEqual,

ComparisonOperator::Greater => IntCC::SignedGreaterThan,

ComparisonOperator::GreaterOrEqual => IntCC::SignedGreaterThanOrEqual,

};

let val = self.builder.ins().icmp(cond, operand_one, operand_two);

self.stack.push(JitValue::Bool(val));

Ok(())

}

(JitValue::Float(a), JitValue::Float(b)) => {

let cond = match op {

ComparisonOperator::Equal => FloatCC::Equal,

ComparisonOperator::NotEqual => FloatCC::NotEqual,

ComparisonOperator::Less => FloatCC::LessThan,

ComparisonOperator::LessOrEqual => FloatCC::LessThanOrEqual,

ComparisonOperator::Greater => FloatCC::GreaterThan,

ComparisonOperator::GreaterOrEqual => FloatCC::GreaterThanOrEqual,

};

let val = self.builder.ins().fcmp(cond, a, b);

self.stack.push(JitValue::Bool(val));

Ok(())

}

_ => Err(JitCompileError::NotSupported),

}

}

Instruction::ExtendedArg

| Instruction::Cache

| Instruction::MakeCell { .. }

| Instruction::CopyFreeVars { .. } => Ok(()),

Instruction::JumpBackward { .. }

| Instruction::JumpBackwardNoInterrupt { .. }

| Instruction::JumpForward { .. } => {

let target = Self::instruction_target(offset, instruction, arg)?

.ok_or(JitCompileError::BadBytecode)?;

let target_block = self.get_or_create_block(target);

self.builder.ins().jump(target_block, &[]);

Ok(())

}

Instruction::LoadConst { consti } => {

let val =

self.prepare_const(bytecode.constants[consti.get(arg)].borrow_constant())?;

self.stack.push(val);

Ok(())

}

Instruction::LoadSmallInt { i } => {

let small_int = i.get(arg) as i64;

let val = self.builder.ins().iconst(types::I64, small_int);

self.stack.push(JitValue::Int(val));

Ok(())

}

Instruction::LoadFast { var_num } | Instruction::LoadFastBorrow { var_num } => {

let local = self.variables[var_num.get(arg)]

.as_ref()

.ok_or(JitCompileError::BadBytecode)?;

self.stack.push(JitValue::from_type_and_value(

local.ty.clone(),

self.builder.use_var(local.var),

));

Ok(())

}

Instruction::LoadFastLoadFast { var_nums }

| Instruction::LoadFastBorrowLoadFastBorrow { var_nums } => {

let oparg = var_nums.get(arg);

let (idx1, idx2) = oparg.indexes();

for idx in [idx1, idx2] {

let local = self.variables[idx]

.as_ref()

.ok_or(JitCompileError::BadBytecode)?;

self.stack.push(JitValue::from_type_and_value(

local.ty.clone(),

self.builder.use_var(local.var),

));

}

Ok(())

}

Instruction::LoadGlobal { namei } => {

let oparg = namei.get(arg);

let name = &bytecode.names[(oparg >> 1) as usize];

if name.as_ref() != bytecode.obj_name.as_ref() {

Err(JitCompileError::NotSupported)

} else {

self.stack.push(JitValue::FuncRef(func_ref));

if (oparg & 1) != 0 {

self.stack.push(JitValue::Null);

}

Ok(())

}

}

Instruction::Nop | Instruction::NotTaken => Ok(()),

Instruction::PopJumpIfFalse { .. } => {

let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let val = self.boolean_val(cond)?;

let then_label = Self::instruction_target(offset, instruction, arg)?

.ok_or(JitCompileError::BadBytecode)?;

let then_block = self.get_or_create_block(then_label);

let else_block = self.builder.create_block();

self.builder

.ins()

.brif(val, else_block, &[], then_block, &[]);

self.builder.switch_to_block(else_block);

Ok(())

}

Instruction::PopJumpIfTrue { .. } => {

let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

let val = self.boolean_val(cond)?;

let then_label = Self::instruction_target(offset, instruction, arg)?

.ok_or(JitCompileError::BadBytecode)?;

let then_block = self.get_or_create_block(then_label);

let else_block = self.builder.create_block();

self.builder

.ins()

.brif(val, then_block, &[], else_block, &[]);

self.builder.switch_to_block(else_block);

Ok(())

}

Instruction::PopTop => {

self.stack.pop();

Ok(())

}

Instruction::Resume { .. } => {

// TODO: Implement the resume instruction

Ok(())

}

Instruction::ReturnValue => {

let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

self.return_value(val)

}

Instruction::StoreFast { var_num } => {

let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;

self.store_variable(var_num.get(arg), val)

}

Instruction::StoreFastLoadFast { var_nums } => {

let oparg = var_nums.get(arg);