starkware-libs · ilyalesokhin-starkware · Aug 7, 2024
diff --git a/crates/cairo-lang-lowering/src/lower/context.rs b/crates/cairo-lang-lowering/src/lower/context.rs
@@ -6,7 +6,6 @@ use cairo_lang_diagnostics::{DiagnosticAdded, Maybe};
 use cairo_lang_semantic::expr::fmt::ExprFormatter;
 use cairo_lang_semantic::items::enm::SemanticEnumEx;
 use cairo_lang_semantic::items::imp::ImplLookupContext;
-use cairo_lang_semantic::usage::Usages;
 use cairo_lang_syntax::node::ids::SyntaxStablePtrId;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
@@ -111,8 +110,6 @@ pub struct EncapsulatingLoweringContext<'db> {
     pub semantic_defs: UnorderedHashMap<semantic::VarId, semantic::Variable>,
     /// Expression formatter of the free function.
     pub expr_formatter: ExprFormatter<'db>,
-    /// Block usages for the entire encapsulating function.
-    pub usages: Usages,
     /// Lowerings of generated functions.
     pub lowerings: OrderedHashMap<semantic::ExprId, FlatLowered>,
 }
@@ -122,14 +119,12 @@ impl<'db> EncapsulatingLoweringContext<'db> {
         semantic_function_id: defs::ids::FunctionWithBodyId,
     ) -> Maybe<Self> {
         let function_body = db.function_body(semantic_function_id)?;
-        let usages = Usages::from_function_body(&function_body);
         Ok(Self {
             db,
             semantic_function_id,
             function_body,
             semantic_defs: Default::default(),
             expr_formatter: ExprFormatter { db: db.upcast(), function_id: semantic_function_id },
-            usages,
             lowerings: Default::default(),
         })
     }

diff --git a/crates/cairo-lang-lowering/src/lower/mod.rs b/crates/cairo-lang-lowering/src/lower/mod.rs
@@ -1387,7 +1387,7 @@ fn lower_expr_loop(
         _ => unreachable!("Loop expression must be either loop, while or for."),
     };
 
-    let usage = &ctx.usages.usages[&loop_expr_id];
+    let usage = ctx.db.function_expr_usage(ctx.semantic_function_id, loop_expr_id).unwrap();
 
     // Determine signature.
     let params = usage
@@ -1428,7 +1428,7 @@ fn lower_expr_loop(
     }
     .intern(ctx.db);
 
-    let snap_usage = ctx.usages.usages[&loop_expr_id].snap_usage.clone();
+    let snap_usage = usage.snap_usage.clone();
 
     // Generate the function.
     let encapsulating_ctx = std::mem::take(&mut ctx.encapsulating_ctx).unwrap();
@@ -1682,7 +1682,7 @@ fn lower_expr_closure(
     builder: &mut BlockBuilder,
 ) -> LoweringResult<LoweredExpr> {
     log::trace!("Lowering a closure expression: {:?}", expr.debug(&ctx.expr_formatter));
-    let usage = ctx.usages.usages[&expr_id].clone();
+    let usage = ctx.db.function_expr_usage(ctx.semantic_function_id, expr_id).unwrap();
     Ok(LoweredExpr::AtVariable(builder.capture(ctx, usage.clone(), expr)))
 }
 

diff --git a/crates/cairo-lang-semantic/src/db.rs b/crates/cairo-lang-semantic/src/db.rs
@@ -41,7 +41,8 @@ use crate::resolve::{ResolvedConcreteItem, ResolvedGenericItem, ResolverData};
 use crate::substitution::GenericSubstitution;
 use crate::types::{ImplTypeId, TypeSizeInformation};
 use crate::{
-    corelib, items, lsp_helpers, semantic, types, FunctionId, Parameter, SemanticDiagnostic, TypeId,
+    corelib, items, lsp_helpers, semantic, types, usage, ExprId, FunctionId, Parameter,
+    SemanticDiagnostic, TypeId,
 };
 
 /// Helper trait to make sure we can always get a `dyn SemanticGroup + 'static` from a
@@ -1479,6 +1480,15 @@ pub trait SemanticGroup:
     /// Aggregates file level semantic diagnostics.
     fn file_semantic_diagnostics(&self, file_id: FileId) -> Maybe<Diagnostics<SemanticDiagnostic>>;
 
+    // usage.
+    // ========
+    #[salsa::invoke(usage::function_expr_usage)]
+    fn function_expr_usage(
+        &self,
+        function_id: FunctionWithBodyId,
+        expr_id: ExprId,
+    ) -> Maybe<usage::Usage>;
+
     // Corelib.
     // ========
     #[salsa::invoke(corelib::core_crate)]

diff --git a/crates/cairo-lang-semantic/src/usage/mod.rs b/crates/cairo-lang-semantic/src/usage/mod.rs
@@ -1,18 +1,20 @@
 //! Introduces [Usages], which is responsible for computing variables usage in semantic blocks\
 //! of a function.
 
