diff --git a/Cargo.lock b/Cargo.lock index b4507ae3b..2709320c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1758,6 +1758,7 @@ dependencies = [ "sha3", "socket2", "statrs", + "thiserror", "thread_local", "timsort", "uname", diff --git a/jit/src/lib.rs b/jit/src/lib.rs index 48d98080d..a484e15fa 100644 --- a/jit/src/lib.rs +++ b/jit/src/lib.rs @@ -20,6 +20,12 @@ pub enum JitCompileError { CraneliftError(#[from] ModuleError), } +#[derive(Debug, thiserror::Error)] +pub enum JitArgumentError { + #[error("argument is of wrong type")] + ArgumentTypeMismatch, +} + struct Jit { builder_context: FunctionBuilderContext, ctx: codegen::Context, @@ -215,8 +221,14 @@ impl<'a> ArgsBuilder<'a> { } } - pub fn set(&mut self, idx: usize, value: AbiValue) { - self.values[idx] = Some(value); + pub fn set(&mut self, idx: usize, value: AbiValue) -> Result<(), JitArgumentError> { + match (&self.code.sig.args[idx], &value) { + (JitType::Int, AbiValue::Int(_)) | (JitType::Float, AbiValue::Float(_)) => { + self.values[idx] = Some(value); + Ok(()) + } + _ => Err(JitArgumentError::ArgumentTypeMismatch), + } } pub fn is_set(&self, idx: usize) -> bool { diff --git a/tests/snippets/jit.py b/tests/snippets/jit.py index 9a5428a56..4b92fa235 100644 --- a/tests/snippets/jit.py +++ b/tests/snippets/jit.py @@ -17,6 +17,7 @@ def tests(): assert foo() == 15 assert bar() == 2e5 assert baz(17, 20) == 49 + assert baz(17, 22.5) == 51.5 tests() diff --git a/vm/Cargo.toml b/vm/Cargo.toml index 2d43fc09d..e94454360 100644 --- a/vm/Cargo.toml +++ b/vm/Cargo.toml @@ -77,6 +77,7 @@ parking_lot = "0.11" thread_local = "1.0" cfg-if = "0.1.10" timsort = "0.1" +thiserror = "1.0" ## unicode stuff unicode_names2 = "0.4" diff --git a/vm/src/obj/objfunction.rs b/vm/src/obj/objfunction.rs index c6a3f359c..308466534 100644 --- a/vm/src/obj/objfunction.rs +++ b/vm/src/obj/objfunction.rs @@ -238,8 +238,15 @@ impl PyFunction { ) -> PyResult { #[cfg(feature = "jit")] if let Some(jitted_code) = self.jitted_code.get() { - if let Some(args) = jitfunc::get_jit_args(self, &func_args, jitted_code, vm) { - return Ok(jitted_code.invoke(&args).into_pyobject(vm)); + match jitfunc::get_jit_args(self, &func_args, jitted_code, vm) { + Ok(args) => { + return Ok(jitted_code.invoke(&args).into_pyobject(vm)); + } + Err(err) => info!( + "jit: function `{}` is falling back to being interpreted because of the \ + error: {}", + self.code.obj_name, err + ), } } diff --git a/vm/src/obj/objfunction/jitfunc.rs b/vm/src/obj/objfunction/jitfunc.rs index ec57fa3dd..cc3cbbe55 100644 --- a/vm/src/obj/objfunction/jitfunc.rs +++ b/vm/src/obj/objfunction/jitfunc.rs @@ -10,7 +10,25 @@ use crate::pyobject::{ use crate::VirtualMachine; use num_traits::ToPrimitive; use rustpython_bytecode::bytecode::CodeFlags; -use rustpython_jit::{AbiValue, Args, CompiledCode, JitType}; +use rustpython_jit::{AbiValue, Args, CompiledCode, JitArgumentError, JitType}; + +#[derive(Debug, thiserror::Error)] +pub enum ArgsError { + #[error("wrong number of arguments passed")] + WrongNumberOfArgs, + #[error("argument passed multiple times")] + ArgPassedMultipleTimes, + #[error("not a keyword argument")] + NotAKeywordArg, + #[error("not all arguments passed")] + NotAllArgsPassed, + #[error("integer can't fit into a machine integer")] + IntOverflow, + #[error("type can't be used in a jit function")] + NonJitType, + #[error("{0}")] + JitError(#[from] JitArgumentError), +} impl IntoPyObject for AbiValue { fn into_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { @@ -85,15 +103,18 @@ pub fn get_jit_arg_types(func: &PyFunctionRef, vm: &VirtualMachine) -> PyResult< } } -fn get_jit_value(vm: &VirtualMachine, obj: &PyObjectRef) -> Option { +fn get_jit_value(vm: &VirtualMachine, obj: &PyObjectRef) -> Result { // This does exact type checks as subclasses of int/float can't be passed to jitted functions let cls = obj.lease_class(); if cls.is(&vm.ctx.types.int_type) { - objint::get_value(&obj).to_i64().map(AbiValue::Int) + objint::get_value(&obj) + .to_i64() + .map(AbiValue::Int) + .ok_or(ArgsError::IntOverflow) } else if cls.is(&vm.ctx.types.float_type) { - Some(AbiValue::Float(objfloat::get_value(&obj))) + Ok(AbiValue::Float(objfloat::get_value(&obj))) } else { - None + Err(ArgsError::NonJitType) } } @@ -107,35 +128,35 @@ pub(crate) fn get_jit_args<'a>( func_args: &PyFuncArgs, jitted_code: &'a CompiledCode, vm: &VirtualMachine, -) -> Option> { +) -> Result, ArgsError> { let mut jit_args = jitted_code.args_builder(); let nargs = func_args.args.len(); if nargs > func.code.arg_names.len() || nargs < func.code.posonlyarg_count { - return None; + return Err(ArgsError::WrongNumberOfArgs); } // Add positional arguments for i in 0..nargs { - jit_args.set(i, get_jit_value(vm, &func_args.args[i])?); + jit_args.set(i, get_jit_value(vm, &func_args.args[i])?)?; } // Handle keyword arguments for (name, value) in &func_args.kwargs { if let Some(arg_idx) = func.code.arg_names.iter().position(|arg| arg == name) { if jit_args.is_set(arg_idx) { - return None; + return Err(ArgsError::ArgPassedMultipleTimes); } - jit_args.set(arg_idx, get_jit_value(vm, &value)?); + jit_args.set(arg_idx, get_jit_value(vm, &value)?)?; } else if let Some(kwarg_idx) = func.code.kwonlyarg_names.iter().position(|arg| arg == name) { let arg_idx = kwarg_idx + func.code.arg_names.len(); if jit_args.is_set(arg_idx) { - return None; + return Err(ArgsError::ArgPassedMultipleTimes); } - jit_args.set(arg_idx, get_jit_value(vm, &value)?); + jit_args.set(arg_idx, get_jit_value(vm, &value)?)?; } else { - return None; + return Err(ArgsError::NotAKeywordArg); } } @@ -145,7 +166,7 @@ pub(crate) fn get_jit_args<'a>( for (i, default) in defaults.iter().enumerate() { let arg_idx = i + func.code.arg_names.len() - defaults.len(); if !jit_args.is_set(arg_idx) { - jit_args.set(arg_idx, get_jit_value(vm, default)?); + jit_args.set(arg_idx, get_jit_value(vm, default)?)?; } } } @@ -157,12 +178,12 @@ pub(crate) fn get_jit_args<'a>( if !jit_args.is_set(arg_idx) { let default = kw_only_defaults .get_item(name.as_str(), vm) - .ok() + .map_err(|_| ArgsError::NotAllArgsPassed) .and_then(|obj| get_jit_value(vm, &obj))?; - jit_args.set(arg_idx, default); + jit_args.set(arg_idx, default)?; } } } - jit_args.into_args() + jit_args.into_args().ok_or(ArgsError::NotAllArgsPassed) }