Skip to content

Commit

Permalink
Add captured values to closure types. (#6187)
Browse files Browse the repository at this point in the history
  • Loading branch information
ilyalesokhin-starkware authored Aug 14, 2024
1 parent 86b6b55 commit c883b9d
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 8 deletions.
20 changes: 17 additions & 3 deletions crates/cairo-lang-semantic/src/expr/compute.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down
4 changes: 4 additions & 0 deletions crates/cairo-lang-semantic/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
21 changes: 16 additions & 5 deletions crates/cairo-lang-semantic/src/usage/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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) => {
Expand Down Expand Up @@ -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);
}
Expand Down

0 comments on commit c883b9d

Please sign in to comment.