From 019496e3a82340300670a8c932634a5778317258 Mon Sep 17 00:00:00 2001
From: Jeong YunWon <jeong@youknowone.org>
Date: Wed, 24 Apr 2024 02:36:52 +0900
Subject: [PATCH] Instruction::ReturnConst

---
 compiler/codegen/src/compile.rs               | 130 ++++++++++--------
 ...thon_codegen__compile__tests__if_ands.snap |   4 +-
 ...hon_codegen__compile__tests__if_mixed.snap |   4 +-
 ...ython_codegen__compile__tests__if_ors.snap |   4 +-
 ...pile__tests__nested_double_async_with.snap |   4 +-
 compiler/core/src/bytecode.rs                 |  45 +++---
 jit/src/instructions.rs                       |  63 +++++----
 jit/tests/common.rs                           |   5 +
 vm/src/frame.rs                               |   4 +
 9 files changed, 149 insertions(+), 114 deletions(-)

diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
index b04821b80..4ad592a0d 100644
--- a/compiler/codegen/src/compile.rs
+++ b/compiler/codegen/src/compile.rs
@@ -351,7 +351,7 @@ impl Compiler {
 
         let (doc, statements) = split_doc(body, &self.opts);
         if let Some(value) = doc {
-            self.emit_constant(ConstantData::Str { value });
+            self.emit_load_const(ConstantData::Str { value });
             let doc = self.name("__doc__");
             emit!(self, Instruction::StoreGlobal(doc))
         }
@@ -365,8 +365,7 @@ impl Compiler {
         assert_eq!(self.code_stack.len(), size_before);
 
         // Emit None at end:
-        self.emit_constant(ConstantData::None);
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_const(ConstantData::None);
         Ok(())
     }
 
@@ -393,13 +392,13 @@ impl Compiler {
                 emit!(self, Instruction::PrintExpr);
             } else {
                 self.compile_statement(last)?;
-                self.emit_constant(ConstantData::None);
+                self.emit_load_const(ConstantData::None);
             }
         } else {
-            self.emit_constant(ConstantData::None);
+            self.emit_load_const(ConstantData::None);
         };
 
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_value();
         Ok(())
     }
 
@@ -424,10 +423,10 @@ impl Compiler {
                     emit!(self, Instruction::Duplicate);
                     self.current_block().instructions.push(store_inst);
                 }
-                _ => self.emit_constant(ConstantData::None),
+                _ => self.emit_load_const(ConstantData::None),
             }
         }
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_value();
 
         Ok(())
     }
@@ -440,7 +439,7 @@ impl Compiler {
     ) -> CompileResult<()> {
         self.symbol_table_stack.push(symbol_table);
         self.compile_expression(expression)?;
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_value();
         Ok(())
     }
 
