From 6d87bd24f2f5e88757d697a21cc13398f0f932a1 Mon Sep 17 00:00:00 2001
From: Jort <jordan.jennings@mystenlabs.com>
Date: Wed, 16 Oct 2024 16:01:58 -0700
Subject: [PATCH] [move] add return locations to source map (#19885)

## Description

store the return loc's in source maps

## Test plan


---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol:
- [ ] Nodes (Validators and Full nodes):
- [ ] Indexer:
- [ ] JSON-RPC:
- [ ] GraphQL:
- [ ] CLI:
- [ ] Rust SDK:
- [ ] REST API:
---
 .../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())?;