compiler/rustc_next_trait_solver/src/solve/eval_ctxt/probe.rs

use std::marker::PhantomData;

use rustc_type_ir::search_graph::CandidateHeadUsages;
use rustc_type_ir::solve::{
    AccessedOpaques, CanonicalResponse, NoSolutionOrRerunNonErased, RerunResultExt,
};
use rustc_type_ir::{InferCtxtLike, Interner};
use tracing::{instrument, warn};

use crate::delegate::SolverDelegate;
use crate::solve::assembly::Candidate;
use crate::solve::{
    BuiltinImplSource, CandidateSource, EvalCtxt, Goal, GoalSource, GoalStalledOn, NoSolution,
    QueryResult, inspect,
};

pub(in crate::solve) struct ProbeCtxt<'me, 'a, D, I, F, T>
where
    D: SolverDelegate<Interner = I>,
    I: Interner,
{
    ecx: &'me mut EvalCtxt<'a, D, I>,
    probe_kind: F,
    _result: PhantomData<T>,
}

impl<D, I, F, T> ProbeCtxt<'_, '_, D, I, F, T>
where
    F: FnOnce(&Result<T, NoSolution>) -> inspect::ProbeKind<I>,
    D: SolverDelegate<Interner = I>,
    I: Interner,
{
    pub(in crate::solve) fn enter_single_candidate(
        self,
        f: impl FnOnce(&mut EvalCtxt<'_, D>) -> Result<T, NoSolutionOrRerunNonErased>,
    ) -> (Result<T, NoSolutionOrRerunNonErased>, CandidateHeadUsages) {
        let mut candidate_usages = CandidateHeadUsages::default();

        if let Err(e) = self.ecx.opaque_accesses.should_bail() {
            return (Err(e.into()), candidate_usages);
        }

        self.ecx.search_graph.enter_single_candidate();
        let result = self.enter(|ecx| {
            let result = f(ecx);
            candidate_usages = ecx.search_graph.finish_single_candidate();
            result
        });
        (result, candidate_usages)
    }

    pub(in crate::solve) fn enter(
        self,
        f: impl FnOnce(&mut EvalCtxt<'_, D>) -> Result<T, NoSolutionOrRerunNonErased>,
    ) -> Result<T, NoSolutionOrRerunNonErased> {
        let nested_goals = self.ecx.nested_goals.clone();
        self.enter_inner(f, nested_goals)
    }

    pub(in crate::solve) fn enter_without_propagated_nested_goals(
        self,
        f: impl FnOnce(&mut EvalCtxt<'_, D>) -> Result<T, NoSolutionOrRerunNonErased>,
    ) -> Result<T, NoSolutionOrRerunNonErased> {
        self.enter_inner(f, Default::default())
    }

    pub(in crate::solve) fn enter_inner(
        self,
        f: impl FnOnce(&mut EvalCtxt<'_, D>) -> Result<T, NoSolutionOrRerunNonErased>,
        propagated_nested_goals: Vec<(GoalSource, Goal<I, I::Predicate>, Option<GoalStalledOn<I>>)>,
    ) -> Result<T, NoSolutionOrRerunNonErased> {
        let ProbeCtxt { ecx: outer, probe_kind, _result } = self;

        outer.opaque_accesses.should_bail()?;

        let delegate = outer.delegate;
        let max_input_universe = outer.max_input_universe;
        let mut nested = EvalCtxt {
            delegate,
            var_kinds: outer.var_kinds,
            var_values: outer.var_values,
            current_goal_kind: outer.current_goal_kind,
            max_input_universe,
            initial_opaque_types_storage_num_entries: outer
                .initial_opaque_types_storage_num_entries,
            search_graph: outer.search_graph,
            nested_goals: propagated_nested_goals,
            origin_span: outer.origin_span,
            tainted: outer.tainted,
            inspect: outer.inspect.take_and_enter_probe(),
            opaque_accesses: AccessedOpaques::default(),
        };
        let r = nested.delegate.probe(|| {
            let r = f(&mut nested);
            nested.inspect.probe_final_state(delegate, max_input_universe);
            r
        });

        outer.opaque_accesses.update(nested.opaque_accesses)?;

        // Unwrap is unreachable, we would have returned on the line above.
        let r = r.map_err_to_rerun().unwrap();

        if !nested.inspect.is_noop() {
            let probe_kind = probe_kind(&r);
            nested.inspect.probe_kind(probe_kind);
            outer.inspect = nested.inspect.finish_probe();
        }

        r.map_err(Into::into)
    }
}

pub(in crate::solve) struct TraitProbeCtxt<'me, 'a, D, I, F>
where
    D: SolverDelegate<Interner = I>,
    I: Interner,
{
    cx: ProbeCtxt<'me, 'a, D, I, F, CanonicalResponse<I>>,
    source: CandidateSource<I>,
}

impl<D, I, F> TraitProbeCtxt<'_, '_, D, I, F>
where
    D: SolverDelegate<Interner = I>,
    I: Interner,
    F: FnOnce(&QueryResult<I>) -> inspect::ProbeKind<I>,
{
    #[instrument(level = "debug", skip_all, fields(source = ?self.source))]
    pub(in crate::solve) fn enter(
        self,
        f: impl FnOnce(&mut EvalCtxt<'_, D>) -> Result<CanonicalResponse<I>, NoSolutionOrRerunNonErased>,
    ) -> Result<Candidate<I>, NoSolutionOrRerunNonErased> {
        let (result, head_usages) = self.cx.enter_single_candidate(f);
        Ok(Candidate { source: self.source, result: result?, head_usages })
    }
}

impl<'a, D, I> EvalCtxt<'a, D, I>
where
    D: SolverDelegate<Interner = I>,
    I: Interner,
{
    /// `probe_kind` is only called when proof tree building is enabled so it can be
    /// as expensive as necessary to output the desired information.
    pub(in crate::solve) fn probe<F, T>(&mut self, probe_kind: F) -> ProbeCtxt<'_, 'a, D, I, F, T>
    where
        F: FnOnce(&Result<T, NoSolution>) -> inspect::ProbeKind<I>,
    {
        ProbeCtxt { ecx: self, probe_kind, _result: PhantomData }
    }

    pub(in crate::solve) fn probe_builtin_trait_candidate(
        &mut self,
        source: BuiltinImplSource,
    ) -> TraitProbeCtxt<'_, 'a, D, I, impl FnOnce(&QueryResult<I>) -> inspect::ProbeKind<I>> {
        self.probe_trait_candidate(CandidateSource::BuiltinImpl(source))
    }

    pub(in crate::solve) fn probe_trait_candidate(
        &mut self,
        source: CandidateSource<I>,
    ) -> TraitProbeCtxt<'_, 'a, D, I, impl FnOnce(&QueryResult<I>) -> inspect::ProbeKind<I>> {
        TraitProbeCtxt {
            cx: ProbeCtxt {
                ecx: self,
                probe_kind: move |result: &QueryResult<I>| inspect::ProbeKind::TraitCandidate {
                    source,
                    result: *result,
                },
                _result: PhantomData,
            },
            source,
        }
    }
}