Added inlining location information. (starkware-libs#5367)

crates/cairo-lang-lowering/src/db.rs

@@ -642,12 +642,8 @@ fn function_with_body_lowering_diagnostics(
             && !lowered.signature.panicable
             && db.in_cycle(function_id, DependencyType::Cost)?
         {
-            let location = Location {
-                stable_location: function_id
-                    .base_semantic_function(db)
-                    .stable_location(db.upcast()),
-                notes: vec![],
-            };
+            let location =
+                Location::new(function_id.base_semantic_function(db).stable_location(db.upcast()));
             diagnostics.add(LoweringDiagnostic {
                 location,
                 kind: LoweringDiagnosticKind::NoPanicFunctionCycle,

crates/cairo-lang-lowering/src/inline/mod.rs

@@ -5,6 +5,7 @@ mod statements_weights;
 
 use std::collections::{HashMap, VecDeque};
 
+use cairo_lang_defs::diagnostic_utils::StableLocation;
 use cairo_lang_defs::ids::LanguageElementId;
 use cairo_lang_diagnostics::{Diagnostics, Maybe};
 use cairo_lang_semantic::items::functions::InlineConfiguration;
@@ -21,7 +22,8 @@ use crate::ids::{ConcreteFunctionWithBodyId, FunctionWithBodyId};
 use crate::lower::context::{VarRequest, VariableAllocator};
 use crate::utils::{Rebuilder, RebuilderEx};
 use crate::{
-    BlockId, FlatBlock, FlatBlockEnd, FlatLowered, Statement, VarRemapping, VarUsage, VariableId,
+    BlockId, FlatBlock, FlatBlockEnd, FlatLowered, Statement, StatementCall, VarRemapping,
+    VariableId,
 };
 
 pub fn get_inline_diagnostics(
@@ -164,6 +166,7 @@ pub struct Mapper<'a, 'b> {
     renamed_vars: HashMap<VariableId, VariableId>,
     return_block_id: BlockId,
     outputs: &'a [id_arena::Id<crate::Variable>],
+    inlining_location: StableLocation,
 
     /// An offset that is added to all the block IDs in order to translate them into the new
     /// lowering representation.
@@ -178,7 +181,9 @@ impl<'a, 'b> Rebuilder for Mapper<'a, 'b> {
         *self.renamed_vars.entry(orig_var_id).or_insert_with(|| {
             self.variables.new_var(VarRequest {
                 ty: self.lowered.variables[orig_var_id].ty,
-                location: self.lowered.variables[orig_var_id].location,
+                location: self.lowered.variables[orig_var_id]
+                    .location
+                    .inlined(self.variables.db, self.inlining_location),
             })
         })
     }
@@ -200,10 +205,20 @@ impl<'a, 'b> Rebuilder for Mapper<'a, 'b> {
                 };
                 *end = FlatBlockEnd::Goto(self.return_block_id, remapping);
             }
-            FlatBlockEnd::Panic(_) | FlatBlockEnd::Goto(_, _) | FlatBlockEnd::Match { .. } => {}
+            FlatBlockEnd::Panic(_) | FlatBlockEnd::Goto(_, _) => {}
+            FlatBlockEnd::Match { info } => {
+                let location = info.location_mut();
+                *location = location.inlined(self.variables.db, self.inlining_location);
+            }
             FlatBlockEnd::NotSet => unreachable!(),
         }
     }
+
+    fn transform_statement(&mut self, statement: &mut Statement) {
+        if let Some(location) = statement.location_mut() {
+            *location = location.inlined(self.variables.db, self.inlining_location);
+        }
+    }
 }
 
 impl<'db> FunctionInlinerRewriter<'db> {
@@ -264,7 +279,7 @@ impl<'db> FunctionInlinerRewriter<'db> {
                 if called_func != self.calling_function_id
                     && self.variables.db.priv_should_inline(called_func)?
                 {
-                    return self.inline_function(called_func, &stmt.inputs, &stmt.outputs);
+                    return self.inline_function(called_func, stmt);
                 }
             }
         }
@@ -273,16 +288,11 @@ impl<'db> FunctionInlinerRewriter<'db> {
         Ok(())
     }
 
-    /// Inlines the given function, with the given input and output variables.
-    /// The statements that need to replace the call statement in the original block
-    /// are pushed into the `statement_rewrite_stack`.
-    /// May also push additional blocks to the block queue.
-    /// The function takes an optional return block id to handle early returns.
+    /// Inlines the given function call.
     pub fn inline_function(
         &mut self,
         function_id: ConcreteFunctionWithBodyId,
-        inputs: &[VarUsage],
-        outputs: &[VariableId],
+        call_stmt: &StatementCall,
     ) -> Maybe<()> {
         let lowered = self.variables.db.inlined_function_with_body_lowered(function_id)?;
         lowered.blocks.has_root()?;
@@ -303,16 +313,20 @@ impl<'db> FunctionInlinerRewriter<'db> {
         // The input variables need to be renamed to match the inputs to the function call.
         let renamed_vars = HashMap::<VariableId, VariableId>::from_iter(izip!(
             lowered.parameters.iter().cloned(),
-            inputs.iter().map(|var_usage| var_usage.var_id)
+            call_stmt.inputs.iter().map(|var_usage| var_usage.var_id)
         ));
 
+        let db = self.variables.db;
+        let inlining_location = db.lookup_intern_location(call_stmt.location).stable_location;
+
         let mut mapper = Mapper {
             variables: &mut self.variables,
             lowered: &lowered,
             renamed_vars,
             block_id_offset: BlockId(return_block_id.0 + 1),
             return_block_id,
-            outputs,
+            outputs: &call_stmt.outputs,
+            inlining_location,
         };
 
         // The current block should Goto to the root block of the inlined function.

crates/cairo-lang-lowering/src/objects.rs

@@ -35,10 +35,12 @@ pub struct Location {
     /// auto-generated by the compiler.
     /// New notes are appended to the end of the vector.
     pub notes: Vec<DiagnosticNote>,
+    /// Function call locations where this value was inlined from.
+    pub inline_locations: Vec<StableLocation>,
 }
 impl Location {
     pub fn new(stable_location: StableLocation) -> Self {
-        Self { stable_location, notes: vec![] }
+        Self { stable_location, notes: vec![], inline_locations: vec![] }
     }
 
     /// Creates a new Location with the given note as the last note.
@@ -83,6 +85,22 @@ impl DebugWithDb<dyn LoweringGroup> for Location {
     }
 }
 
+impl LocationId {
+    /// Returns the location with the added inlining location of it.
+    pub fn inlined(self, db: &dyn LoweringGroup, inlining_location: StableLocation) -> Self {
+        let mut location = db.lookup_intern_location(self);
+        location.inline_locations.push(inlining_location);
+        db.intern_location(location)
+    }
+    /// Returns all relevant stable pointers of the location.
+    pub fn all_locations(self, db: &dyn LoweringGroup) -> Vec<StableLocation> {
+        let location = db.lookup_intern_location(self);
+        let mut all_locations = vec![location.stable_location];
+        all_locations.extend(location.inline_locations.iter().cloned());
+        all_locations
+    }
+}
+
 pub type VariableId = Id<Variable>;
 
 /// Represents a usage of a variable.
@@ -268,6 +286,17 @@ impl Statement {
             Statement::Desnap(stmt) => Some(stmt.input.location),
         }
     }
+    pub fn location_mut(&mut self) -> Option<&mut LocationId> {
+        match self {
+            Statement::Const(_) => None,
+            Statement::Call(stmt) => Some(&mut stmt.location),
+            Statement::StructConstruct(_) => None,
+            Statement::StructDestructure(stmt) => Some(&mut stmt.input.location),
+            Statement::EnumConstruct(stmt) => Some(&mut stmt.input.location),
+            Statement::Snapshot(stmt) => Some(&mut stmt.input.location),
+            Statement::Desnap(stmt) => Some(&mut stmt.input.location),
+        }
+    }
 }
 
 /// A statement that binds a const value to a variable.
@@ -441,6 +470,13 @@ impl MatchInfo {
             MatchInfo::Value(s) => &s.location,
         }
     }
+    pub fn location_mut(&mut self) -> &mut LocationId {
+        match self {
+            MatchInfo::Enum(s) => &mut s.location,
+            MatchInfo::Extern(s) => &mut s.location,
+            MatchInfo::Value(s) => &mut s.location,
+        }
+    }
 }
 
 /// Used in graph algorithms, and describes how to construct the edges in function dependency graph.

crates/cairo-lang-sierra-generator/src/block_generator.rs

@@ -48,9 +48,10 @@ pub fn generate_block_body_code(
     // Process the statements.
     for (i, statement) in block.statements.iter().enumerate() {
         let statement_lowering_location = (block_id, i);
-        let statement_cairo_location = statement.location().map(|location_id| {
-            context.get_db().lookup_intern_location(location_id).stable_location
-        });
+        let statement_cairo_location = statement
+            .location()
+            .map(|location_id| location_id.all_locations(context.get_db().upcast()))
+            .unwrap_or_default();
         context.maybe_set_cairo_location(statement_cairo_location);
         generate_statement_code(context, statement, &statement_lowering_location)?;
         let drop_location = &DropLocation::PostStatement(statement_lowering_location);
@@ -137,8 +138,7 @@ pub fn generate_block_code(
         // Process the block end if it's a match.
         lowering::FlatBlockEnd::Match { info } => {
             let statement_location = (block_id, block.statements.len());
-            let statement_cairo_location =
-                Some(context.get_db().lookup_intern_location(*info.location()).stable_location);
+            let statement_cairo_location = info.location().all_locations(context.get_db().upcast());
             if context.should_enable_ap_tracking(&block_id) {
                 context.set_ap_tracking(true);
                 context.push_statement(simple_basic_statement(
@@ -194,9 +194,11 @@ fn generate_push_values_statement_for_remapping(
             dup,
         })
     }
-    let location = remapping.iter().last().map(|(_, inner_output)| {
-        context.get_db().lookup_intern_location(inner_output.location).stable_location
-    });
+    let location = remapping
+        .iter()
+        .last()
+        .map(|(_, inner_output)| inner_output.location.all_locations(context.get_db().upcast()))
+        .unwrap_or_default();
     Ok(StatementWithLocation {
         statement: pre_sierra::Statement::PushValues(push_values),
         location,
@@ -232,7 +234,8 @@ pub fn generate_return_code(
     }
     let location = returned_variables
         .last()
-        .map(|var| context.get_db().lookup_intern_location(var.location).stable_location);
+        .map(|var| var.location.all_locations(context.get_db().upcast()))
+        .unwrap_or_default();
     context.maybe_set_cairo_location(location);
     context.push_statement(pre_sierra::Statement::PushValues(push_values));
     context.push_statement(return_statement(return_variables_on_stack));

crates/cairo-lang-sierra-generator/src/expr_generator_context.rs

@@ -34,7 +34,7 @@ pub struct ExprGeneratorContext<'a> {
     ap_tracking_configuration: ApTrackingConfiguration,
 
     /// The current location for adding statements.
-    curr_cairo_location: Option<StableLocation>,
+    curr_cairo_location: Vec<StableLocation>,
     /// The accumulated statements for the expression.
     statements: Vec<pre_sierra::StatementWithLocation>,
 }
@@ -59,7 +59,7 @@ impl<'a> ExprGeneratorContext<'a> {
             ap_tracking_enabled: true,
             ap_tracking_configuration,
             statements: vec![],
-            curr_cairo_location: None,
+            curr_cairo_location: vec![],
         }
     }
 
@@ -189,14 +189,14 @@ impl<'a> ExprGeneratorContext<'a> {
     pub fn push_statement(&mut self, statement: pre_sierra::Statement) {
         self.statements.push(pre_sierra::StatementWithLocation {
             statement,
-            location: self.curr_cairo_location,
+            location: self.curr_cairo_location.clone(),
         });
     }
 
     /// Sets up a location for the next pushed statements.
-    pub fn maybe_set_cairo_location(&mut self, location: Option<StableLocation>) {
-        if let Some(location) = location {
-            self.curr_cairo_location = Some(location);
+    pub fn maybe_set_cairo_location(&mut self, location: Vec<StableLocation>) {
+        if !location.is_empty() {
+            self.curr_cairo_location = location;
         }
     }
 

crates/cairo-lang-sierra-generator/src/pre_sierra.rs

@@ -71,7 +71,7 @@ pub enum Statement {
 }
 impl Statement {
     pub fn into_statement_without_location(self) -> StatementWithLocation {
-        StatementWithLocation { statement: self, location: None }
+        StatementWithLocation { statement: self, location: vec![] }
     }
     pub fn to_string(&self, db: &dyn SierraGenGroup) -> String {
         StatementWithDb { db, statement: self.clone() }.to_string()
@@ -82,11 +82,11 @@ impl Statement {
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct StatementWithLocation {
     pub statement: Statement,
-    pub location: Option<StableLocation>,
+    pub location: Vec<StableLocation>,
 }
 
 impl StatementWithLocation {
-    pub fn set_location(&mut self, location: Option<StableLocation>) {
+    pub fn set_location(&mut self, location: Vec<StableLocation>) {
         self.location = location;
     }
 }

crates/cairo-lang-sierra-generator/src/program_generator.rs

@@ -205,8 +205,10 @@ impl DebugWithDb<dyn SierraGenGroup> for SierraProgramWithDebug {
                 if let Some(loc) =
                     &self.debug_info.statements_locations.locations.get(&StatementIdx(i))
                 {
-                    let loc =
-                        get_location_marks(db.upcast(), &loc.diagnostic_location(db.upcast()));
+                    let loc = get_location_marks(
+                        db.upcast(),
+                        &loc.first().unwrap().diagnostic_location(db.upcast()),
+                    );
                     println!("{}", loc.split('\n').map(|l| format!("// {l}")).join("\n"));
                 }
             }

crates/cairo-lang-sierra-generator/src/resolve_labels.rs

@@ -12,7 +12,7 @@ use crate::pre_sierra;
 pub fn resolve_labels_and_extract_locations(
     statements: Vec<pre_sierra::StatementWithLocation>,
     label_replacer: &LabelReplacer,
-) -> (Vec<program::Statement>, Vec<Option<StableLocation>>) {
+) -> (Vec<program::Statement>, Vec<Vec<StableLocation>>) {
     statements
         .into_iter()
         .filter_map(|statement| match statement.statement {