-use cairo_lang_defs::ids::MemberId;
+use cairo_lang_defs::ids::{FunctionWithBodyId, MemberId};
+use cairo_lang_diagnostics::Maybe;
 use cairo_lang_proc_macros::DebugWithDb;
 use cairo_lang_utils::extract_matches;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 use cairo_lang_utils::ordered_hash_set::OrderedHashSet;
 use id_arena::Arena;
 
+use crate::db::SemanticGroup;
 use crate::expr::fmt::ExprFormatter;
 use crate::expr::objects::Arenas;
 use crate::{
     ConcreteStructId, Condition, Expr, ExprFunctionCallArg, ExprId, ExprVarMemberPath,
-    FixedSizeArrayItems, FunctionBody, Pattern, PatternId, Statement, VarId,
+    FixedSizeArrayItems, Pattern, PatternId, Statement, VarId,
 };
 
 #[cfg(test)]
@@ -50,7 +52,7 @@ impl From<&ExprVarMemberPath> for MemberPath {
 }
 
 /// Usages of variables and member paths in semantic code.
-#[derive(Clone, Debug, Default, DebugWithDb)]
+#[derive(Clone, Debug, Default, PartialEq, Eq, DebugWithDb)]
 #[debug_db(ExprFormatter<'a>)]
 pub struct Usage {
     /// Member paths that are read.
@@ -151,21 +153,12 @@ impl Usage {
 }
 
 /// Usages of member paths in expressions of interest, currently loops and closures.
-#[derive(Debug, DebugWithDb)]
-#[debug_db(ExprFormatter<'a>)]
-pub struct Usages {
-    /// Mapping from an [ExprId] to its [Usage].
-    pub usages: OrderedHashMap<ExprId, Usage>,
+pub struct Usages<'db> {
+    db: &'db dyn SemanticGroup,
+    function_id: FunctionWithBodyId,
 }
-impl Usages {
-    pub fn from_function_body(function_body: &FunctionBody) -> Self {
-        let mut current = Usage::default();
-        let mut usages = Self { usages: Default::default() };
-        usages.handle_expr(&function_body.arenas, function_body.body_expr, &mut current);
-        usages
-    }
-
-    fn handle_expr(&mut self, arenas: &Arenas, expr_id: ExprId, current: &mut Usage) {
+impl Usages<'_> {
+    fn handle_expr(&self, arenas: &Arenas, expr_id: ExprId, current: &mut Usage) {
         match &arenas.exprs[expr_id] {
             Expr::Tuple(expr) => {
                 for expr_id in &expr.items {
@@ -240,58 +233,24 @@ impl Usages {
                 usage.finalize_as_scope();
                 current.add_usage_and_changes(&usage);
             }
-            Expr::Loop(expr) => {
-                let mut usage = Default::default();
-                self.handle_expr(arenas, expr.body, &mut usage);
+            Expr::Loop(_expr) => {
+                let usage = self.db.function_expr_usage(self.function_id, expr_id).unwrap();
                 current.add_usage_and_changes(&usage);
-                self.usages.insert(expr_id, usage);
             }
-            Expr::While(expr) => {
-                let mut usage = Default::default();
-                match &expr.condition {
-                    Condition::BoolExpr(expr) => {
-                        self.handle_expr(arenas, *expr, &mut usage);
-                    }
-                    Condition::Let(expr, patterns) => {
-                        self.handle_expr(arenas, *expr, &mut usage);
-                        for pattern in patterns {
-                            Self::handle_pattern(&arenas.patterns, *pattern, &mut usage);
-                        }
-                    }
-                }
-                self.handle_expr(arenas, expr.body, &mut usage);
-                usage.finalize_as_scope();
+            Expr::While(_expr) => {
+                let usage = self.db.function_expr_usage(self.function_id, expr_id).unwrap();
                 current.add_usage_and_changes(&usage);
-
-                self.usages.insert(expr_id, usage);
             }
             Expr::For(expr) => {
                 current.introductions.insert(
                     extract_matches!(&expr.into_iter_member_path, ExprVarMemberPath::Var).var,
                 );
-                let mut usage: Usage = Default::default();
-                usage.usage.insert(
-                    (&expr.into_iter_member_path).into(),
-                    expr.into_iter_member_path.clone(),
-                );
-                usage.changes.insert(
-                    (&expr.into_iter_member_path).into(),
-                    expr.into_iter_member_path.clone(),
-                );
-                Self::handle_pattern(&arenas.patterns, expr.pattern, &mut usage);
-                self.handle_expr(arenas, expr.body, &mut usage);
-                usage.finalize_as_scope();
+                let usage = self.db.function_expr_usage(self.function_id, expr_id).unwrap();
                 current.add_usage_and_changes(&usage);
-                self.usages.insert(expr_id, usage);
             }
-            Expr::ExprClosure(expr) => {
-                let mut usage: Usage = Default::default();
-
-                usage.introductions.extend(expr.param_ids.iter().map(|id| VarId::Param(*id)));
-                self.handle_expr(arenas, expr.body, &mut usage);
-                usage.finalize_as_scope();
+            Expr::ExprClosure(_expr) => {
+                let usage = self.db.function_expr_usage(self.function_id, expr_id).unwrap();
                 current.add_usage_and_changes(&usage);
-                self.usages.insert(expr_id, usage);
             }
             Expr::FunctionCall(expr) => {
                 for arg in &expr.args {
@@ -393,3 +352,55 @@ impl Usages {
         }
     }
 }
+
+/// Query implementation of [crate::db::SemanticGroup::function_expr_usage].
+pub fn function_expr_usage(
+    db: &dyn SemanticGroup,
+    function_id: FunctionWithBodyId,
+    expr_id: ExprId,
+) -> Maybe<Usage> {
+    let function_body = db.function_body(function_id)?;
+
+    let usages = Usages { db, function_id };
+    let arenas = &function_body.arenas;
+    let mut usage = Default::default();
+    match &arenas.exprs[expr_id] {
+        Expr::Loop(expr) => {
+            // Since expr.body is a block, we don't need to call `finalize_as_scope` on usage.
+            usages.handle_expr(arenas, expr.body, &mut usage);
+        }
+        Expr::While(expr) => {
+            match &expr.condition {
+                Condition::BoolExpr(expr) => {
+                    usages.handle_expr(arenas, *expr, &mut usage);
+                }
+                Condition::Let(expr, patterns) => {
+                    usages.handle_expr(arenas, *expr, &mut usage);
+                    for pattern in patterns {
+                        Usages::handle_pattern(&arenas.patterns, *pattern, &mut usage);
+                    }
+                }
+            }
+            usages.handle_expr(arenas, expr.body, &mut usage);
+            usage.finalize_as_scope();
+        }
+        Expr::For(expr) => {
+            usage
+                .usage
+                .insert((&expr.into_iter_member_path).into(), expr.into_iter_member_path.clone());
+            usage
+                .changes
+                .insert((&expr.into_iter_member_path).into(), expr.into_iter_member_path.clone());
+            Usages::handle_pattern(&arenas.patterns, expr.pattern, &mut usage);
+            usages.handle_expr(arenas, expr.body, &mut usage);
+            usage.finalize_as_scope();
+        }
+        Expr::ExprClosure(expr) => {
+            usage.introductions.extend(expr.param_ids.iter().map(|id| VarId::Param(*id)));
+            usages.handle_expr(arenas, expr.body, &mut usage);
+            usage.finalize_as_scope();
+        }
+        _ => unreachable!("Usage query is only supported for loops, whiles, fors and closures"),
+    }
+    Ok(usage)
+}
diff --git a/crates/cairo-lang-semantic/src/usage/test.rs b/crates/cairo-lang-semantic/src/usage/test.rs
@@ -7,7 +7,6 @@ use cairo_lang_syntax::node::TypedStablePtr;
 use cairo_lang_test_utils::parse_test_file::TestRunnerResult;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
 
-use super::Usages;
 use crate::db::SemanticGroup;
 use crate::expr::fmt::ExprFormatter;
 use crate::test_utils::{setup_test_function, SemanticDatabaseForTesting};
@@ -40,22 +39,22 @@ fn test_function_usage(
     let expr_formatter = ExprFormatter { db, function_id: test_function.function_id };
     let function_def =
         db.function_body(test_function.concrete_function_id.function_with_body_id(db)).unwrap();
-    let usages = Usages::from_function_body(&function_def);
 
     let mut usages_str = String::new();
-    for (expr_id, usage) in usages.usages.iter() {
-        let expr = &function_def.arenas.exprs[*expr_id];
-        let stable_ptr = expr.stable_ptr();
-        let node = stable_ptr.untyped().lookup(db);
-        let position = node.span_start_without_trivia(db).position_in_file(db, file_id).unwrap();
-
+    for (expr_id, expr) in function_def.arenas.exprs.iter() {
         match expr {
             Expr::Loop(_) => write!(usages_str, "Loop").unwrap(),
             Expr::While(_) => write!(usages_str, "While").unwrap(),
             Expr::For(_) => write!(usages_str, "For").unwrap(),
             Expr::ExprClosure(_) => write!(usages_str, "Closure").unwrap(),
-            _ => unreachable!(),
+            _ => continue,
         }
+
+        let usage = db.function_expr_usage(test_function.function_id, expr_id).unwrap();
+        let stable_ptr = expr.stable_ptr();
+        let node = stable_ptr.untyped().lookup(db);
+        let position = node.span_start_without_trivia(db).position_in_file(db, file_id).unwrap();
+
         writeln!(usages_str, " {}:{}:", position.line, position.col).unwrap();
         write!(usages_str, "  Usage:").unwrap();
         for (_, expr) in usage.usage.iter() {