crates/hir-ty/src/opaques.rs

//! Handling of opaque types, detection of defining scope and hidden type.

use hir_def::{
    AssocItemId, AssocItemLoc, DefWithBodyId, FunctionId, HasModule, ItemContainerId, TypeAliasId,
    signatures::ImplSignature,
};
use hir_expand::name::Name;
use la_arena::{Arena, ArenaMap};
use rustc_type_ir::inherent::Ty as _;
use syntax::ast;

use crate::{
    ImplTraitId, InferBodyId, InferenceResult,
    db::{HirDatabase, InternedOpaqueTyId},
    lower::{ImplTrait, ImplTraitIdx},
    next_solver::{
        DbInterner, ErrorGuaranteed, SolverDefId, StoredEarlyBinder, StoredTy, Ty, TypingMode,
        infer::{DbInternerInferExt, traits::ObligationCause},
        obligation_ctxt::ObligationCtxt,
    },
};

pub(crate) fn opaque_types_defined_by<'db>(
    db: &'db dyn HirDatabase,
    def_id: InferBodyId<'_>,
    result: &mut Vec<SolverDefId<'db>>,
) {
    if let Some(func) = def_id.as_function() {
        // A function may define its own RPITs.
        extend_with_opaques(
            db,
            ImplTrait::return_type_impl_traits(db, func),
            |opaque_idx| ImplTraitId::ReturnTypeImplTrait(func, opaque_idx),
            result,
        );
    }

    let extend_with_taits = |type_alias| {
        extend_with_opaques(
            db,
            ImplTrait::type_alias_impl_traits(db, type_alias),
            |opaque_idx| ImplTraitId::TypeAliasImplTrait(type_alias, opaque_idx),
            result,
        );
    };

    // Collect opaques from assoc items.
    let extend_with_atpit_from_assoc_items = |assoc_items: &[(Name, AssocItemId)]| {
        assoc_items
            .iter()
            .filter_map(|&(_, assoc_id)| match assoc_id {
                AssocItemId::TypeAliasId(it) => Some(it),
                AssocItemId::FunctionId(_) | AssocItemId::ConstId(_) => None,
            })
            .for_each(extend_with_taits);
    };
    let extend_with_atpit_from_container = |container| match container {
        ItemContainerId::ImplId(impl_id)
            if ImplSignature::of(db, impl_id).target_trait.is_some() =>
        {
            extend_with_atpit_from_assoc_items(&impl_id.impl_items(db).items);
        }
        ItemContainerId::TraitId(trait_id) => {
            extend_with_atpit_from_assoc_items(&trait_id.trait_items(db).items);
        }
        _ => {}
    };
    match def_id {
        InferBodyId::DefWithBodyId(DefWithBodyId::ConstId(id)) => {
            extend_with_atpit_from_container(id.loc(db).container)
        }
        InferBodyId::DefWithBodyId(DefWithBodyId::FunctionId(id)) => {
            extend_with_atpit_from_container(id.loc(db).container)
        }
        InferBodyId::DefWithBodyId(DefWithBodyId::StaticId(_))
        | InferBodyId::DefWithBodyId(DefWithBodyId::VariantId(_))
        | InferBodyId::AnonConstId(_) => {}
    }

    // FIXME: Collect opaques from `#[define_opaque]`.

    fn extend_with_opaques<'db>(
        db: &'db dyn HirDatabase,
        opaques: &Arena<ImplTrait>,
        mut make_impl_trait: impl FnMut(ImplTraitIdx) -> ImplTraitId,
        result: &mut Vec<SolverDefId<'db>>,
    ) {
        for (opaque_idx, _) in opaques.iter() {
            let opaque_id = InternedOpaqueTyId::new(db, make_impl_trait(opaque_idx));
            result.push(opaque_id.into());
        }
    }
}

// These are firewall queries to prevent drawing dependencies between infers:

