Skip to content

Commit b79d7ed

Browse files
committed
Fix not type checking arguments to jit function.
1 parent bd3c4fc commit b79d7ed

6 files changed

Lines changed: 64 additions & 21 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎jit/src/lib.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,12 @@ pub enum JitCompileError {
2020
CraneliftError(#[from] ModuleError),
2121
}
2222

23+
#[derive(Debug, thiserror::Error)]
24+
pub enum JitArgumentError {
25+
#[error("argument is of wrong type")]
26+
ArgumentTypeMismatch,
27+
}
28+
2329
struct Jit {
2430
builder_context: FunctionBuilderContext,
2531
ctx: codegen::Context,
@@ -215,8 +221,14 @@ impl<'a> ArgsBuilder<'a> {
215221
}
216222
}
217223

218-
pub fn set(&mut self, idx: usize, value: AbiValue) {
219-
self.values[idx] = Some(value);
224+
pub fn set(&mut self, idx: usize, value: AbiValue) -> Result<(), JitArgumentError> {
225+
match (&self.code.sig.args[idx], &value) {
226+
(JitType::Int, AbiValue::Int(_)) | (JitType::Float, AbiValue::Float(_)) => {
227+
self.values[idx] = Some(value);
228+
Ok(())
229+
}
230+
_ => Err(JitArgumentError::ArgumentTypeMismatch),
231+
}
220232
}
221233

222234
pub fn is_set(&self, idx: usize) -> bool {

‎tests/snippets/jit.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ def tests():
1717
assert foo() == 15
1818
assert bar() == 2e5
1919
assert baz(17, 20) == 49
20+
assert baz(17, 22.5) == 51.5
2021

2122

2223
tests()

‎vm/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ parking_lot = "0.11"
7777
thread_local = "1.0"
7878
cfg-if = "0.1.10"
7979
timsort = "0.1"
80+
thiserror = "1.0"
8081

8182
## unicode stuff
8283
unicode_names2 = "0.4"

‎vm/src/obj/objfunction.rs‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -238,8 +238,15 @@ impl PyFunction {
238238
) -> PyResult {
239239
#[cfg(feature = "jit")]
240240
if let Some(jitted_code) = self.jitted_code.get() {
241-
if let Some(args) = jitfunc::get_jit_args(self, &func_args, jitted_code, vm) {
242-
return Ok(jitted_code.invoke(&args).into_pyobject(vm));
241+
match jitfunc::get_jit_args(self, &func_args, jitted_code, vm) {
242+
Ok(args) => {
243+
return Ok(jitted_code.invoke(&args).into_pyobject(vm));
244+
}
245+
Err(err) => info!(
246+
"jit: function `{}` is falling back to being interpreted because of the \
247+
error: {}",
248+
self.code.obj_name, err
249+
),
243250
}
244251
}
245252

‎vm/src/obj/objfunction/jitfunc.rs‎

Lines changed: 38 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,25 @@ use crate::pyobject::{
1010
use crate::VirtualMachine;
1111
use num_traits::ToPrimitive;
1212
use rustpython_bytecode::bytecode::CodeFlags;
13-
use rustpython_jit::{AbiValue, Args, CompiledCode, JitType};
13+
use rustpython_jit::{AbiValue, Args, CompiledCode, JitArgumentError, JitType};
14+
15+
#[derive(Debug, thiserror::Error)]
16+
pub enum ArgsError {
17+
#[error("wrong number of arguments passed")]
18+
WrongNumberOfArgs,
19+
#[error("argument passed multiple times")]
20+
ArgPassedMultipleTimes,
21+
#[error("not a keyword argument")]
22+
NotAKeywordArg,
23+
#[error("not all arguments passed")]
24+
NotAllArgsPassed,
25+
#[error("integer can't fit into a machine integer")]
26+
IntOverflow,
27+
#[error("type can't be used in a jit function")]
28+
NonJitType,
29+
#[error("{0}")]
30+
JitError(#[from] JitArgumentError),
31+
}
1432

1533
impl IntoPyObject for AbiValue {
1634
fn into_pyobject(self, vm: &VirtualMachine) -> PyObjectRef {
@@ -85,15 +103,18 @@ pub fn get_jit_arg_types(func: &PyFunctionRef, vm: &VirtualMachine) -> PyResult<
85103
}
86104
}
87105

88-
fn get_jit_value(vm: &VirtualMachine, obj: &PyObjectRef) -> Option<AbiValue> {
106+
fn get_jit_value(vm: &VirtualMachine, obj: &PyObjectRef) -> Result<AbiValue, ArgsError> {
89107
// This does exact type checks as subclasses of int/float can't be passed to jitted functions
90108
let cls = obj.lease_class();
91109
if cls.is(&vm.ctx.types.int_type) {
92-
objint::get_value(&obj).to_i64().map(AbiValue::Int)
110+
objint::get_value(&obj)
111+
.to_i64()
112+
.map(AbiValue::Int)
113+
.ok_or(ArgsError::IntOverflow)
93114
} else if cls.is(&vm.ctx.types.float_type) {
94-
Some(AbiValue::Float(objfloat::get_value(&obj)))
115+
Ok(AbiValue::Float(objfloat::get_value(&obj)))
95116
} else {
96-
None
117+
Err(ArgsError::NonJitType)
97118
}
98119
}
99120

@@ -107,35 +128,35 @@ pub(crate) fn get_jit_args<'a>(
107128
func_args: &PyFuncArgs,
108129
jitted_code: &'a CompiledCode,
109130
vm: &VirtualMachine,
110-
) -> Option<Args<'a>> {
131+
) -> Result<Args<'a>, ArgsError> {
111132
let mut jit_args = jitted_code.args_builder();
112133
let nargs = func_args.args.len();
113134

114135
if nargs > func.code.arg_names.len() || nargs < func.code.posonlyarg_count {
115-
return None;
136+
return Err(ArgsError::WrongNumberOfArgs);
116137
}
117138

118139
// Add positional arguments
119140
for i in 0..nargs {
120-
jit_args.set(i, get_jit_value(vm, &func_args.args[i])?);
141+
jit_args.set(i, get_jit_value(vm, &func_args.args[i])?)?;
121142
}
122143

123144
// Handle keyword arguments
124145
for (name, value) in &func_args.kwargs {
125146
if let Some(arg_idx) = func.code.arg_names.iter().position(|arg| arg == name) {
126147
if jit_args.is_set(arg_idx) {
127-
return None;
148+
return Err(ArgsError::ArgPassedMultipleTimes);
128149
}
129-
jit_args.set(arg_idx, get_jit_value(vm, &value)?);
150+
jit_args.set(arg_idx, get_jit_value(vm, &value)?)?;
130151
} else if let Some(kwarg_idx) = func.code.kwonlyarg_names.iter().position(|arg| arg == name)
131152
{
132153
let arg_idx = kwarg_idx + func.code.arg_names.len();
133154
if jit_args.is_set(arg_idx) {
134-
return None;
155+
return Err(ArgsError::ArgPassedMultipleTimes);
135156
}
136-
jit_args.set(arg_idx, get_jit_value(vm, &value)?);
157+
jit_args.set(arg_idx, get_jit_value(vm, &value)?)?;
137158
} else {
138-
return None;
159+
return Err(ArgsError::NotAKeywordArg);
139160
}
140161
}
141162

@@ -145,7 +166,7 @@ pub(crate) fn get_jit_args<'a>(
145166
for (i, default) in defaults.iter().enumerate() {
146167
let arg_idx = i + func.code.arg_names.len() - defaults.len();
147168
if !jit_args.is_set(arg_idx) {
148-
jit_args.set(arg_idx, get_jit_value(vm, default)?);
169+
jit_args.set(arg_idx, get_jit_value(vm, default)?)?;
149170
}
150171
}
151172
}
@@ -157,12 +178,12 @@ pub(crate) fn get_jit_args<'a>(
157178
if !jit_args.is_set(arg_idx) {
158179
let default = kw_only_defaults
159180
.get_item(name.as_str(), vm)
160-
.ok()
181+
.map_err(|_| ArgsError::NotAllArgsPassed)
161182
.and_then(|obj| get_jit_value(vm, &obj))?;
162-
jit_args.set(arg_idx, default);
183+
jit_args.set(arg_idx, default)?;
163184
}
164185
}
165186
}
166187

167-
jit_args.into_args()
188+
jit_args.into_args().ok_or(ArgsError::NotAllArgsPassed)
168189
}

0 commit comments

Comments
 (0)