Instruction::ReturnConst

RustPython · May 9, 2024 · 61d0f80 · 61d0f80
1 parent 85fa157
commit 61d0f80
Show file tree

Hide file tree

Showing 9 changed files with 149 additions and 114 deletions.
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
diff --git a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap
diff --git a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap
diff --git a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap
diff --git a/...r/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/...r/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
@@ -486,6 +486,9 @@ pub enum Instruction {
         target: Arg<Label>,
     },
     ReturnValue,
+    ReturnConst {
+        idx: Arg<u32>,
+    },
     YieldValue,
     YieldFrom,
     SetupAnnotation,
@@ -1141,7 +1144,12 @@ impl Instruction {
     pub fn unconditional_branch(&self) -> bool {
         matches!(
             self,
-            Jump { .. } | Continue { .. } | Break { .. } | ReturnValue | Raise { .. }
+            Jump { .. }
+                | Continue { .. }
+                | Break { .. }
+                | ReturnValue
+                | ReturnConst { .. }
+                | Raise { .. }
         )
     }
 
@@ -1219,6 +1227,7 @@ impl Instruction {
                 }
             }
             ReturnValue => -1,
+            ReturnConst { .. } => 0,
             YieldValue => 0,
             YieldFrom => -1,
             SetupAnnotation | SetupLoop | SetupFinally { .. } | EnterFinally | EndFinally => 0,
