use std::assert_matches::assert_matches;
use rustc_ast_ir::try_visit;
use rustc_ast_ir::visit::VisitorResult;
use rustc_infer::infer::{DefineOpaqueTypes, InferCtxt, InferOk};
use rustc_macros::extension;
use rustc_middle::traits::solve::{Certainty, Goal, GoalSource, NoSolution, QueryResult};
use rustc_middle::traits::ObligationCause;
use rustc_middle::ty::{TyCtxt, TypeFoldable};
use rustc_middle::{bug, ty};
use rustc_next_trait_solver::resolve::EagerResolver;
use rustc_next_trait_solver::solve::inspect::{self, instantiate_canonical_state};
use rustc_next_trait_solver::solve::{GenerateProofTree, MaybeCause, SolverDelegateEvalExt as _};
use rustc_span::{Span, DUMMY_SP};
use tracing::instrument;
use crate::solve::delegate::SolverDelegate;
use crate::traits::ObligationCtxt;
pub struct InspectConfig {
pub max_depth: usize,
}
pub struct InspectGoal<'a, 'tcx> {
infcx: &'a SolverDelegate<'tcx>,
depth: usize,
orig_values: Vec<ty::GenericArg<'tcx>>,
goal: Goal<'tcx, ty::Predicate<'tcx>>,
result: Result<Certainty, NoSolution>,
evaluation_kind: inspect::CanonicalGoalEvaluationKind<TyCtxt<'tcx>>,
normalizes_to_term_hack: Option<NormalizesToTermHack<'tcx>>,
source: GoalSource,
}
#[derive(Copy, Clone)]
struct NormalizesToTermHack<'tcx> {
term: ty::Term<'tcx>,
unconstrained_term: ty::Term<'tcx>,
}
impl<'tcx> NormalizesToTermHack<'tcx> {
fn constrain(
self,
infcx: &InferCtxt<'tcx>,
span: Span,
param_env: ty::ParamEnv<'tcx>,
) -> Result<Certainty, NoSolution> {
infcx
.at(&ObligationCause::dummy_with_span(span), param_env)
.eq(DefineOpaqueTypes::Yes, self.term, self.unconstrained_term)
.map_err(|_| NoSolution)
.and_then(|InferOk { value: (), obligations }| {
let ocx = ObligationCtxt::new(infcx);
ocx.register_obligations(obligations);
let errors = ocx.select_all_or_error();
if errors.is_empty() {
Ok(Certainty::Yes)
} else if errors.iter().all(|e| !e.is_true_error()) {
Ok(Certainty::AMBIGUOUS)
} else {
Err(NoSolution)
}
})
}
}
pub struct InspectCandidate<'a, 'tcx> {
goal: &'a InspectGoal<'a, 'tcx>,
kind: inspect::ProbeKind<TyCtxt<'tcx>>,
steps: Vec<&'a inspect::ProbeStep<TyCtxt<'tcx>>>,
final_state: inspect::CanonicalState<TyCtxt<'tcx>, ()>,
result: QueryResult<'tcx>,
shallow_certainty: Certainty,
}
impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
pub fn kind(&self) -> inspect::ProbeKind<TyCtxt<'tcx>> {
self.kind
}
pub fn result(&self) -> Result<Certainty, NoSolution> {
self.result.map(|c| c.value.certainty)
}
pub fn goal(&self) -> &'a InspectGoal<'a, 'tcx> {
self.goal
}
pub fn shallow_certainty(&self) -> Certainty {
self.shallow_certainty
}
pub fn visit_nested_no_probe<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
for goal in self.instantiate_nested_goals(visitor.span()) {
try_visit!(goal.visit_with(visitor));
}
V::Result::output()
}
pub fn instantiate_nested_goals(&self, span: Span) -> Vec<InspectGoal<'a, 'tcx>> {
self.instantiate_nested_goals_and_opt_impl_args(span).0
}
#[instrument(
level = "debug",
skip_all,
fields(goal = ?self.goal.goal, steps = ?self.steps)
)]
pub fn instantiate_nested_goals_and_opt_impl_args(
&self,
span: Span,
) -> (Vec<InspectGoal<'a, 'tcx>>, Option<ty::GenericArgsRef<'tcx>>) {
let infcx = self.goal.infcx;
let param_env = self.goal.goal.param_env;
let mut orig_values = self.goal.orig_values.to_vec();
let mut instantiated_goals = vec![];
let mut opt_impl_args = None;
for step in &self.steps {
match **step {
inspect::ProbeStep::AddGoal(source, goal) => instantiated_goals.push((
source,
instantiate_canonical_state(infcx, span, param_env, &mut orig_values, goal),
)),
inspect::ProbeStep::RecordImplArgs { impl_args } => {
opt_impl_args = Some(instantiate_canonical_state(
infcx,
span,
param_env,
&mut orig_values,
impl_args,
));
}
inspect::ProbeStep::MakeCanonicalResponse { .. }
| inspect::ProbeStep::NestedProbe(_) => unreachable!(),
}
}
let () =
instantiate_canonical_state(infcx, span, param_env, &mut orig_values, self.final_state);
if let Some(term_hack) = self.goal.normalizes_to_term_hack {
let _ = term_hack.constrain(infcx, span, param_env);
}
let opt_impl_args =
opt_impl_args.map(|impl_args| impl_args.fold_with(&mut EagerResolver::new(infcx)));
let goals = instantiated_goals
.into_iter()
.map(|(source, goal)| match goal.predicate.kind().no_bound_vars() {
Some(ty::PredicateKind::NormalizesTo(ty::NormalizesTo { alias, term })) => {
let unconstrained_term = match term.unpack() {
ty::TermKind::Ty(_) => infcx.next_ty_var(span).into(),
ty::TermKind::Const(_) => infcx.next_const_var(span).into(),
};
let goal =
goal.with(infcx.tcx, ty::NormalizesTo { alias, term: unconstrained_term });
let proof_tree = infcx
.probe(|_| infcx.evaluate_root_goal_raw(goal, GenerateProofTree::Yes).1);
InspectGoal::new(
infcx,
self.goal.depth + 1,
proof_tree.unwrap(),
Some(NormalizesToTermHack { term, unconstrained_term }),
source,
)
}
_ => {
let proof_tree = infcx
.probe(|_| infcx.evaluate_root_goal(goal, GenerateProofTree::Yes).1)
.unwrap();
InspectGoal::new(infcx, self.goal.depth + 1, proof_tree, None, source)
}
})
.collect();
(goals, opt_impl_args)
}
pub fn visit_nested_in_probe<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
self.goal.infcx.probe(|_| self.visit_nested_no_probe(visitor))
}
}
impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
pub fn infcx(&self) -> &'a InferCtxt<'tcx> {
self.infcx
}
pub fn goal(&self) -> Goal<'tcx, ty::Predicate<'tcx>> {
self.goal
}
pub fn result(&self) -> Result<Certainty, NoSolution> {
self.result
}
pub fn source(&self) -> GoalSource {
self.source
}
pub fn depth(&self) -> usize {
self.depth
}
fn candidates_recur(
&'a self,
candidates: &mut Vec<InspectCandidate<'a, 'tcx>>,
steps: &mut Vec<&'a inspect::ProbeStep<TyCtxt<'tcx>>>,
probe: &'a inspect::Probe<TyCtxt<'tcx>>,
) {
let mut shallow_certainty = None;
for step in &probe.steps {
match *step {
inspect::ProbeStep::AddGoal(..) | inspect::ProbeStep::RecordImplArgs { .. } => {
steps.push(step)
}
inspect::ProbeStep::MakeCanonicalResponse { shallow_certainty: c } => {
assert_matches!(
shallow_certainty.replace(c),
None | Some(Certainty::Maybe(MaybeCause::Ambiguity))
);
}
inspect::ProbeStep::NestedProbe(ref probe) => {
match probe.kind {
inspect::ProbeKind::UpcastProjectionCompatibility
| inspect::ProbeKind::ShadowedEnvProbing => continue,
inspect::ProbeKind::NormalizedSelfTyAssembly
| inspect::ProbeKind::UnsizeAssembly
| inspect::ProbeKind::Root { .. }
| inspect::ProbeKind::TryNormalizeNonRigid { .. }
| inspect::ProbeKind::TraitCandidate { .. }
| inspect::ProbeKind::OpaqueTypeStorageLookup { .. } => {
let num_steps = steps.len();
self.candidates_recur(candidates, steps, probe);
steps.truncate(num_steps);
}
}
}
}
}
match probe.kind {
inspect::ProbeKind::UpcastProjectionCompatibility
| inspect::ProbeKind::ShadowedEnvProbing => bug!(),
inspect::ProbeKind::NormalizedSelfTyAssembly | inspect::ProbeKind::UnsizeAssembly => {}
inspect::ProbeKind::Root { result }
| inspect::ProbeKind::TryNormalizeNonRigid { result }
| inspect::ProbeKind::TraitCandidate { source: _, result }
| inspect::ProbeKind::OpaqueTypeStorageLookup { result } => {
if let Some(shallow_certainty) = shallow_certainty {
candidates.push(InspectCandidate {
goal: self,
kind: probe.kind,
steps: steps.clone(),
final_state: probe.final_state,
shallow_certainty,
result,
});
}
}
}
}
pub fn candidates(&'a self) -> Vec<InspectCandidate<'a, 'tcx>> {
let mut candidates = vec![];
let last_eval_step = match &self.evaluation_kind {
inspect::CanonicalGoalEvaluationKind::Overflow => return vec![],
inspect::CanonicalGoalEvaluationKind::Evaluation { final_revision } => final_revision,
};
let mut nested_goals = vec![];
self.candidates_recur(&mut candidates, &mut nested_goals, &last_eval_step.evaluation);
candidates
}
pub fn unique_applicable_candidate(&'a self) -> Option<InspectCandidate<'a, 'tcx>> {
let mut candidates = self.candidates();
candidates.retain(|c| c.result().is_ok());
candidates.pop().filter(|_| candidates.is_empty())
}
fn new(
infcx: &'a InferCtxt<'tcx>,
depth: usize,
root: inspect::GoalEvaluation<TyCtxt<'tcx>>,
normalizes_to_term_hack: Option<NormalizesToTermHack<'tcx>>,
source: GoalSource,
) -> Self {
let infcx = <&SolverDelegate<'tcx>>::from(infcx);
let inspect::GoalEvaluation { uncanonicalized_goal, orig_values, evaluation } = root;
let result = evaluation.result.and_then(|ok| {
if let Some(term_hack) = normalizes_to_term_hack {
infcx
.probe(|_| term_hack.constrain(infcx, DUMMY_SP, uncanonicalized_goal.param_env))
.map(|certainty| ok.value.certainty.unify_with(certainty))
} else {
Ok(ok.value.certainty)
}
});
InspectGoal {
infcx,
depth,
orig_values,
goal: uncanonicalized_goal.fold_with(&mut EagerResolver::new(infcx)),
result,
evaluation_kind: evaluation.kind,
normalizes_to_term_hack,
source,
}
}
pub(crate) fn visit_with<V: ProofTreeVisitor<'tcx>>(&self, visitor: &mut V) -> V::Result {
if self.depth < visitor.config().max_depth {
try_visit!(visitor.visit_goal(self));
}
V::Result::output()
}
}
pub trait ProofTreeVisitor<'tcx> {
type Result: VisitorResult = ();
fn span(&self) -> Span;
fn config(&self) -> InspectConfig {
InspectConfig { max_depth: 10 }
}
fn visit_goal(&mut self, goal: &InspectGoal<'_, 'tcx>) -> Self::Result;
}
#[extension(pub trait ProofTreeInferCtxtExt<'tcx>)]
impl<'tcx> InferCtxt<'tcx> {
fn visit_proof_tree<V: ProofTreeVisitor<'tcx>>(
&self,
goal: Goal<'tcx, ty::Predicate<'tcx>>,
visitor: &mut V,
) -> V::Result {
self.visit_proof_tree_at_depth(goal, 0, visitor)
}
fn visit_proof_tree_at_depth<V: ProofTreeVisitor<'tcx>>(
&self,
goal: Goal<'tcx, ty::Predicate<'tcx>>,
depth: usize,
visitor: &mut V,
) -> V::Result {
let (_, proof_tree) =
<&SolverDelegate<'tcx>>::from(self).evaluate_root_goal(goal, GenerateProofTree::Yes);
let proof_tree = proof_tree.unwrap();
visitor.visit_goal(&InspectGoal::new(self, depth, proof_tree, None, GoalSource::Misc))
}
}