From 39d21ae1d2aaa76fb5ab8ba11e11c5e7e955589f Mon Sep 17 00:00:00 2001
From: Jordan Jennings <jordan.jennings@mystenlabs.com>
Date: Tue, 15 Oct 2024 16:17:27 -0700
Subject: [PATCH] add return locations to source map

---
 .../src/source_map.rs                         | 20 +++++++++++++++++++
 .../move-ir-to-bytecode/src/compiler.rs       | 18 +++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/external-crates/move/crates/move-bytecode-source-map/src/source_map.rs b/external-crates/move/crates/move-bytecode-source-map/src/source_map.rs
index 92d5cb220c191..32ae09de384b3 100644
--- a/external-crates/move/crates/move-bytecode-source-map/src/source_map.rs
+++ b/external-crates/move/crates/move-bytecode-source-map/src/source_map.rs
@@ -67,6 +67,9 @@ pub struct FunctionSourceMap {
     /// The names of the parameters to the function.
     pub parameters: Vec<SourceName>,
 
+    /// The locations of the return values
+    pub returns: Vec<Loc>,
+
     /// The index into the vector is the local's index. The corresponding `(Identifier, Location)` tuple
     /// is the name and location of the local.
     pub locals: Vec<SourceName>,
@@ -206,6 +209,7 @@ impl FunctionSourceMap {
             definition_location,
             type_parameters: Vec::new(),
             parameters: Vec::new(),
+            returns: Vec::new(),
             locals: Vec::new(),
             code_map: BTreeMap::new(),
             is_native,
@@ -252,6 +256,10 @@ impl FunctionSourceMap {
         self.parameters.push(name)
     }
 
+    /// add the locations of return values
+    pub fn add_return_mapping(&mut self, loc: Loc) {
+        self.returns.push(loc);
+    }
     /// Recall that we are using a segment tree. We therefore lookup the location for the code
     /// offset by performing a range query for the largest number less than or equal to the code
     /// offset passed in.
@@ -451,6 +459,18 @@ impl SourceMap {
         Ok(())
     }
 
+    pub fn add_return_mapping(
+        &mut self,
+        fdef_idx: FunctionDefinitionIndex,
+        loc: Loc,
+    ) -> Result<()> {
+        let func_entry = self.function_map.get_mut(&fdef_idx.0).ok_or_else(|| {
+            format_err!("Tried to add return mapping to undefined function index")
+        })?;
+        func_entry.add_return_mapping(loc);
+        Ok(())
+    }
+
     pub fn get_parameter_or_local_name(
         &self,
         fdef_idx: FunctionDefinitionIndex,
diff --git a/external-crates/move/crates/move-ir-to-bytecode/src/compiler.rs b/external-crates/move/crates/move-ir-to-bytecode/src/compiler.rs
index efb218beb89aa..adaa773b53614 100644
--- a/external-crates/move/crates/move-ir-to-bytecode/src/compiler.rs
+++ b/external-crates/move/crates/move-ir-to-bytecode/src/compiler.rs
@@ -47,6 +47,11 @@ macro_rules! record_src_loc {
             .source_map
             .add_parameter_mapping($context.current_function_definition_index(), source_name)?;
     }};
+    (return_: $context:expr, $_type:expr) => {{
+        $context
+            .source_map
+            .add_return_mapping($context.current_function_definition_index(), $_type.loc)?;
+    }};
     (field: $context:expr, $idx: expr, $field:expr) => {{
         $context
             .source_map
@@ -882,6 +887,7 @@ fn compile_function_body_impl(
                 context,
                 m,
                 ast_function.signature.formals,
+                ast_function.signature.return_type,
                 locals,
                 code,
             )?)
@@ -898,6 +904,7 @@ fn compile_function_body_impl(
                 context,
                 m,
                 ast_function.signature.formals,
+                ast_function.signature.return_type,
                 locals,
                 code,
             )?)
@@ -907,6 +914,9 @@ fn compile_function_body_impl(
             for (var, _) in ast_function.signature.formals.into_iter() {
                 record_src_loc!(parameter: context, var)
             }
+            for _type in ast_function.signature.return_type.into_iter() {
+                record_src_loc!(return_: context, _type)
+            }
             None
         }
     })
@@ -954,6 +964,7 @@ fn compile_function_body(
     context: &mut Context,
     type_parameters: HashMap<TypeVar_, TypeParameterIndex>,
     formals: Vec<(Var, Type)>,
+    return_type: Vec<Type>,
     locals: Vec<(Var, Type)>,
     blocks: Vec<Block>,
 ) -> Result<CodeUnit> {
@@ -964,6 +975,9 @@ fn compile_function_body(
         record_src_loc!(parameter: context, var);
     }
 
+    for _type in return_type {
+        record_src_loc!(return_: context, _type);
+    }
     let mut locals_signature = Signature(vec![]);
     for (var_, t) in locals {
         let sig = compile_type(context, function_frame.type_parameters(), &t)?;
@@ -1716,6 +1730,7 @@ fn compile_function_body_bytecode(
     context: &mut Context,
     type_parameters: HashMap<TypeVar_, TypeParameterIndex>,
     formals: Vec<(Var, Type)>,
+    return_type: Vec<Type>,
     locals: Vec<(Var, Type)>,
     blocks: BytecodeBlocks,
 ) -> Result<CodeUnit> {
@@ -1726,6 +1741,9 @@ fn compile_function_body_bytecode(
         function_frame.define_local(&var.value, sig.clone())?;
         record_src_loc!(parameter: context, var);
     }
+    for _type in return_type {
+        record_src_loc!(return_: context, _type);
+    }
     for (var_, t) in locals {
         let sig = compile_type(context, function_frame.type_parameters(), &t)?;
         function_frame.define_local(&var_.value, sig.clone())?;