Skip to content

Commit 3ebdeb7

Browse files
committed
Update rustpython_jit to work with codeobj.constants vec
1 parent 5176864 commit 3ebdeb7

3 files changed

Lines changed: 39 additions & 25 deletions

File tree

jit/src/instructions.rs

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
use cranelift::prelude::*;
22
use num_traits::cast::ToPrimitive;
33
use rustpython_bytecode::bytecode::{
4-
BinaryOperator, CodeObject, ComparisonOperator, Constant, Instruction, Label, NameScope,
5-
UnaryOperator,
4+
self, BinaryOperator, BorrowedConstant, CodeObject, ComparisonOperator, Instruction, Label,
5+
NameScope, UnaryOperator,
66
};
77
use std::collections::HashMap;
88

@@ -105,7 +105,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
105105
.or_insert_with(|| builder.create_block())
106106
}
107107

108-
pub fn compile(&mut self, bytecode: &CodeObject) -> Result<(), JitCompileError> {
108+
pub fn compile<C: bytecode::Constant>(
109+
&mut self,
110+
bytecode: &CodeObject<C>,
111+
) -> Result<(), JitCompileError> {
109112
let offset_to_label: HashMap<&usize, &Label> =
110113
bytecode.label_map.iter().map(|(k, v)| (v, k)).collect();
111114

@@ -128,15 +131,18 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
128131
continue;
129132
}
130133

131-
self.add_instruction(&instruction)?;
134+
self.add_instruction(&instruction, &bytecode.constants)?;
132135
}
133136

134137
Ok(())
135138
}
136139

137-
fn load_const(&mut self, constant: &Constant) -> Result<(), JitCompileError> {
140+
fn load_const<C: bytecode::Constant>(
141+
&mut self,
142+
constant: BorrowedConstant<C>,
143+
) -> Result<(), JitCompileError> {
138144
match constant {
139-
Constant::Integer { value } => {
145+
BorrowedConstant::Integer { value } => {
140146
let val = self.builder.ins().iconst(
141147
types::I64,
142148
value.to_i64().ok_or(JitCompileError::NotSupported)?,
@@ -147,16 +153,16 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
147153
});
148154
Ok(())
149155
}
150-
Constant::Float { value } => {
151-
let val = self.builder.ins().f64const(*value);
156+
BorrowedConstant::Float { value } => {
157+
let val = self.builder.ins().f64const(value);
152158
self.stack.push(JitValue {
153159
val,
154160
ty: JitType::Float,
155161
});
156162
Ok(())
157163
}
158-
Constant::Boolean { value } => {
159-
let val = self.builder.ins().iconst(types::I8, *value as i64);
164+
BorrowedConstant::Boolean { value } => {
165+
let val = self.builder.ins().iconst(types::I8, value as i64);
160166
self.stack.push(JitValue {
161167
val,
162168
ty: JitType::Bool,
@@ -167,7 +173,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
167173
}
168174
}
169175

170-
fn add_instruction(&mut self, instruction: &Instruction) -> Result<(), JitCompileError> {
176+
pub fn add_instruction<C: bytecode::Constant>(
177+
&mut self,
178+
instruction: &Instruction,
179+
constants: &[C],
180+
) -> Result<(), JitCompileError> {
171181
match instruction {
172182
Instruction::JumpIfFalse { target } => {
173183
let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
@@ -222,7 +232,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
222232
let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
223233
self.store_variable(name.clone(), val)
224234
}
225-
Instruction::LoadConst { value } => self.load_const(value),
235+
Instruction::LoadConst { idx } => self.load_const(constants[*idx].borrow_constant()),
226236
Instruction::ReturnValue => {
227237
let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
228238
if let Some(ref ty) = self.sig.ret {

jit/src/lib.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,9 +46,9 @@ impl Jit {
4646
}
4747
}
4848

49-
fn build_function(
49+
fn build_function<C: bytecode::Constant>(
5050
&mut self,
51-
bytecode: &bytecode::CodeObject,
51+
bytecode: &bytecode::CodeObject<C>,
5252
args: &[JitType],
5353
) -> Result<(FuncId, JitSig), JitCompileError> {
5454
for arg in args {
@@ -92,8 +92,8 @@ impl Jit {
9292
}
9393
}
9494

95-
pub fn compile(
96-
bytecode: &bytecode::CodeObject,
95+
pub fn compile<C: bytecode::Constant>(
96+
bytecode: &bytecode::CodeObject<C>,
9797
args: &[JitType],
9898
) -> Result<CompiledCode, JitCompileError> {
9999
let mut jit = Jit::new();

jit/tests/common.rs

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use std::collections::HashMap;
22

3-
use rustpython_bytecode::bytecode::{CodeObject, Constant, Instruction, NameScope};
3+
use rustpython_bytecode::bytecode::{CodeObject, ConstantData, Instruction, NameScope};
44
use rustpython_jit::{CompiledCode, JitType};
55

66
#[derive(Debug, Clone)]
@@ -39,12 +39,12 @@ enum StackValue {
3939
Function(Function),
4040
}
4141

42-
impl From<Constant> for StackValue {
43-
fn from(value: Constant) -> Self {
42+
impl From<ConstantData> for StackValue {
43+
fn from(value: ConstantData) -> Self {
4444
match value {
45-
Constant::String { value } => StackValue::String(value),
46-
Constant::None => StackValue::None,
47-
Constant::Code { code } => StackValue::Code(code),
45+
ConstantData::Str { value } => StackValue::String(value),
46+
ConstantData::None => StackValue::None,
47+
ConstantData::Code { code } => StackValue::Code(code),
4848
c => unimplemented!("constant {:?} isn't yet supported in py_function!", c),
4949
}
5050
}
@@ -65,15 +65,19 @@ impl StackMachine {
6565

6666
pub fn run(&mut self, code: CodeObject) {
6767
for instruction in code.instructions {
68-
if self.process_instruction(instruction) {
68+
if self.process_instruction(instruction, &code.constants) {
6969
break;
7070
}
7171
}
7272
}
7373

74-
fn process_instruction(&mut self, instruction: Instruction) -> bool {
74+
fn process_instruction(
75+
&mut self,
76+
instruction: Instruction,
77+
constants: &[ConstantData],
78+
) -> bool {
7579
match instruction {
76-
Instruction::LoadConst { value } => self.stack.push(value.into()),
80+
Instruction::LoadConst { idx } => self.stack.push(constants[idx].clone().into()),
7781
Instruction::LoadName {
7882
name,
7983
scope: NameScope::Free,

0 commit comments

Comments
 (0)