@@ -512,7 +511,7 @@ impl Compiler {
         };
 
         if NameUsage::Load == usage && name == "__debug__" {
-            self.emit_constant(ConstantData::Boolean {
+            self.emit_load_const(ConstantData::Boolean {
                 value: self.opts.optimize == 0,
             });
             return Ok(());
@@ -577,10 +576,10 @@ impl Compiler {
                 // import a, b, c as d
                 for name in names {
                     let name = &name;
-                    self.emit_constant(ConstantData::Integer {
+                    self.emit_load_const(ConstantData::Integer {
                         value: num_traits::Zero::zero(),
                     });
-                    self.emit_constant(ConstantData::None);
+                    self.emit_load_const(ConstantData::None);
                     let idx = self.name(&name.name);
                     emit!(self, Instruction::ImportName { idx });
                     if let Some(alias) = &name.asname {
@@ -624,10 +623,10 @@ impl Compiler {
                 let module_idx = module.as_ref().map(|s| self.name(s.as_str()));
 
                 // from .... import (*fromlist)
-                self.emit_constant(ConstantData::Integer {
+                self.emit_load_const(ConstantData::Integer {
                     value: level.as_ref().map_or(0, |level| level.to_u32()).into(),
                 });
-                self.emit_constant(ConstantData::Tuple {
+                self.emit_load_const(ConstantData::Tuple {
                     elements: from_list,
                 });
                 if let Some(idx) = module_idx {
@@ -852,13 +851,12 @@ impl Compiler {
                             ));
                         }
                         self.compile_expression(v)?;
+                        self.emit_return_value();
                     }
                     None => {
-                        self.emit_constant(ConstantData::None);
+                        self.emit_return_const(ConstantData::None);
                     }
                 }
-
-                emit!(self, Instruction::ReturnValue);
             }
             Stmt::Assign(StmtAssign { targets, value, .. }) => {
                 self.compile_expression(value)?;
@@ -942,7 +940,7 @@ impl Compiler {
         if !kw_with_defaults.is_empty() {
             let default_kw_count = kw_with_defaults.len();
             for (arg, default) in kw_with_defaults.iter() {
-                self.emit_constant(ConstantData::Str {
+                self.emit_load_const(ConstantData::Str {
                     value: arg.arg.to_string(),
                 });
                 self.compile_expression(default)?;
@@ -1201,8 +1199,7 @@ impl Compiler {
                 // the last instruction is a ReturnValue already, we don't need to emit it
             }
             _ => {
-                self.emit_constant(ConstantData::None);
-                emit!(self, Instruction::ReturnValue);
+                self.emit_return_const(ConstantData::None);
             }
         }
 
@@ -1217,7 +1214,7 @@ impl Compiler {
         // Return annotation:
         if let Some(annotation) = returns {
             // key:
-            self.emit_constant(ConstantData::Str {
+            self.emit_load_const(ConstantData::Str {
                 value: "return".to_owned(),
             });
             // value:
@@ -1234,7 +1231,7 @@ impl Compiler {
             .chain(args.kwarg.as_deref());
         for arg in args_iter {
             if let Some(annotation) = &arg.annotation {
-                self.emit_constant(ConstantData::Str {
+                self.emit_load_const(ConstantData::Str {
                     value: self.mangle(arg.arg.as_str()).into_owned(),
                 });
                 self.compile_annotation(annotation)?;
@@ -1256,10 +1253,10 @@ impl Compiler {
             func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
         }
 
-        self.emit_constant(ConstantData::Code {
+        self.emit_load_const(ConstantData::Code {
             code: Box::new(code),
         });
-        self.emit_constant(ConstantData::Str {
+        self.emit_load_const(ConstantData::Str {
             value: qualified_name,
         });
 
@@ -1268,7 +1265,7 @@ impl Compiler {
 
         if let Some(value) = doc_str {
             emit!(self, Instruction::Duplicate);
-            self.emit_constant(ConstantData::Str { value });
+            self.emit_load_const(ConstantData::Str { value });
             emit!(self, Instruction::Rotate2);
             let doc = self.name("__doc__");
             emit!(self, Instruction::StoreAttr { idx: doc });
@@ -1389,7 +1386,7 @@ impl Compiler {
         emit!(self, Instruction::LoadGlobal(dunder_name));
         let dunder_module = self.name("__module__");
         emit!(self, Instruction::StoreLocal(dunder_module));
-        self.emit_constant(ConstantData::Str {
+        self.emit_load_const(ConstantData::Str {
             value: qualified_name,
         });
         let qualname = self.name("__qualname__");
@@ -1417,10 +1414,10 @@ impl Compiler {
             let classcell = self.name("__classcell__");
             emit!(self, Instruction::StoreLocal(classcell));
         } else {
-            self.emit_constant(ConstantData::None);
+            self.emit_load_const(ConstantData::None);
         }
 
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_value();
 
         let code = self.pop_code_object();
 
@@ -1435,17 +1432,17 @@ impl Compiler {
             func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
         }
 
-        self.emit_constant(ConstantData::Code {
+        self.emit_load_const(ConstantData::Code {
             code: Box::new(code),
         });
-        self.emit_constant(ConstantData::Str {
+        self.emit_load_const(ConstantData::Str {
             value: name.to_owned(),
         });
 
         // Turn code object into function object:
         emit!(self, Instruction::MakeFunction(func_flags));
 
-        self.emit_constant(ConstantData::Str {
+        self.emit_load_const(ConstantData::Str {
             value: name.to_owned(),
         });
 
@@ -1462,7 +1459,7 @@ impl Compiler {
         // Duplicate top of stack (the function or class object)
 
         // Doc string value:
-        self.emit_constant(match doc_str {
+        self.emit_load_const(match doc_str {
             Some(doc) => ConstantData::Str { value: doc },
             None => ConstantData::None, // set docstring None if not declared
         });
@@ -1519,7 +1516,7 @@ impl Compiler {
             if is_async {
                 emit!(self, Instruction::BeforeAsyncWith);
                 emit!(self, Instruction::GetAwaitable);
-                self.emit_constant(ConstantData::None);
+                self.emit_load_const(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);
                 emit!(self, Instruction::SetupAsyncWith { end: final_block });
             } else {
@@ -1560,7 +1557,7 @@ impl Compiler {
 
         if is_async {
             emit!(self, Instruction::GetAwaitable);
-            self.emit_constant(ConstantData::None);
+            self.emit_load_const(ConstantData::None);
             emit!(self, Instruction::YieldFrom);
         }
 
@@ -1598,7 +1595,7 @@ impl Compiler {
                 }
             );
             emit!(self, Instruction::GetANext);
-            self.emit_constant(ConstantData::None);
+            self.emit_load_const(ConstantData::None);
             emit!(self, Instruction::YieldFrom);
             self.compile_store(target)?;
             emit!(self, Instruction::PopBlock);
@@ -1732,7 +1729,7 @@ impl Compiler {
 
     fn compile_annotation(&mut self, annotation: &located_ast::Expr) -> CompileResult<()> {
         if self.future_annotations {
-            self.emit_constant(ConstantData::Str {
+            self.emit_load_const(ConstantData::Str {
                 value: annotation.to_string(),
             });
         } else {
@@ -1764,7 +1761,7 @@ impl Compiler {
             // Store as dict entry in __annotations__ dict:
             let annotations = self.name("__annotations__");
             emit!(self, Instruction::LoadNameAny(annotations));
-            self.emit_constant(ConstantData::Str {
+            self.emit_load_const(ConstantData::Str {
                 value: self.mangle(id.as_str()).into_owned(),
             });
             emit!(self, Instruction::StoreSubscript);
@@ -2145,7 +2142,7 @@ impl Compiler {
                 self.compile_chained_comparison(left, ops, comparators)?;
             }
             Expr::Constant(ExprConstant { value, .. }) => {
-                self.emit_constant(compile_constant(value));
+                self.emit_load_const(compile_constant(value));
             }
             Expr::List(ExprList { elts, .. }) => {
                 let (size, unpack) = self.gather_elements(0, elts)?;
@@ -2180,7 +2177,7 @@ impl Compiler {
                 let mut compile_bound = |bound: Option<&located_ast::Expr>| match bound {
                     Some(exp) => self.compile_expression(exp),
                     None => {
-                        self.emit_constant(ConstantData::None);
+                        self.emit_load_const(ConstantData::None);
                         Ok(())
                     }
                 };
@@ -2199,7 +2196,7 @@ impl Compiler {
                 self.mark_generator();
                 match value {
                     Some(expression) => self.compile_expression(expression)?,
-                    Option::None => self.emit_constant(ConstantData::None),
+                    Option::None => self.emit_load_const(ConstantData::None),
                 };
                 emit!(self, Instruction::YieldValue);
             }
@@ -2209,7 +2206,7 @@ impl Compiler {
                 }
                 self.compile_expression(value)?;
                 emit!(self, Instruction::GetAwaitable);
-                self.emit_constant(ConstantData::None);
+                self.emit_load_const(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);
             }
             Expr::YieldFrom(ExprYieldFrom { value, .. }) => {
@@ -2225,12 +2222,12 @@ impl Compiler {
                 self.mark_generator();
                 self.compile_expression(value)?;
                 emit!(self, Instruction::GetIter);
-                self.emit_constant(ConstantData::None);
+                self.emit_load_const(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);
             }
             Expr::JoinedStr(ExprJoinedStr { values, .. }) => {
                 if let Some(value) = try_get_constant_string(values) {
-                    self.emit_constant(ConstantData::Str { value })
+                    self.emit_load_const(ConstantData::Str { value })
                 } else {
                     for value in values {
                         self.compile_expression(value)?;
@@ -2251,7 +2248,7 @@ impl Compiler {
             }) => {
                 match format_spec {
                     Some(spec) => self.compile_expression(spec)?,
-                    None => self.emit_constant(ConstantData::Str {
+                    None => self.emit_load_const(ConstantData::Str {
                         value: String::new(),
                     }),
                 };
@@ -2281,15 +2278,15 @@ impl Compiler {
                     .insert_full(ConstantData::None);
 
                 self.compile_expression(body)?;
-                emit!(self, Instruction::ReturnValue);
+                self.emit_return_value();
                 let code = self.pop_code_object();
                 if self.build_closure(&code) {
                     func_flags |= bytecode::MakeFunctionFlags::CLOSURE;
                 }
-                self.emit_constant(ConstantData::Code {
+                self.emit_load_const(ConstantData::Code {
                     code: Box::new(code),
                 });
-                self.emit_constant(ConstantData::Str { value: name });
+                self.emit_load_const(ConstantData::Str { value: name });
                 // Turn code object into function object:
                 emit!(self, Instruction::MakeFunction(func_flags));
 
@@ -2430,7 +2427,7 @@ impl Compiler {
                 let mut sub_size = 0;
                 for keyword in sub_keywords {
                     if let Some(name) = &keyword.arg {
-                        self.emit_constant(ConstantData::Str {
+                        self.emit_load_const(ConstantData::Str {
                             value: name.to_string(),
                         });
                         self.compile_expression(&keyword.value)?;
@@ -2540,7 +2537,7 @@ impl Compiler {
                 self.compile_expression(&keyword.value)?;
             }
 
-            self.emit_constant(ConstantData::Tuple {
+            self.emit_load_const(ConstantData::Tuple {
                 elements: kwarg_names,
             });
             CallType::Keyword { nargs: count }
@@ -2693,7 +2690,7 @@ impl Compiler {
                     }
                 );
                 emit!(self, Instruction::GetANext);
-                self.emit_constant(ConstantData::None);
+                self.emit_load_const(ConstantData::None);
                 emit!(self, Instruction::YieldFrom);
                 self.compile_store(&generator.target)?;
                 emit!(self, Instruction::PopBlock);
@@ -2727,11 +2724,11 @@ impl Compiler {
         }
 
         if return_none {
-            self.emit_constant(ConstantData::None)
+            self.emit_load_const(ConstantData::None)
         }
 
         // Return freshly filled list:
-        emit!(self, Instruction::ReturnValue);
+        self.emit_return_value();
 
         // Fetch code for listcomp function:
         let code = self.pop_code_object();
@@ -2744,12 +2741,12 @@ impl Compiler {
         }
 
         // List comprehension code:
-        self.emit_constant(ConstantData::Code {
+        self.emit_load_const(ConstantData::Code {
             code: Box::new(code),
         });
 
         // List comprehension function name:
-        self.emit_constant(ConstantData::Str {
+        self.emit_load_const(ConstantData::Str {
             value: name.to_owned(),
         });
 
@@ -2770,7 +2767,7 @@ impl Compiler {
         emit!(self, Instruction::CallFunctionPositional { nargs: 1 });
         if is_async {
             emit!(self, Instruction::GetAwaitable);
-            self.emit_constant(ConstantData::None);
+            self.emit_load_const(ConstantData::None);
             emit!(self, Instruction::YieldFrom);
         }
         Ok(())
@@ -2825,12 +2822,31 @@ impl Compiler {
 
     // fn block_done()
 
-    fn emit_constant(&mut self, constant: ConstantData) {
+    fn arg_constant(&mut self, constant: ConstantData) -> u32 {
         let info = self.current_code_info();
-        let idx = info.constants.insert_full(constant).0.to_u32();
+        info.constants.insert_full(constant).0.to_u32()
+    }
+
+    fn emit_load_const(&mut self, constant: ConstantData) {
+        let idx = self.arg_constant(constant);
         self.emit_arg(idx, |idx| Instruction::LoadConst { idx })
     }
 
+    fn emit_return_const(&mut self, constant: ConstantData) {
+        let idx = self.arg_constant(constant);
+        self.emit_arg(idx, |idx| Instruction::ReturnConst { idx })
+    }
+
+    fn emit_return_value(&mut self) {
+        if let Some(inst) = self.current_block().instructions.last_mut() {
+            if let Instruction::LoadConst { idx } = inst.instr {
+                inst.instr = Instruction::ReturnConst { idx };
+                return;
+            }
+        }
+        emit!(self, Instruction::ReturnValue)
+    }
+
     fn current_code_info(&mut self) -> &mut ir::CodeInfo {
         self.code_stack.last_mut().expect("no code on stack")
     }
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
index 93eb0ea6e..cf6223be1 100644
--- a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap
+++ b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ands.snap
@@ -9,6 +9,4 @@ expression: "compile_exec(\"\\\nif True and False and False:\n    pass\n\")"
               4 LoadConst            (False)
               5 JumpIfFalse          (6)
 
-  2     >>    6 LoadConst            (None)
-              7 ReturnValue
-
+  2     >>    6 ReturnConst          (None)
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
index 9594af69e..332d0fa98 100644
--- a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap
+++ b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_mixed.snap
@@ -11,6 +11,4 @@ expression: "compile_exec(\"\\\nif (True and False) or (False and True):\n    pa
               6 LoadConst            (True)
               7 JumpIfFalse          (8)
 
-  2     >>    8 LoadConst            (None)
-              9 ReturnValue
-
+  2     >>    8 ReturnConst          (None)
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
index dd582b7d7..d4a4d8f2c 100644
--- a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap
+++ b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__if_ors.snap
@@ -9,6 +9,4 @@ expression: "compile_exec(\"\\\nif True or False or False:\n    pass\n\")"
               4 LoadConst            (False)
               5 JumpIfFalse          (6)
 
-  2     >>    6 LoadConst            (None)
-              7 ReturnValue
-
+  2     >>    6 ReturnConst          (None)
diff --git a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
index dcc6f4c2c..91523b258 100644
--- a/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
+++ b/compiler/codegen/src/snapshots/rustpython_codegen__compile__tests__nested_double_async_with.snap
@@ -80,6 +80,4 @@ expression: "compile_exec(\"\\\nfor stop_exc in (StopIteration('spam'), StopAsyn
              66 WithCleanupFinish
              67 Jump                 (9)
         >>   68 PopBlock
-             69 LoadConst            (None)
-             70 ReturnValue
-
+             69 ReturnConst          (None)
diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
index 8ada60a9e..b00989ffe 100644
--- 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
index ee1cc0487..24585eba3 100644
--- 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
index 2f56e5db1..ef5e8cfbc 100644
--- 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
index 86c13c072..cc3e26585 100644
--- 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) {