#[salsa::tracked(returns(ref))]
pub(crate) fn rpit_hidden_types(
    db: &dyn HirDatabase,
    function: FunctionId,
) -> ArenaMap<ImplTraitIdx, StoredEarlyBinder<StoredTy>> {
    let infer = InferenceResult::of(db, DefWithBodyId::from(function));
    let mut result = ArenaMap::new();
    for (opaque, hidden_type) in infer.return_position_impl_trait_types(db) {
        result.insert(opaque, StoredEarlyBinder::bind(hidden_type.store()));
    }
    result.shrink_to_fit();
    result
}

#[salsa::tracked(returns(ref))]
pub(crate) fn tait_hidden_types(
    db: &dyn HirDatabase,
    type_alias: TypeAliasId,
) -> ArenaMap<ImplTraitIdx, StoredEarlyBinder<StoredTy>> {
    // Call this first, to not perform redundant work if there are no TAITs.
    let taits_count = ImplTrait::type_alias_impl_traits(db, type_alias).len();

    let loc = type_alias.loc(db);
    let module = loc.module(db);
    let interner = DbInterner::new_with(db, module.krate(db));
    let infcx = interner.infer_ctxt().build(TypingMode::non_body_analysis());
    let mut ocx = ObligationCtxt::new(&infcx);
    let cause = ObligationCause::dummy();
    let param_env = db.trait_environment(type_alias.into());

    let defining_bodies = tait_defining_bodies(db, loc);

    let mut result = ArenaMap::with_capacity(taits_count);
    for defining_body in defining_bodies {
        let infer = InferenceResult::of(db, defining_body);
        for (&opaque, hidden_type) in &infer.type_of_opaque {
            let ImplTraitId::TypeAliasImplTrait(opaque_owner, opaque_idx) = opaque.loc(db) else {
                continue;
            };
            if opaque_owner != type_alias {
                continue;
            }
            // In the presence of errors, we attempt to create a unified type from all
            // types. rustc doesn't do that, but this should improve the experience.
            let hidden_type = infcx.insert_type_vars(hidden_type.as_ref());
            match result.entry(opaque_idx) {
                la_arena::Entry::Vacant(entry) => {
                    entry.insert(StoredEarlyBinder::bind(hidden_type.store()));
                }
                la_arena::Entry::Occupied(entry) => {
                    _ = ocx.eq(
                        &cause,
                        param_env,
                        entry.get().get().instantiate_identity().skip_norm_wip(),
                        hidden_type,
                    );
                }
            }
        }
    }

    _ = ocx.try_evaluate_obligations();

    // Fill missing entries.
    for idx in 0..taits_count {
        let idx = la_arena::Idx::from_raw(la_arena::RawIdx::from_u32(idx as u32));
        match result.entry(idx) {
            la_arena::Entry::Vacant(entry) => {
                entry.insert(StoredEarlyBinder::bind(
                    Ty::new_error(interner, ErrorGuaranteed).store(),
                ));
            }
            la_arena::Entry::Occupied(mut entry) => {
                let hidden_type = entry.get().get().skip_binder();
                let hidden_type =
                    infcx.resolve_vars_if_possible(hidden_type).replace_infer_with_error(interner);
                *entry.get_mut() = StoredEarlyBinder::bind(hidden_type.store());
            }
        }
    }

    result
}

fn tait_defining_bodies(
    db: &dyn HirDatabase,
    loc: &AssocItemLoc<ast::TypeAlias>,
) -> Vec<DefWithBodyId> {
    let from_assoc_items = |assoc_items: &[(Name, AssocItemId)]| {
        // Associated Type Position Impl Trait.
        assoc_items
            .iter()
            .filter_map(|&(_, assoc_id)| match assoc_id {
                AssocItemId::FunctionId(it) => Some(it.into()),
                AssocItemId::ConstId(it) => Some(it.into()),
                AssocItemId::TypeAliasId(_) => None,
            })
            .collect()
    };
    match loc.container {
        ItemContainerId::ImplId(impl_id)
            if ImplSignature::of(db, impl_id).target_trait.is_some() =>
        {
            return from_assoc_items(&impl_id.impl_items(db).items);
        }
        ItemContainerId::TraitId(trait_id) => {
            return from_assoc_items(&trait_id.trait_items(db).items);
        }
        _ => {}
    }

    // FIXME: Support general TAITs, or decisively decide not to.
    Vec::new()
}