Cache deref chain

Author: ilyalesokhin-starkware

Diff for: crates/cairo-lang-semantic/src/db.rs

@@ -926,6 +926,11 @@ pub trait SemanticGroup:
         impl_type_def_id: ImplTypeDefId,
     ) -> Maybe<GenericParamsData>;
 
+    /// Returns the deref chain and diagnostics for a given type.
+    #[salsa::invoke(items::imp::deref_chain)]
+    #[salsa::cycle(items::imp::deref_chain_cycle)]
+    fn deref_chain(&self, ty: TypeId, try_deref_mut: bool) -> Maybe<items::imp::DerefChain>;
+
     // Impl type.
     // ================
     /// Returns the implized impl type if the impl is concrete. Returns a TypeId that's not an impl

Diff for: crates/cairo-lang-semantic/src/expr/compute.rs

@@ -31,7 +31,6 @@ use cairo_lang_syntax::node::kind::SyntaxKind;
 use cairo_lang_syntax::node::{Terminal, TypedStablePtr, TypedSyntaxNode, ast};
 use cairo_lang_utils as utils;
 use cairo_lang_utils::ordered_hash_map::{Entry, OrderedHashMap};
-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::{Intern, LookupIntern, OptionHelper, extract_matches, try_extract_matches};
@@ -51,9 +50,9 @@ use super::pattern::{
     PatternOtherwise, PatternTuple, PatternVariable,
 };
 use crate::corelib::{
-    CoreTraitContext, core_binary_operator, core_bool_ty, core_unary_operator, deref_mut_trait,
-    deref_trait, false_literal_expr, get_core_trait, get_usize_ty, never_ty, numeric_literal_trait,
-    true_literal_expr, try_get_core_ty_by_name, unit_expr, unit_ty, unwrap_error_propagation_type,
+    CoreTraitContext, core_binary_operator, core_bool_ty, core_unary_operator, false_literal_expr,
+    get_core_trait, get_usize_ty, never_ty, numeric_literal_trait, true_literal_expr,
+    try_get_core_ty_by_name, unit_expr, unit_ty, unwrap_error_propagation_type,
 };
 use crate::db::SemanticGroup;
 use crate::diagnostic::SemanticDiagnosticKind::{self, *};
