From 3b4f5935e5816bddd0a97f13a501f0d31f0daf9c Mon Sep 17 00:00:00 2001 From: kyokuping Date: Mon, 10 Aug 2026 17:33:33 +0900 Subject: [PATCH] Add JIT support for returning None --- crates/jit/src/instructions.rs | 39 ++++++++++++++++++---------------- crates/jit/src/lib.rs | 32 ++++++++++++++-------------- crates/jit/tests/misc_tests.rs | 19 ++++++++--------- 3 files changed, 46 insertions(+), 44 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 9b4656da5ad..67cf07f6e7f 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -39,6 +39,7 @@ impl JitValue { JitType::Int => Self::Int(val), JitType::Float => Self::Float(val), JitType::Bool => Self::Bool(val), + JitType::None => unreachable!("None cannot be used as an argument type"), } } @@ -47,7 +48,8 @@ impl JitValue { Self::Int(_) => Some(JitType::Int), Self::Float(_) => Some(JitType::Float), Self::Bool(_) => Some(JitType::Bool), - Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, + Self::None => Some(JitType::None), + Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, } } @@ -112,8 +114,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { #[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 cranelift_ty = ty.to_cranelift().ok_or(JitCompileError::NotSupported)?; let local = self.variables[idx].get_or_insert_with(|| { - let var = builder.declare_var(ty.to_cranelift()); + let var = builder.declare_var(cranelift_ty); Local { var, ty: ty.clone(), @@ -328,27 +331,27 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } 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) { + let val_type = val.to_jit_type().ok_or(JitCompileError::NotSupported)?; + if let Some(ref ret_type) = self.sig.ret { + if ret_type != &val_type { 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())); + self.sig.ret = Some(val_type.clone()); + if let Some(val_type) = val_type.to_cranelift() { + self.builder + .func + .signature + .returns + .push(AbiParam::new(val_type)); + } } - // 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]); + if let Some(cr_val) = val.into_value() { + self.builder.ins().return_(&[cr_val]); + } else { + self.builder.ins().return_(&[]); + } Ok(()) } diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index dbaa4a3eb26..0c700e93cf8 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -61,19 +61,12 @@ impl Jit { ret: Option, ) -> Result<(FuncId, JitSig), JitCompileError> { for arg in args { - self.ctx - .func - .signature - .params - .push(AbiParam::new(arg.to_cranelift())); + let arg = arg.to_cranelift().ok_or(JitCompileError::NotSupported)?; + self.ctx.func.signature.params.push(AbiParam::new(arg)); } - if ret.is_some() { - self.ctx - .func - .signature - .returns - .push(AbiParam::new(ret.clone().unwrap().to_cranelift())); + if let Some(ret) = ret.as_ref().and_then(JitType::to_cranelift) { + self.ctx.func.signature.returns.push(AbiParam::new(ret)); } let id = self.module.declare_function( @@ -167,7 +160,10 @@ impl CompiledCode { libffi::middle::CodePtr::from_ptr(self.code as *const _), cif_args, ); - self.sig.ret.as_ref().map(|ty| value.to_typed(ty)) + match self.sig.ret.as_ref() { + Some(JitType::None) | None => None, + Some(ty) => Some(value.to_typed(ty)), + } } } } @@ -193,14 +189,16 @@ pub enum JitType { Int, Float, Bool, + None, } impl JitType { - fn to_cranelift(&self) -> types::Type { + fn to_cranelift(&self) -> Option { match self { - Self::Int => types::I64, - Self::Float => types::F64, - Self::Bool => types::I8, + Self::Int => Some(types::I64), + Self::Float => Some(types::F64), + Self::Bool => Some(types::I8), + Self::None => None, } } @@ -209,6 +207,7 @@ impl JitType { Self::Int => libffi::middle::Type::i64(), Self::Float => libffi::middle::Type::f64(), Self::Bool => libffi::middle::Type::u8(), + Self::None => libffi::middle::Type::void(), } } } @@ -306,6 +305,7 @@ impl UnTypedAbiValue { JitType::Int => AbiValue::Int(self.int), JitType::Float => AbiValue::Float(self.float), JitType::Bool => AbiValue::Bool(self.boolean != 0), + JitType::None => unreachable!("None has no ABI value"), } } } diff --git a/crates/jit/tests/misc_tests.rs b/crates/jit/tests/misc_tests.rs index b73100ad6ec..5404df0a769 100644 --- a/crates/jit/tests/misc_tests.rs +++ b/crates/jit/tests/misc_tests.rs @@ -2,16 +2,15 @@ mod tests { use rustpython_jit::{AbiValue, JitArgumentError}; - // TODO currently broken - // #[test] - // fn test_no_return_value() { - // let func = jit_function! { func() => r##" - // def func(): - // pass - // "## }; - // - // assert_eq!(func(), Ok(())); - // } + #[test] + fn no_return_value() { + let func = jit_function! { func() => r##" + def func(): + pass + "## }; + + assert_eq!(func(), Ok(())); + } #[test] fn invoke() {