1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
use std::fmt;

use rustc_errors::ErrorGuaranteed;
use rustc_infer::infer::region_constraints::RegionConstraintData;
use rustc_middle::traits::query::NoSolution;
use rustc_middle::ty::{TyCtxt, TypeFoldable};
use rustc_span::Span;
use tracing::info;

use crate::infer::canonical::query_response;
use crate::infer::InferCtxt;
use crate::traits::query::type_op::TypeOpOutput;
use crate::traits::ObligationCtxt;

pub struct CustomTypeOp<F> {
    closure: F,
    description: &'static str,
}

impl<F> CustomTypeOp<F> {
    pub fn new<'tcx, R>(closure: F, description: &'static str) -> Self
    where
        F: FnOnce(&ObligationCtxt<'_, 'tcx>) -> Result<R, NoSolution>,
    {
        CustomTypeOp { closure, description }
    }
}

impl<'tcx, F, R> super::TypeOp<'tcx> for CustomTypeOp<F>
where
    F: FnOnce(&ObligationCtxt<'_, 'tcx>) -> Result<R, NoSolution>,
    R: fmt::Debug + TypeFoldable<TyCtxt<'tcx>>,
{
    type Output = R;
    /// We can't do any custom error reporting for `CustomTypeOp`, so
    /// we can use `!` to enforce that the implementation never provides it.
    type ErrorInfo = !;

    /// Processes the operation and all resulting obligations,
    /// returning the final result along with any region constraints
    /// (they will be given over to the NLL region solver).
    fn fully_perform(
        self,
        infcx: &InferCtxt<'tcx>,
        span: Span,
    ) -> Result<TypeOpOutput<'tcx, Self>, ErrorGuaranteed> {
        if cfg!(debug_assertions) {
            info!("fully_perform({:?})", self);
        }

        Ok(scrape_region_constraints(infcx, self.closure, self.description, span)?.0)
    }
}

impl<F> fmt::Debug for CustomTypeOp<F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.description.fmt(f)
    }
}

/// Executes `op` and then scrapes out all the "old style" region
/// constraints that result, creating query-region-constraints.
pub fn scrape_region_constraints<'tcx, Op, R>(
    infcx: &InferCtxt<'tcx>,
    op: impl FnOnce(&ObligationCtxt<'_, 'tcx>) -> Result<R, NoSolution>,
    name: &'static str,
    span: Span,
) -> Result<(TypeOpOutput<'tcx, Op>, RegionConstraintData<'tcx>), ErrorGuaranteed>
where
    R: TypeFoldable<TyCtxt<'tcx>>,
    Op: super::TypeOp<'tcx, Output = R>,
{
    // During NLL, we expect that nobody will register region
    // obligations **except** as part of a custom type op (and, at the
    // end of each custom type op, we scrape out the region
    // obligations that resulted). So this vector should be empty on
    // entry.
    let pre_obligations = infcx.take_registered_region_obligations();
    assert!(
        pre_obligations.is_empty(),
        "scrape_region_constraints: incoming region obligations = {pre_obligations:#?}",
    );

    let value = infcx.commit_if_ok(|_| {
        let ocx = ObligationCtxt::new(infcx);
        let value = op(&ocx).map_err(|_| {
            infcx.dcx().span_delayed_bug(span, format!("error performing operation: {name}"))
        })?;
        let errors = ocx.select_all_or_error();
        if errors.is_empty() {
            Ok(value)
        } else {
            Err(infcx
                .dcx()
                .delayed_bug(format!("errors selecting obligation during MIR typeck: {errors:?}")))
        }
    })?;

    // Next trait solver performs operations locally, and normalize goals should resolve vars.
    let value = infcx.resolve_vars_if_possible(value);

    let region_obligations = infcx.take_registered_region_obligations();
    let region_constraint_data = infcx.take_and_reset_region_constraints();
    let region_constraints = query_response::make_query_region_constraints(
        infcx.tcx,
        region_obligations
            .iter()
            .map(|r_o| (r_o.sup_type, r_o.sub_region, r_o.origin.to_constraint_category()))
            .map(|(ty, r, cc)| (infcx.resolve_vars_if_possible(ty), r, cc)),
        &region_constraint_data,
    );

    if region_constraints.is_empty() {
        Ok((
            TypeOpOutput { output: value, constraints: None, error_info: None },
            region_constraint_data,
        ))
    } else {
        Ok((
            TypeOpOutput {
                output: value,
                constraints: Some(infcx.tcx.arena.alloc(region_constraints)),
                error_info: None,
            },
            region_constraint_data,
        ))
    }
}