Skip to content

Commit 49ec351

Browse files
youknowoneyouknowone
committed
Fix jit failures
1 parent 1cb1439 commit 49ec351

File tree

1 file changed

+73
-10
lines changed

1 file changed

+73
-10
lines changed

‎crates/jit/src/instructions.rs‎

Lines changed: 73 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
// spell-checker: disable
22
use super::{JitCompileError, JitSig, JitType};
3+
use alloc::collections::BTreeSet;
34
use cranelift::codegen::ir::FuncRef;
45
use cranelift::prelude::*;
56
use num_traits::cast::ToPrimitive;
@@ -154,12 +155,67 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
154155
.or_insert_with(|| builder.create_block())
155156
}
156157

158+
fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result<Label, JitCompileError> {
159+
let after = offset
160+
.checked_add(1)
161+
.and_then(|i| i.checked_add(caches))
162+
.ok_or(JitCompileError::BadBytecode)?;
163+
let target = after
164+
.checked_add(u32::from(arg))
165+
.ok_or(JitCompileError::BadBytecode)?;
166+
Ok(Label(target))
167+
}
168+
169+
fn jump_target_backward(
170+
offset: u32,
171+
caches: u32,
172+
arg: OpArg,
173+
) -> Result<Label, JitCompileError> {
174+
let after = offset
175+
.checked_add(1)
176+
.and_then(|i| i.checked_add(caches))
177+
.ok_or(JitCompileError::BadBytecode)?;
178+
let target = after
179+
.checked_sub(u32::from(arg))
180+
.ok_or(JitCompileError::BadBytecode)?;
181+
Ok(Label(target))
182+
}
183+
184+
fn instruction_target(
185+
offset: u32,
186+
instruction: Instruction,
187+
arg: OpArg,
188+
) -> Result<Option<Label>, JitCompileError> {
189+
let target = match instruction {
190+
Instruction::JumpForward { .. } => Some(Self::jump_target_forward(offset, 0, arg)?),
191+
Instruction::JumpBackward { .. } => Some(Self::jump_target_backward(offset, 1, arg)?),
192+
Instruction::JumpBackwardNoInterrupt { .. } => {
193+
Some(Self::jump_target_backward(offset, 0, arg)?)
194+
}
195+
Instruction::PopJumpIfFalse { .. }
196+
| Instruction::PopJumpIfTrue { .. }
197+
| Instruction::PopJumpIfNone { .. }
198+
| Instruction::PopJumpIfNotNone { .. }
199+
| Instruction::ForIter { .. }
200+
| Instruction::Send { .. } => Some(Self::jump_target_forward(offset, 1, arg)?),
201+
_ => None,
202+
};
203+
Ok(target)
204+
}
205+
157206
pub fn compile<C: bytecode::Constant>(
158207
&mut self,
159208
func_ref: FuncRef,
160209
bytecode: &CodeObject<C>,
161210
) -> Result<(), JitCompileError> {
162-
let label_targets = bytecode.label_targets();
211+
let mut label_targets = BTreeSet::new();
212+
let mut target_arg_state = OpArgState::default();
213+
for (offset, &raw_instr) in bytecode.instructions.iter().enumerate() {
214+
let (instruction, arg) = target_arg_state.get(raw_instr);
215+
if let Some(target) = Self::instruction_target(offset as u32, instruction, arg)? {
216+
label_targets.insert(target);
217+
}
218+
}
163219
let mut arg_state = OpArgState::default();
164220

165221
// Track whether we have "returned" in the current block
@@ -206,7 +262,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
206262
}
207263

208264
// Actually compile this instruction:
209-
self.add_instruction(func_ref, bytecode, instruction, arg)?;
265+
self.add_instruction(func_ref, bytecode, offset as u32, instruction, arg)?;
210266

211267
// If that was an unconditional branch or return, mark future instructions unreachable
212268
match instruction {
@@ -288,6 +344,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
288344
&mut self,
289345
func_ref: FuncRef,
290346
bytecode: &CodeObject<C>,
347+
offset: u32,
291348
instruction: Instruction,
292349
arg: OpArg,
293350
) -> Result<(), JitCompileError> {
@@ -559,10 +616,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
559616
}
560617
Instruction::ExtendedArg | Instruction::Cache => Ok(()),
561618

562-
Instruction::JumpBackward { target }
563-
| Instruction::JumpBackwardNoInterrupt { target }
564-
| Instruction::JumpForward { target } => {
565-
let target_block = self.get_or_create_block(target.get(arg));
619+
Instruction::JumpBackward { .. }
620+
| Instruction::JumpBackwardNoInterrupt { .. }
621+
| Instruction::JumpForward { .. } => {
622+
let target = Self::instruction_target(offset, instruction, arg)?
623+
.ok_or(JitCompileError::BadBytecode)?;
624+
let target_block = self.get_or_create_block(target);
566625
self.builder.ins().jump(target_block, &[]);
567626
Ok(())
568627
}
@@ -615,10 +674,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
615674
}
616675
}
617676
Instruction::Nop | Instruction::NotTaken => Ok(()),
618-
Instruction::PopJumpIfFalse { target } => {
677+
Instruction::PopJumpIfFalse { .. } => {
619678
let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
620679
let val = self.boolean_val(cond)?;
621-
let then_block = self.get_or_create_block(target.get(arg));
680+
let then_label = Self::instruction_target(offset, instruction, arg)?
681+
.ok_or(JitCompileError::BadBytecode)?;
682+
let then_block = self.get_or_create_block(then_label);
622683
let else_block = self.builder.create_block();
623684

624685
self.builder
@@ -628,10 +689,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
628689

629690
Ok(())
630691
}
631-
Instruction::PopJumpIfTrue { target } => {
692+
Instruction::PopJumpIfTrue { .. } => {
632693
let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
633694
let val = self.boolean_val(cond)?;
634-
let then_block = self.get_or_create_block(target.get(arg));
695+
let then_label = Self::instruction_target(offset, instruction, arg)?
696+
.ok_or(JitCompileError::BadBytecode)?;
697+
let then_block = self.get_or_create_block(then_label);
635698
let else_block = self.builder.create_block();
636699

637700
self.builder

0 commit comments

Comments
 (0)