@@ -1322,6 +1331,23 @@ impl Instruction {
         let name = |i: u32| ctx.get_name(i as usize);
         let cell_name = |i: u32| ctx.get_cell_name(i as usize);
 
+        let fmt_const =
+            |op: &str, arg: OpArg, f: &mut fmt::Formatter, idx: &Arg<u32>| -> fmt::Result {
+                let value = ctx.get_constant(idx.get(arg) as usize);
+                match value.borrow_constant() {
+                    BorrowedConstant::Code { code } if expand_code_objects => {
+                        write!(f, "{:pad$}({:?}):", op, code)?;
+                        code.display_inner(f, true, level + 1)?;
+                        Ok(())
+                    }
+                    c => {
+                        write!(f, "{:pad$}(", op)?;
+                        c.fmt_display(f)?;
+                        write!(f, ")")
+                    }
+                }
+            };
+
         match self {
             ImportName { idx } => w!(ImportName, name = idx),
             ImportNameless => w!(ImportNameless),
@@ -1346,21 +1372,7 @@ impl Instruction {
             DeleteSubscript => w!(DeleteSubscript),
             StoreAttr { idx } => w!(StoreAttr, name = idx),
             DeleteAttr { idx } => w!(DeleteAttr, name = idx),
-            LoadConst { idx } => {
-                let value = ctx.get_constant(idx.get(arg) as usize);
-                match value.borrow_constant() {
-                    BorrowedConstant::Code { code } if expand_code_objects => {
-                        write!(f, "{:pad$}({:?}):", "LoadConst", code)?;
-                        code.display_inner(f, true, level + 1)?;
-                        Ok(())
-                    }
-                    c => {
-                        write!(f, "{:pad$}(", "LoadConst")?;
-                        c.fmt_display(f)?;
-                        write!(f, ")")
-                    }
-                }
-            }
+            LoadConst { idx } => fmt_const("LoadConst", arg, f, idx),
             UnaryOperation { op } => w!(UnaryOperation, ?op),
             BinaryOperation { op } => w!(BinaryOperation, ?op),
             BinaryOperationInplace { op } => w!(BinaryOperationInplace, ?op),
@@ -1390,6 +1402,7 @@ impl Instruction {
             CallMethodEx { has_kwargs } => w!(CallMethodEx, has_kwargs),
             ForIter { target } => w!(ForIter, target),
             ReturnValue => w!(ReturnValue),
+            ReturnConst { idx } => fmt_const("ReturnConst", arg, f, idx),
             YieldValue => w!(YieldValue),
             YieldFrom => w!(YieldFrom),
             SetupAnnotation => w!(SetupAnnotation),

diff --git a/jit/src/instructions.rs b/jit/src/instructions.rs
@@ -183,35 +183,48 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         Ok(())
     }
 
-    fn load_const<C: bytecode::Constant>(
+    fn prepare_const<C: bytecode::Constant>(
         &mut self,
         constant: BorrowedConstant<C>,
-    ) -> Result<(), JitCompileError> {
-        match constant {
+    ) -> Result<JitValue, JitCompileError> {
+        let value = match constant {
             BorrowedConstant::Integer { value } => {
                 let val = self.builder.ins().iconst(
                     types::I64,
                     value.to_i64().ok_or(JitCompileError::NotSupported)?,
                 );
-                self.stack.push(JitValue::Int(val));
-                Ok(())
+                JitValue::Int(val)
             }
             BorrowedConstant::Float { value } => {
                 let val = self.builder.ins().f64const(value);
-                self.stack.push(JitValue::Float(val));
-                Ok(())
+                JitValue::Float(val)
             }
             BorrowedConstant::Boolean { value } => {
                 let val = self.builder.ins().iconst(types::I8, value as i64);
-                self.stack.push(JitValue::Bool(val));
-                Ok(())
+                JitValue::Bool(val)
             }
-            BorrowedConstant::None => {
-                self.stack.push(JitValue::None);
-                Ok(())
+            BorrowedConstant::None => JitValue::None,
+            _ => return Err(JitCompileError::NotSupported),
+        };
+        Ok(value)
+    }
+
+    fn return_value(&mut self, val: JitValue) -> Result<(), JitCompileError> {
+        if let Some(ref ty) = self.sig.ret {
+            if val.to_jit_type().as_ref() != Some(ty) {
+                return Err(JitCompileError::NotSupported);
             }
-            _ => Err(JitCompileError::NotSupported),
+        } else {
+            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.builder.ins().return_(&[val.into_value().unwrap()]);
+        Ok(())
     }
 
     pub fn add_instruction<C: bytecode::Constant>(
@@ -269,7 +282,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 self.store_variable(idx.get(arg), val)
             }
             Instruction::LoadConst { idx } => {
-                self.load_const(constants[idx.get(arg) as usize].borrow_constant())
+                let val = self.prepare_const(constants[idx.get(arg) as usize].borrow_constant())?;
+                self.stack.push(val);
+                Ok(())
             }
             Instruction::BuildTuple { size } => {
                 let elements = self.pop_multiple(size.get(arg) as usize);
@@ -293,21 +308,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             }
             Instruction::ReturnValue => {
                 let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
-                if let Some(ref ty) = self.sig.ret {
-                    if val.to_jit_type().as_ref() != Some(ty) {
-                        return Err(JitCompileError::NotSupported);
-                    }
-                } else {
-                    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.builder.ins().return_(&[val.into_value().unwrap()]);
-                Ok(())
+                self.return_value(val)
+            }
+            Instruction::ReturnConst { idx } => {
+                let val = self.prepare_const(constants[idx.get(arg) as usize].borrow_constant())?;
+                self.return_value(val)
             }
             Instruction::CompareOperation { op, .. } => {
                 let op = op.get(arg);

diff --git a/jit/tests/common.rs b/jit/tests/common.rs
@@ -141,6 +141,11 @@ impl StackMachine {
                 let i = self.stack.len() - 3;
                 self.stack[i..].rotate_right(1);
             }
+            Instruction::ReturnConst { idx } => {
+                let idx = idx.get(arg);
+                self.stack.push(constants[idx as usize].clone().into());
+                return ControlFlow::Break(());
+            }
             Instruction::ReturnValue => return ControlFlow::Break(()),
             Instruction::ExtendedArg => {}
             _ => unimplemented!(

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -806,6 +806,10 @@ impl ExecutingFrame<'_> {
                 let value = self.pop_value();
                 self.unwind_blocks(vm, UnwindReason::Returning { value })
             }
+            bytecode::Instruction::ReturnConst { idx } => {
+                let value = self.code.constants[idx.get(arg) as usize].clone().into();
+                self.unwind_blocks(vm, UnwindReason::Returning { value })
+            }
             bytecode::Instruction::YieldValue => {
                 let value = self.pop_value();
                 let value = if self.code.flags.contains(bytecode::CodeFlags::IS_COROUTINE) {