From c883b9dfd2d9f97b3d7f6712832e0d2a5258c541 Mon Sep 17 00:00:00 2001
From: ilyalesokhin-starkware <ilya@starkware.co>
Date: Wed, 14 Aug 2024 13:00:06 +0300
Subject: [PATCH] Add captured values to closure types. (#6187)

---
 .../cairo-lang-semantic/src/expr/compute.rs   | 20 +++++++++++++++---
 crates/cairo-lang-semantic/src/types.rs       |  4 ++++
 crates/cairo-lang-semantic/src/usage/mod.rs   | 21 ++++++++++++++-----
 3 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/crates/cairo-lang-semantic/src/expr/compute.rs b/crates/cairo-lang-semantic/src/expr/compute.rs
index 1261838b2d8..5a770da1035 100644
--- a/crates/cairo-lang-semantic/src/expr/compute.rs
+++ b/crates/cairo-lang-semantic/src/expr/compute.rs
@@ -32,9 +32,10 @@ use cairo_lang_utils::ordered_hash_set::OrderedHashSet;
 use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
 use cairo_lang_utils::unordered_hash_set::UnorderedHashSet;
 use cairo_lang_utils::{extract_matches, try_extract_matches, LookupIntern, OptionHelper};
-use itertools::{zip_eq, Itertools};
+use itertools::{chain, zip_eq, Itertools};
 use num_bigint::BigInt;
 use num_traits::ToPrimitive;
+use salsa::InternKey;
 use smol_str::SmolStr;
 use utils::Intern;
 
@@ -76,6 +77,7 @@ use crate::types::{
     peel_snapshots_ex, resolve_type, verify_fixed_size_array_size, wrap_in_snapshots,
     ClosureTypeLongId, ConcreteTypeId,
 };
+use crate::usage::Usages;
 use crate::{
     ConcreteEnumId, ConcreteTraitLongId, GenericArgumentId, GenericParam, Member, Mutability,
     Parameter, PatternStringLiteral, PatternStruct, Signature,
@@ -1531,14 +1533,26 @@ fn compute_expr_closure_semantic(
     if matches!(ctx.function_id, ContextFunction::Global) {
         ctx.diagnostics.report(syntax, ClosureInGlobalScope);
     }
+
+    let param_ids = params.iter().map(|param| param.id).collect_vec();
+    let mut usages = Usages { usages: Default::default() };
+    let usage = usages.handle_closure(&ctx.arenas, &param_ids, body);
+
+    let captured_types = chain!(
+        usage.snap_usage.values().map(|item| wrap_in_snapshots(ctx.db, item.ty(), 1)),
+        chain!(usage.usage.values(), usage.changes.values()).map(|item| item.ty())
+    )
+    .sorted_by_key(|ty| ty.as_intern_id())
+    .dedup()
+    .collect_vec();
     Ok(Expr::ExprClosure(ExprClosure {
         body,
-        param_ids: params.iter().map(|param| param.id).collect(),
-
+        param_ids,
         stable_ptr: syntax.stable_ptr().into(),
         ty: TypeLongId::Closure(ClosureTypeLongId {
             param_tys: params.iter().map(|param| param.ty).collect(),
             ret_ty,
+            captured_types,
             wrapper_location: StableLocation::new(syntax.wrapper(syntax_db).stable_ptr().into()),
         })
         .intern(ctx.db),
diff --git a/crates/cairo-lang-semantic/src/types.rs b/crates/cairo-lang-semantic/src/types.rs
index e71625f1319..6ed218ebf77 100644
--- a/crates/cairo-lang-semantic/src/types.rs
+++ b/crates/cairo-lang-semantic/src/types.rs
@@ -413,6 +413,10 @@ impl ConcreteExternTypeId {
 pub struct ClosureTypeLongId {
     pub param_tys: Vec<TypeId>,
     pub ret_ty: TypeId,
+    /// The set of types captured by the closure, this field is used to determined if the
+    /// closure has Drop, Destruct or PanicDestruct.
+    /// A vector as the fields needs to be hashable.
+    pub captured_types: Vec<TypeId>,
     /// Every closure has a unique type that is based on the stable location of its wrapper.
     #[dont_rewrite]
     pub wrapper_location: StableLocation,
diff --git a/crates/cairo-lang-semantic/src/usage/mod.rs b/crates/cairo-lang-semantic/src/usage/mod.rs
index 00dcd1a1c9f..16a1c5e15be 100644
--- a/crates/cairo-lang-semantic/src/usage/mod.rs
+++ b/crates/cairo-lang-semantic/src/usage/mod.rs
@@ -1,7 +1,7 @@
 //! 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::{MemberId, ParamId};
 use cairo_lang_proc_macros::DebugWithDb;
 use cairo_lang_utils::extract_matches;
 use cairo_lang_utils::ordered_hash_map::OrderedHashMap;
@@ -165,6 +165,20 @@ impl Usages {
         usages
     }
 
+    pub fn handle_closure(
+        &mut self,
+        arenas: &Arenas,
+        param_ids: &[ParamId],
+        body: ExprId,
+    ) -> Usage {
+        let mut usage: Usage = Default::default();
+
+        usage.introductions.extend(param_ids.iter().map(|id| VarId::Param(*id)));
+        self.handle_expr(arenas, body, &mut usage);
+        usage.finalize_as_scope();
+        usage
+    }
+
     fn handle_expr(&mut self, arenas: &Arenas, expr_id: ExprId, current: &mut Usage) {
         match &arenas.exprs[expr_id] {
             Expr::Tuple(expr) => {
@@ -286,11 +300,8 @@ impl Usages {
                 self.usages.insert(expr_id, usage);
             }
             Expr::ExprClosure(expr) => {
-                let mut usage: Usage = Default::default();
+                let usage = self.handle_closure(arenas, &expr.param_ids, expr.body);
 
-                usage.introductions.extend(expr.param_ids.iter().map(|id| VarId::Param(*id)));
-                self.handle_expr(arenas, expr.body, &mut usage);
-                usage.finalize_as_scope();
                 current.add_usage_and_changes(&usage);
                 self.usages.insert(expr_id, usage);
             }