@@ -3040,12 +3039,13 @@ fn get_enriched_type_member_access(
     accessed_member_name: &str,
 ) -> Maybe<Option<EnrichedTypeMemberAccess>> {
     let (_, mut long_ty) = peel_snapshots(ctx.db, expr.ty());
-    if matches!(long_ty, TypeLongId::Var(_)) {
-        // Save some work. ignore the result. The error, if any, will be reported later.
-        ctx.resolver.inference().solve().ok();
-        long_ty = ctx.resolver.inference().rewrite(long_ty).no_err();
-    }
+    // Run solver to get as much info on the type as possible.
+    // Ignore the result of the `solve()` call - the error, if any, will be
+    // reported later.
+    ctx.resolver.inference().solve().ok();
+    ctx.resolver.inference().internal_rewrite(&mut long_ty).no_err();
     let (_, long_ty) = peel_snapshots_ex(ctx.db, long_ty);
+    let ty = long_ty.clone().intern(ctx.db);
     let base_var = match &expr.expr {
         Expr::Var(expr_var) => Some(expr_var.var),
         Expr::MemberAccess(ExprMemberAccess { member_path: Some(member_path), .. }) => {
@@ -3056,15 +3056,14 @@ fn get_enriched_type_member_access(
     let is_mut_var = base_var
         .filter(|var_id| matches!(ctx.semantic_defs.get(var_id), Some(var) if var.is_mut()))
         .is_some();
-    let ty = long_ty.clone().intern(ctx.db);
     let key = (ty, is_mut_var);
     let mut enriched_members = match ctx.resolver.type_enriched_members.entry(key) {
         Entry::Occupied(entry) => {
             let e = entry.get();
             match e.get_member(accessed_member_name) {
                 Some(value) => return Ok(Some(value)),
                 None => {
-                    if e.exploration_tail.is_none() {
+                    if e.exploration_done {
                         // There's no further exploration to be done, and member was not found.
                         return Ok(None);
                     }
@@ -3088,10 +3087,10 @@ fn get_enriched_type_member_access(
                 } else {
                     Default::default()
                 };
-            EnrichedMembers { members, deref_functions: vec![], exploration_tail: Some(expr.id) }
+            EnrichedMembers { members, deref_functions: vec![], exploration_done: false }
         }
     };
-    enrich_members(ctx, &mut enriched_members, is_mut_var, stable_ptr, accessed_member_name)?;
+    enrich_members(ctx, &mut enriched_members, ty, is_mut_var, stable_ptr, accessed_member_name)?;
     let e = ctx.resolver.type_enriched_members.entry(key).or_insert(enriched_members);
     Ok(e.get_member(accessed_member_name))
 }
@@ -3103,66 +3102,25 @@ fn get_enriched_type_member_access(
 fn enrich_members(
     ctx: &mut ComputationContext<'_>,
     enriched_members: &mut EnrichedMembers,
+    ty: TypeId,
     is_mut_var: bool,
     stable_ptr: ast::ExprPtr,
     accessed_member_name: &str,
 ) -> Maybe<()> {
-    let EnrichedMembers { members: enriched, deref_functions, exploration_tail } = enriched_members;
-    let mut visited_types: OrderedHashSet<TypeId> = OrderedHashSet::default();
-
-    let expr_id =
-        exploration_tail.expect("`enrich_members` should be called with a `calc_tail` value.");
-    let mut expr = ExprAndId { expr: ctx.arenas.exprs[expr_id].clone(), id: expr_id };
-
-    let deref_mut_trait_id = deref_mut_trait(ctx.db);
-    let deref_trait_id = deref_trait(ctx.db);
-
-    let compute_deref_method_function_call_data =
-        |ctx: &mut ComputationContext<'_>, expr: ExprAndId, use_deref_mut: bool| {
-            let deref_trait = if use_deref_mut { deref_mut_trait_id } else { deref_trait_id };
-            compute_method_function_call_data(
-                ctx,
-                &[deref_trait],
-                if use_deref_mut { "deref_mut".into() } else { "deref".into() },
-                expr.clone(),
-                stable_ptr.0,
-                None,
-                |_, _, _| None,
-                |_, _, _| None,
-            )
-        };
-
-    // If the variable is mutable, and implements DerefMut, we use DerefMut in the first iteration.
-    let mut use_deref_mut = deref_functions.is_empty()
-        && is_mut_var
-        && compute_deref_method_function_call_data(ctx, expr.clone(), true).is_ok();
+    let EnrichedMembers { members: enriched, deref_functions, exploration_done } = enriched_members;
 
-    // This function either finds a member and sets `exploration_tail` or finishes the exploration
-    // and leaves that exploration tail as `None`.
-    *exploration_tail = None;
+    let deref_chain = ctx.db.deref_chain(ty, is_mut_var)?;
 
     // Add members of derefed types.
-    while let Ok((function_id, _, cur_expr, mutability)) =
-        compute_deref_method_function_call_data(ctx, expr, use_deref_mut)
-    {
-        deref_functions.push((function_id, mutability));
-        use_deref_mut = false;
+    for deref_info in deref_chain.derefs.iter().skip(deref_functions.len()).cloned() {
+        deref_functions.push((deref_info.function_id, deref_info.self_mutability));
         let n_deref = deref_functions.len();
-        expr = cur_expr;
-        let derefed_expr = expr_function_call(
-            ctx,
-            function_id,
-            vec![NamedArg(expr, None, mutability)],
-            stable_ptr,
-            stable_ptr,
-        )?;
-        let ty = ctx.reduce_ty(derefed_expr.ty());
-        let (_, long_ty) = finalized_snapshot_peeled_ty(ctx, ty, stable_ptr)?;
+
+        let (_, long_ty) = finalized_snapshot_peeled_ty(ctx, deref_info.target_ty, stable_ptr)?;
         // If the type is still a variable we stop looking for derefed members.
         if let TypeLongId::Var(_) = long_ty {
             break;
         }
-        expr = ExprAndId { expr: derefed_expr.clone(), id: ctx.arenas.exprs.alloc(derefed_expr) };
         if let TypeLongId::Concrete(ConcreteTypeId::Struct(concrete_struct_id)) = long_ty {
             let members = ctx.db.concrete_struct_members(concrete_struct_id)?;
             for (member_name, member) in members.iter() {
@@ -3171,17 +3129,12 @@ fn enrich_members(
             }
             // If member is contained we can stop the calculation post the lookup.
             if members.contains_key(accessed_member_name) {
-                // Found member, so exploration isn't done - setting up the tail.
-                *exploration_tail = Some(expr.id);
-                break;
+                // Found member, so exploration isn't done.
+                return Ok(());
             }
         }
-        if !visited_types.insert(long_ty.intern(ctx.db)) {
-            // Break if we have a cycle. A diagnostic will be reported from the impl and not from
-            // member access.
-            break;
-        }
     }
+    *exploration_done = true;
     Ok(())
 }
 

Diff for: crates/cairo-lang-semantic/src/expr/inference/solver.rs

@@ -144,7 +144,7 @@ pub fn enrich_lookup_context(
 }
 
 /// Adds the defining module of the type to the lookup context.
-fn enrich_lookup_context_with_ty(
+pub fn enrich_lookup_context_with_ty(
     db: &dyn SemanticGroup,
     ty: TypeId,
     lookup_context: &mut ImplLookupContext,

Diff for: crates/cairo-lang-semantic/src/items/imp.rs

@@ -61,8 +61,8 @@ use super::visibility::peek_visible_in;
 use super::{TraitOrImplContext, resolve_trait_path};
 use crate::corelib::{
     CoreTraitContext, concrete_destruct_trait, concrete_drop_trait, copy_trait, core_crate,
-    deref_trait, destruct_trait, drop_trait, fn_once_trait, fn_trait, get_core_trait,
-    panic_destruct_trait,
+    deref_mut_trait, deref_trait, destruct_trait, drop_trait, fn_once_trait, fn_trait,
+    get_core_trait, panic_destruct_trait,
 };
 use crate::db::{SemanticGroup, get_resolver_data_options};
 use crate::diagnostic::SemanticDiagnosticKind::{self, *};
@@ -71,7 +71,7 @@ use crate::expr::compute::{ComputationContext, ContextFunction, Environment, com
 use crate::expr::inference::canonic::ResultNoErrEx;
 use crate::expr::inference::conform::InferenceConform;
 use crate::expr::inference::infers::InferenceEmbeddings;
-use crate::expr::inference::solver::SolutionSet;
+use crate::expr::inference::solver::{SolutionSet, enrich_lookup_context_with_ty};
 use crate::expr::inference::{
     ImplVarId, ImplVarTraitItemMappings, Inference, InferenceError, InferenceId,
 };
@@ -768,6 +768,106 @@ pub fn impl_semantic_definition_diagnostics(
     diagnostics.build()
 }
 
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct DerefChain {
+    pub derefs: Arc<[DerefInfo]>,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct DerefInfo {
+    pub function_id: FunctionId,
+    pub self_mutability: Mutability,
+    pub target_ty: TypeId,
+}
+
+pub fn deref_chain_cycle(
+    _db: &dyn SemanticGroup,
+    _cycle: &salsa::Cycle,
+    _ty: &TypeId,
+    _try_deref_mut: &bool,
+) -> Maybe<DerefChain> {
+    // `SemanticDiagnosticKind::DerefCycle` will be reported by `deref_impl_diagnostics`.
+    Maybe::Err(skip_diagnostic())
+}
+
+pub fn deref_chain(db: &dyn SemanticGroup, ty: TypeId, try_deref_mut: bool) -> Maybe<DerefChain> {
+    let mut opt_deref = None;
+    if try_deref_mut {
+        opt_deref = get_deref_func_and_target(db, ty, deref_mut_trait(db), "deref_mut".into())?;
+    }
+    let self_mutability = if opt_deref.is_some() {
+        Mutability::Reference
+    } else {
+        opt_deref = get_deref_func_and_target(db, ty, deref_trait(db), "deref".into())?;
+        Mutability::Immutable
+    };
+
+    let Some((function_id, target_ty)) = opt_deref else {
+        return Ok(DerefChain { derefs: Arc::new([]) });
+    };
+
+    let inner_chain = db.deref_chain(target_ty, false)?;
+
+    Ok(DerefChain {
+        derefs: chain!(
+            [DerefInfo { function_id, target_ty, self_mutability }],
+            inner_chain.derefs.iter().cloned()
+        )
+        .collect(),
+    })
+}
+
+pub fn get_deref_func_and_target(
+    db: &dyn SemanticGroup,
+    ty: TypeId,
+    deref_trait_id: TraitId,
+    deref_method: SmolStr,
+) -> Result<Option<(FunctionId, TypeId)>, DiagnosticAdded> {
+    let defs_db = db.upcast();
+    let mut lookup_context =
+        ImplLookupContext::new(deref_trait_id.module_file_id(defs_db).0, vec![]);
+    enrich_lookup_context_with_ty(db, ty, &mut lookup_context);
+    let concrete_trait = ConcreteTraitLongId {
+        trait_id: deref_trait_id,
+        generic_args: vec![GenericArgumentId::Type(ty)],
+    }
+    .intern(db);
+    let Ok(deref_impl) = get_impl_at_context(db, lookup_context, concrete_trait, None) else {
+        return Ok(None);
+    };
+    let concrete_impl_id = match deref_impl.lookup_intern(db) {
+        ImplLongId::Concrete(concrete_impl_id) => concrete_impl_id,
+        _ => panic!("Expected concrete impl"),
+    };
+
+    let deref_trait_func =
+        db.trait_function_by_name(deref_trait_id, deref_method).unwrap().unwrap();
+    let function_id = FunctionLongId {
+        function: ConcreteFunction {
+            generic_function: GenericFunctionId::Impl(ImplGenericFunctionId {
+                impl_id: deref_impl,
+                function: deref_trait_func,
+            }),
+            generic_args: vec![],
+        },
+    }
+    .intern(db);
+
+    let data = db.priv_impl_definition_data(concrete_impl_id.impl_def_id(db)).unwrap();
+    let mut types_iter = data.item_type_asts.iter();
+    let (impl_item_type_id, _) = types_iter.next().unwrap();
+    if types_iter.next().is_some() {
+        panic!(
+            "get_impl_based_on_single_impl_type called with an impl that has more than one type"
+        );
+    }
+    let ty = db.impl_type_def_resolved_type(*impl_item_type_id).unwrap();
+    let ty = SubstitutionRewriter { db, substitution: &concrete_impl_id.substitution(db)? }
+        .rewrite(ty)?;
+
+    Ok(Some((function_id, ty)))
+}
+
 /// Reports diagnostic for a deref impl.
 fn deref_impl_diagnostics(
     db: &dyn SemanticGroup,

Diff for: crates/cairo-lang-semantic/src/resolve/mod.rs

@@ -55,7 +55,7 @@ use crate::items::{TraitOrImplContext, visibility};
 use crate::substitution::{GenericSubstitution, SemanticRewriter};
 use crate::types::{ConcreteEnumLongId, ImplTypeId, are_coupons_enabled, resolve_type};
 use crate::{
-    ConcreteFunction, ConcreteTypeId, ConcreteVariant, ExprId, FunctionId, FunctionLongId,
+    ConcreteFunction, ConcreteTypeId, ConcreteVariant, FunctionId, FunctionLongId,
     GenericArgumentId, GenericParam, Member, Mutability, TypeId, TypeLongId,
 };
 
@@ -124,7 +124,7 @@ pub struct EnrichedMembers {
     /// from this point.
     /// Useful for partial computation of enriching members where a member was already previously
     /// found.
-    pub exploration_tail: Option<ExprId>,
+    pub exploration_done: bool,
 }
 impl EnrichedMembers {
     /// Returns `EnrichedTypeMemberAccess` for a single member if exists.