Skip to main content

flux_middle/rty/
refining.rs

1//! *Refining* is the process of generating a refined version of a rust type.
2//!
3//! Concretely, this module provides functions to go from types in [`flux_rustc_bridge::ty`] to
4//! types in [`rty`].
5
6use flux_arc_interner::{List, SliceInternable};
7use flux_common::bug;
8use flux_rustc_bridge::{ty, ty::GenericArgsExt as _};
9use itertools::Itertools;
10use rustc_abi::VariantIdx;
11use rustc_data_structures::fx::FxHashMap;
12use rustc_hir::def_id::DefId;
13use rustc_middle::ty::ParamTy;
14use rustc_span::Symbol;
15use rustc_type_ir::INNERMOST;
16
17use super::{
18    RefineArgsExt,
19    fold::{TypeFoldable, TypeFolder, TypeVisitable},
20};
21use crate::{
22    global_env::{GlobalEnv, WeakKvarInfo, WeakKvarMap},
23    queries::{QueryErr, QueryResult},
24    query_bug,
25    rty::{self, Expr, fold::TypeSuperFoldable},
26};
27
28pub fn refine_generics(generics: &ty::Generics) -> rty::Generics {
29    let params = generics
30        .params
31        .iter()
32        .map(|param| refine_generic_param_def(false, param))
33        .collect();
34
35    rty::Generics {
36        own_params: params,
37        parent: generics.parent(),
38        parent_count: generics.parent_count(),
39        has_self: generics.orig.has_self,
40    }
41}
42
43pub(crate) fn refine_generic_param_def(
44    as_type: bool,
45    param: &ty::GenericParamDef,
46) -> rty::GenericParamDef {
47    rty::GenericParamDef {
48        kind: refine_generic_param_def_kind(as_type, param.kind),
49        index: param.index,
50        name: param.name,
51        def_id: param.def_id,
52    }
53}
54
55fn refine_generic_param_def_kind(
56    as_type: bool,
57    kind: ty::GenericParamDefKind,
58) -> rty::GenericParamDefKind {
59    match kind {
60        ty::GenericParamDefKind::Lifetime => rty::GenericParamDefKind::Lifetime,
61        ty::GenericParamDefKind::Type { has_default } => {
62            if as_type {
63                rty::GenericParamDefKind::Type { has_default }
64            } else {
65                rty::GenericParamDefKind::Base { has_default }
66            }
67        }
68        ty::GenericParamDefKind::Const { has_default, .. } => {
69            rty::GenericParamDefKind::Const { has_default }
70        }
71    }
72}
73
74/// The type of a `static` without a flux spec: the default refinement of its rust type. Reads of
75/// such a static use this type too, so e.g. the field specs of a struct are assumed when reading
76/// from the static, and thus must be checked in its initializer.
77pub fn default_static_ty(genv: GlobalEnv, def_id: DefId) -> QueryResult<rty::Ty> {
78    genv.lower_type_of(def_id)?
79        .skip_binder()
80        .refine(&Refiner::default_for_item(genv, def_id)?)
81}
82
83pub struct Refiner<'genv, 'tcx> {
84    genv: GlobalEnv<'genv, 'tcx>,
85    def_id: DefId,
86    generics: rty::Generics,
87    refine: fn(rty::BaseTy) -> rty::SubsetTyCtor,
88}
89
90impl<'genv, 'tcx> Refiner<'genv, 'tcx> {
91    pub fn new_for_item(
92        genv: GlobalEnv<'genv, 'tcx>,
93        def_id: DefId,
94        refine: fn(rty::BaseTy) -> rty::SubsetTyCtor,
95    ) -> QueryResult<Self> {
96        let generics = genv.generics_of(def_id)?;
97        Ok(Self { genv, def_id, generics, refine })
98    }
99
100    pub fn default_for_item(genv: GlobalEnv<'genv, 'tcx>, def_id: DefId) -> QueryResult<Self> {
101        Self::new_for_item(genv, def_id, refine_default)
102    }
103
104    pub fn with_holes(genv: GlobalEnv<'genv, 'tcx>, def_id: DefId) -> QueryResult<Self> {
105        Self::new_for_item(genv, def_id, |bty| {
106            let sort = bty.sort();
107            let constr = rty::SubsetTy::new(
108                bty.shift_in_escaping(1),
109                rty::Expr::nu(),
110                rty::Expr::hole(rty::HoleKind::Pred),
111            );
112            rty::Binder::bind_with_sort(constr, sort)
113        })
114    }
115
116    pub fn refine<T: Refine + ?Sized>(&self, t: &T) -> QueryResult<T::Output> {
117        t.refine(self)
118    }
119
120    fn refine_existential_predicate_generic_args(
121        &self,
122        def_id: DefId,
123        args: &ty::GenericArgs,
124    ) -> QueryResult<rty::GenericArgs> {
125        let generics = self.generics_of(def_id)?;
126        args.iter()
127            .enumerate()
128            .map(|(idx, arg)| {
129                // We need to skip the generic for Self
130                let param = generics.param_at(idx + 1, self.genv)?;
131                self.refine_generic_arg(&param, arg)
132            })
133            .try_collect()
134    }
135
136    pub fn refine_variant_def(
137        &self,
138        adt_def_id: DefId,
139        variant_idx: VariantIdx,
140    ) -> QueryResult<rty::PolyVariant> {
141        let adt_def = self.adt_def(adt_def_id)?;
142        let variant_def = adt_def.variant(variant_idx);
143        let fields = variant_def
144            .fields
145            .iter()
146            .map(|fld| {
147                let ty = self.genv.lower_type_of(fld.did)?.instantiate_identity();
148                ty.refine(self)
149            })
150            .try_collect()?;
151
152        let idx = if adt_def.sort_def().is_struct() {
153            rty::Expr::unit_struct(adt_def_id)
154        } else {
155            rty::Expr::ctor_enum(adt_def_id, variant_idx)
156        };
157        let value = rty::VariantSig::new(
158            adt_def,
159            rty::GenericArg::identity_for_item(self.genv, adt_def_id)?,
160            fields,
161            idx,
162            List::empty(),
163        );
164
165        Ok(rty::Binder::bind_with_vars(value, List::empty()))
166    }
167
168    pub fn refine_generic_args(
169        &self,
170        def_id: DefId,
171        args: &ty::GenericArgs,
172    ) -> QueryResult<rty::GenericArgs> {
173        let generics = self.generics_of(def_id)?;
174        args.iter()
175            .enumerate()
176            .map(|(idx, arg)| {
177                let param = generics.param_at(idx, self.genv)?;
178                self.refine_generic_arg(&param, arg)
179            })
180            .collect()
181    }
182
183    pub fn refine_generic_arg(
184        &self,
185        param: &rty::GenericParamDef,
186        arg: &ty::GenericArg,
187    ) -> QueryResult<rty::GenericArg> {
188        match (&param.kind, arg) {
189            (rty::GenericParamDefKind::Type { .. }, ty::GenericArg::Ty(ty)) => {
190                Ok(rty::GenericArg::Ty(ty.refine(self)?))
191            }
192            (rty::GenericParamDefKind::Base { .. }, ty::GenericArg::Ty(ty)) => {
193                let rty::TyOrBase::Base(contr) = self.refine_ty_or_base(ty)? else {
194                    return Err(QueryErr::InvalidGenericArg { def_id: param.def_id });
195                };
196                Ok(rty::GenericArg::Base(contr))
197            }
198            (rty::GenericParamDefKind::Lifetime, ty::GenericArg::Lifetime(re)) => {
199                Ok(rty::GenericArg::Lifetime(*re))
200            }
201            (rty::GenericParamDefKind::Const { .. }, ty::GenericArg::Const(ct)) => {
202                Ok(rty::GenericArg::Const(ct.clone()))
203            }
204            _ => bug!("mismatched generic arg `{arg:?}` `{param:?}`"),
205        }
206    }
207
208    fn refine_alias_term(&self, alias_term: &ty::AliasTerm) -> QueryResult<rty::AliasTerm> {
209        let args = self.refine_generic_args(alias_term.def_id(), &alias_term.args)?;
210        Ok(rty::AliasTerm::new(alias_term.kind, args))
211    }
212
213    fn refine_alias_ty(&self, alias_ty: &ty::AliasTy) -> QueryResult<rty::AliasTy> {
214        match alias_ty.kind {
215            // Only opaque types carry refinement arguments
216            ty::AliasKind::Opaque { def_id } => {
217                let args = self.refine_generic_args(def_id, &alias_ty.args)?;
218                let refine_args = rty::RefineArgs::for_item(self.genv, def_id, |param, _| {
219                    let param = param.instantiate(self.genv.tcx(), &args, &[]);
220                    Ok(rty::Expr::hole(rty::HoleKind::Expr(param.sort)))
221                })?;
222                Ok(rty::AliasTy::new(alias_ty.kind, args, refine_args))
223            }
224            ty::AliasKind::Projection { def_id } | ty::AliasKind::Free { def_id } => {
225                let args = self.refine_generic_args(def_id, &alias_ty.args)?;
226                Ok(rty::AliasTy::new(alias_ty.kind, args, List::empty()))
227            }
228        }
229    }
230
231    pub fn refine_ty_or_base(&self, ty: &ty::Ty) -> QueryResult<rty::TyOrBase> {
232        let bty = match ty.kind() {
233            ty::TyKind::Closure(did, args) => {
234                let no_panic = self.genv.no_panic(*did);
235                let closure_args = args.as_closure();
236                let upvar_tys = closure_args
237                    .upvar_tys()
238                    .iter()
239                    .map(|ty| ty.refine(self))
240                    .try_collect()?;
241                rty::BaseTy::Closure(*did, upvar_tys, args.clone(), no_panic)
242            }
243            ty::TyKind::Coroutine(did, args) => {
244                let coroutine_args = args.as_coroutine();
245                let resume_ty = coroutine_args.resume_ty().refine(self)?;
246                let upvar_tys = coroutine_args
247                    .upvar_tys()
248                    .map(|ty| ty.refine(self))
249                    .try_collect()?;
250                rty::BaseTy::Coroutine(*did, resume_ty, upvar_tys, args.clone())
251            }
252            ty::TyKind::CoroutineWitness(..) => {
253                bug!("implement when we know what this is");
254            }
255            ty::TyKind::Never => rty::BaseTy::Never,
256            ty::TyKind::Ref(r, ty, mutbl) => rty::BaseTy::Ref(*r, ty.refine(self)?, *mutbl),
257            ty::TyKind::Float(float_ty) => rty::BaseTy::Float(*float_ty),
258            ty::TyKind::Tuple(tys) => {
259                let tys = tys.iter().map(|ty| ty.refine(self)).try_collect()?;
260                rty::BaseTy::Tuple(tys)
261            }
262            ty::TyKind::Array(ty, len) => rty::BaseTy::Array(ty.refine(self)?, len.clone()),
263            ty::TyKind::Param(param_ty) => {
264                match self.param(*param_ty)?.kind {
265                    rty::GenericParamDefKind::Type { .. } => {
266                        return Ok(rty::TyOrBase::Ty(rty::Ty::param(*param_ty)));
267                    }
268                    rty::GenericParamDefKind::Base { .. } => rty::BaseTy::Param(*param_ty),
269                    rty::GenericParamDefKind::Lifetime | rty::GenericParamDefKind::Const { .. } => {
270                        bug!()
271                    }
272                }
273            }
274            ty::TyKind::Adt(adt_def, args) => {
275                let adt_def = self.genv.adt_def(adt_def.did())?;
276                let args = self.refine_generic_args(adt_def.did(), args)?;
277                rty::BaseTy::adt(adt_def, args)
278            }
279            ty::TyKind::FnDef(def_id, args) => {
280                let args = self.refine_generic_args(*def_id, args)?;
281                rty::BaseTy::fn_def(*def_id, args)
282            }
283            ty::TyKind::Alias(alias_ty) => {
284                let alias_ty = self.as_default().refine_alias_ty(alias_ty)?;
285                rty::BaseTy::Alias(alias_ty)
286            }
287            ty::TyKind::Bool => rty::BaseTy::Bool,
288            ty::TyKind::Int(int_ty) => rty::BaseTy::Int(*int_ty),
289            ty::TyKind::Uint(uint_ty) => rty::BaseTy::Uint(*uint_ty),
290            ty::TyKind::Foreign(def_id) => rty::BaseTy::Foreign(*def_id),
291            ty::TyKind::Str => rty::BaseTy::Str,
292            ty::TyKind::Slice(ty) => rty::BaseTy::Slice(ty.refine(self)?),
293            ty::TyKind::Char => rty::BaseTy::Char,
294            ty::TyKind::FnPtr(poly_fn_sig) => {
295                rty::BaseTy::FnPtr(poly_fn_sig.refine(&self.as_default())?)
296            }
297            ty::TyKind::RawPtr(ty, mu) => rty::BaseTy::RawPtr(ty.refine(&self.as_default())?, *mu),
298            ty::TyKind::Dynamic(exi_preds, r) => {
299                let exi_preds = exi_preds
300                    .iter()
301                    .map(|pred| pred.refine(self))
302                    .try_collect()?;
303                rty::BaseTy::Dynamic(exi_preds, *r)
304            }
305            ty::TyKind::Pat => rty::BaseTy::Pat,
306        };
307        Ok(rty::TyOrBase::Base((self.refine)(bty)))
308    }
309
310    fn as_default(&self) -> Self {
311        Refiner { refine: refine_default, generics: self.generics.clone(), ..*self }
312    }
313
314    fn adt_def(&self, def_id: DefId) -> QueryResult<rty::AdtDef> {
315        self.genv.adt_def(def_id)
316    }
317
318    fn generics_of(&self, def_id: DefId) -> QueryResult<rty::Generics> {
319        self.genv.generics_of(def_id)
320    }
321
322    fn param(&self, param_ty: ParamTy) -> QueryResult<rty::GenericParamDef> {
323        self.generics.param_at(param_ty.index as usize, self.genv)
324    }
325}
326
327pub trait Refine {
328    type Output;
329
330    fn refine(&self, refiner: &Refiner) -> QueryResult<Self::Output>;
331}
332
333impl Refine for ty::Ty {
334    type Output = rty::Ty;
335
336    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::Ty> {
337        Ok(refiner.refine_ty_or_base(self)?.into_ty())
338    }
339}
340
341impl<T: Refine> Refine for ty::Binder<T> {
342    type Output = rty::Binder<T::Output>;
343
344    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::Binder<T::Output>> {
345        let vars = refine_bound_variables(self.vars());
346        let inner = self.skip_binder_ref().refine(refiner)?;
347        Ok(rty::Binder::bind_with_vars(inner, vars))
348    }
349}
350
351impl Refine for ty::FnSig {
352    type Output = rty::FnSig;
353
354    // TODO(hof2)
355    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::FnSig> {
356        let inputs = self
357            .inputs()
358            .iter()
359            .map(|ty| ty.refine(refiner))
360            .try_collect()?;
361        let ret = self.output().refine(refiner)?.shift_in_escaping(1);
362        let output = rty::Binder::bind_with_vars(rty::FnOutput::new(ret, vec![]), List::empty());
363        // TODO(hof2) make a hoister to hoist all the stuff out of the inputs,
364        // the hoister will have a list of all the variables it hoisted and the
365        // single hole for the "requires"; then we "fill" the hole with a KVAR
366        // and generate a PolyFnSig with the hoisted variables
367        // see `into_bb_env` in `type_env.rs` for an example.
368        Ok(rty::FnSig::new(self.safety, self.abi, List::empty(), inputs, output, Expr::ff(), true))
369    }
370}
371
372impl Refine for ty::Clause {
373    type Output = rty::Clause;
374
375    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::Clause> {
376        Ok(rty::Clause { kind: self.kind.refine(refiner)? })
377    }
378}
379
380impl Refine for ty::TraitRef {
381    type Output = rty::TraitRef;
382
383    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::TraitRef> {
384        Ok(rty::TraitRef {
385            def_id: self.def_id,
386            args: refiner.refine_generic_args(self.def_id, &self.args)?,
387        })
388    }
389}
390
391impl Refine for ty::ClauseKind {
392    type Output = rty::ClauseKind;
393
394    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::ClauseKind> {
395        let kind = match self {
396            ty::ClauseKind::Trait(trait_pred) => {
397                let pred = rty::TraitPredicate { trait_ref: trait_pred.trait_ref.refine(refiner)? };
398                rty::ClauseKind::Trait(pred)
399            }
400            ty::ClauseKind::Projection(proj_pred) => {
401                let rty::TyOrBase::Base(term) = refiner.refine_ty_or_base(&proj_pred.term)? else {
402                    return Err(query_bug!(
403                        refiner.def_id,
404                        "sorry, we can't handle non-base associated types"
405                    ));
406                };
407                let pred = rty::ProjectionPredicate {
408                    projection_term: refiner.refine_alias_term(&proj_pred.projection_term)?,
409                    term,
410                };
411                rty::ClauseKind::Projection(pred)
412            }
413            ty::ClauseKind::RegionOutlives(pred) => {
414                let pred = rty::OutlivesPredicate(pred.0, pred.1);
415                rty::ClauseKind::RegionOutlives(pred)
416            }
417            ty::ClauseKind::TypeOutlives(pred) => {
418                let pred = rty::OutlivesPredicate(pred.0.refine(refiner)?, pred.1);
419                rty::ClauseKind::TypeOutlives(pred)
420            }
421            ty::ClauseKind::ConstArgHasType(const_, ty) => {
422                rty::ClauseKind::ConstArgHasType(const_.clone(), ty.refine(&refiner.as_default())?)
423            }
424            ty::ClauseKind::UnstableFeature(sym) => rty::ClauseKind::UnstableFeature(*sym),
425        };
426        Ok(kind)
427    }
428}
429
430impl Refine for ty::ExistentialPredicate {
431    type Output = rty::ExistentialPredicate;
432
433    fn refine(&self, refiner: &Refiner) -> QueryResult<Self::Output> {
434        let pred = match self {
435            ty::ExistentialPredicate::Trait(trait_ref) => {
436                rty::ExistentialPredicate::Trait(rty::ExistentialTraitRef {
437                    def_id: trait_ref.def_id,
438                    args: refiner.refine_existential_predicate_generic_args(
439                        trait_ref.def_id,
440                        &trait_ref.args,
441                    )?,
442                })
443            }
444            ty::ExistentialPredicate::Projection(projection) => {
445                let rty::TyOrBase::Base(term) = refiner.refine_ty_or_base(&projection.term)? else {
446                    return Err(query_bug!(
447                        refiner.def_id,
448                        "sorry, we can't handle non-base associated types"
449                    ));
450                };
451                rty::ExistentialPredicate::Projection(rty::ExistentialProjection {
452                    def_id: projection.def_id,
453                    args: refiner.refine_existential_predicate_generic_args(
454                        projection.def_id,
455                        &projection.args,
456                    )?,
457                    term,
458                })
459            }
460            ty::ExistentialPredicate::AutoTrait(def_id) => {
461                rty::ExistentialPredicate::AutoTrait(*def_id)
462            }
463        };
464        Ok(pred)
465    }
466}
467
468impl Refine for ty::GenericPredicates {
469    type Output = rty::GenericPredicates;
470
471    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::GenericPredicates> {
472        Ok(rty::GenericPredicates {
473            parent: self.parent,
474            predicates: refiner.refine(&self.predicates)?,
475        })
476    }
477}
478
479impl<T> Refine for List<T>
480where
481    T: SliceInternable,
482    T: Refine<Output: SliceInternable>,
483{
484    type Output = rty::List<T::Output>;
485
486    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::List<T::Output>> {
487        refiner.refine(&self[..])
488    }
489}
490
491impl<T> Refine for [T]
492where
493    T: Refine<Output: SliceInternable>,
494{
495    type Output = rty::List<T::Output>;
496
497    fn refine(&self, refiner: &Refiner) -> QueryResult<rty::List<T::Output>> {
498        self.iter().map(|t| refiner.refine(t)).try_collect()
499    }
500}
501
502fn refine_default(bty: rty::BaseTy) -> rty::SubsetTyCtor {
503    let sort = bty.sort();
504    let constr = rty::SubsetTy::trivial(bty.shift_in_escaping(1), rty::Expr::nu());
505    rty::Binder::bind_with_sort(constr, sort)
506}
507
508pub fn refine_bound_variables(vars: &[ty::BoundVariableKind]) -> List<rty::BoundVariableKind> {
509    vars.iter()
510        .map(|kind| {
511            match kind {
512                ty::BoundVariableKind::Region(kind) => rty::BoundVariableKind::Region(*kind),
513            }
514        })
515        .collect()
516}
517
518impl rty::PolyFnSig {
519    pub fn add_weak_kvars(self, genv: GlobalEnv, def_id: DefId) -> QueryResult<Self> {
520        let refinement_generics = genv.refinement_generics_of(def_id)?;
521        let early_param_sorts: FxHashMap<Symbol, rty::Sort> = refinement_generics
522            .0
523            .own_params
524            .iter()
525            .map(|param| (param.name, param.sort.clone()))
526            .collect();
527        let early_vars = self
528            .early_params()
529            .into_iter()
530            .filter_map(|param| {
531                let sort = early_param_sorts.get(&param.name).unwrap().clone();
532                if is_weak_kvar_sort(&sort) {
533                    Some((rty::Var::EarlyParam(param), sort))
534                } else {
535                    None
536                }
537            })
538            .collect_vec();
539        let late_vars = make_vars_and_sorts_from_bound_vars(self.vars());
540        Ok(self.map(|fn_sig| {
541            let mut params = late_vars.into_iter().chain(early_vars).collect_vec();
542            let mut wkvar_inserter = WeakKVarInserter {
543                wkvar_map: WeakKvarMap::default(),
544                def_id,
545                kvid: rty::KVid::from(0_usize),
546                existential_params: Vec::new(),
547                params: params.clone(),
548            };
549            let requires_wkvar = make_weak_kvar(
550                &mut wkvar_inserter.wkvar_map,
551                def_id,
552                &mut wkvar_inserter.kvid,
553                Vec::new(),
554                params.clone(),
555            );
556            let inputs = fn_sig
557                .inputs
558                .iter()
559                .map(|input| wkvar_inserter.fold_ty(input))
560                .collect();
561            shift_in_vars(&mut params);
562            let output_binder_params = make_vars_and_sorts_from_bound_vars(fn_sig.output.vars());
563            params.extend(output_binder_params);
564            wkvar_inserter.params = params.clone();
565            let ensures = if !fn_sig.output.vars().is_empty()
566                && let Some(ensures_wkvar) = make_weak_kvar(
567                    &mut wkvar_inserter.wkvar_map,
568                    def_id,
569                    &mut wkvar_inserter.kvid,
570                    make_vars_and_sorts_from_bound_vars(fn_sig.output.vars()),
571                    params.clone(),
572                ) {
573                fn_sig
574                    .output
575                    .skip_binder_ref()
576                    .ensures
577                    .iter()
578                    .cloned()
579                    .chain(std::iter::once(rty::Ensures::Pred(rty::Expr::wkvar(ensures_wkvar))))
580                    .collect()
581            } else {
582                fn_sig.output.skip_binder_ref().ensures.clone()
583            };
584            let output = fn_sig
585                .output
586                .map(|output| rty::FnOutput { ret: wkvar_inserter.fold_ty(&output.ret), ensures });
587            genv.feed_weak_kvars(def_id, wkvar_inserter.wkvar_map);
588
589            rty::FnSig {
590                abi: fn_sig.abi,
591                safety: fn_sig.safety,
592                inputs,
593                // NOTE(CK): Not sure whether we can avoid the clone.
594                requires: if let Some(requires_wkvar) = requires_wkvar {
595                    fn_sig
596                        .requires
597                        .iter()
598                        .cloned()
599                        .chain(std::iter::once(rty::Expr::wkvar(requires_wkvar)))
600                        .collect()
601                } else {
602                    fn_sig.requires
603                },
604                output,
605                lifted: fn_sig.lifted,
606                no_panic: fn_sig.no_panic,
607            }
608        }))
609    }
610}
611
612struct WeakKVarInserter {
613    wkvar_map: WeakKvarMap,
614    def_id: DefId,
615    kvid: rty::KVid,
616    existential_params: Vec<Vec<(rty::Var, rty::Sort)>>,
617    params: Vec<(rty::Var, rty::Sort)>,
618}
619
620impl TypeFolder for WeakKVarInserter {
621    fn fold_ty(&mut self, ty: &rty::Ty) -> rty::Ty {
622        use rty::{Expr, Ty, TyKind::*};
623        match ty.kind() {
624            // This is the only recursive case where we need to update the params
625            // since we're going under a binder.
626            //
627            // We handle the shifting in and out explicitly rather than using
628            // the enter_binder and exit_binder methods because we immediately
629            // use the bound vars to make a weak kvar.
630            Exists(bound_ty) => {
631                for v in &mut self.existential_params {
632                    shift_in_vars(v);
633                }
634                shift_in_vars(&mut self.params);
635                let exist_params = make_vars_and_sorts_from_bound_vars(bound_ty.vars());
636                // Take all of the current existential params + the current params,
637                // AFTER shifting in.
638                let params = self
639                    .existential_params
640                    .iter()
641                    .flatten()
642                    .chain(self.params.iter())
643                    .cloned()
644                    .collect();
645                // We pass the params immediately under this binder as the self args.
646                //
647                // The purpose of self args is to ensure that we don't have duplication
648                // of suggestions.
649                //
650                // Suppose after we add weak kvars we have the type
651                //
652                //     fn ({exists v0. Vec<i32>[v0] | $wk1[v0]()}) requires $wk0[]()
653                //
654                // If we are looking to instantiate a weak kvar to the
655                // expression `2 > 1` (for some reason), we can validly put it
656                // in both $wk0 and $wk1. But the self arg ensures that we don't
657                // put it in $wk1, since it requires the expression contain one
658                // of its self args (in this case, just `v0`).
659                let wkvar = make_weak_kvar(
660                    &mut self.wkvar_map,
661                    self.def_id,
662                    &mut self.kvid,
663                    exist_params.clone(),
664                    params,
665                );
666                // Now we add the params for future weak kvars.
667                self.existential_params.push(exist_params);
668                let new_ty = bound_ty.skip_binder_ref().super_fold_with(self);
669                self.existential_params.pop();
670                for v in &mut self.existential_params {
671                    shift_out_vars(v);
672                }
673                shift_out_vars(&mut self.params);
674                Ty::exists(rty::Binder::bind_with_vars(
675                    if let Some(wkvar) = wkvar {
676                        Ty::constr(Expr::wkvar(wkvar), new_ty)
677                    } else {
678                        new_ty
679                    },
680                    bound_ty.vars().clone(),
681                ))
682            }
683            _ => ty.super_fold_with(self),
684        }
685    }
686
687    fn fold_bty(&mut self, bty: &rty::BaseTy) -> rty::BaseTy {
688        use rty::{BaseTy, Expr, GenericArg};
689        match bty {
690            BaseTy::Adt(adt_def, args) => {
691                let new_args = args
692                    .iter()
693                    .map(|arg| {
694                        match arg {
695                            GenericArg::Base(subset_ty) => {
696                                for v in &mut self.existential_params {
697                                    shift_in_vars(v);
698                                }
699                                shift_in_vars(&mut self.params);
700                                let exist_params =
701                                    make_vars_and_sorts_from_bound_vars(subset_ty.vars());
702                                // Take all of the current existential params + the current params,
703                                // AFTER shifting in.
704                                let params = self
705                                    .existential_params
706                                    .iter()
707                                    .flatten()
708                                    .chain(self.params.iter())
709                                    .cloned()
710                                    .collect();
711                                // We pass the params immediately under this binder as the self args.
712                                // see the TyKind::Exists case.
713                                let wkvar = make_weak_kvar(
714                                    &mut self.wkvar_map,
715                                    self.def_id,
716                                    &mut self.kvid,
717                                    exist_params.clone(),
718                                    params,
719                                );
720                                // Now we add the params for future weak kvars.
721                                self.existential_params.push(exist_params);
722                                let new_ty = subset_ty.skip_binder_ref().super_fold_with(self);
723                                let new_ty_maybe_with_wkvar = if let Some(wkvar) = wkvar {
724                                    new_ty.strengthen(Expr::wkvar(wkvar))
725                                } else {
726                                    new_ty
727                                };
728                                self.existential_params.pop();
729                                for v in &mut self.existential_params {
730                                    shift_out_vars(v);
731                                }
732                                shift_out_vars(&mut self.params);
733                                GenericArg::Base(rty::Binder::bind_with_vars(
734                                    new_ty_maybe_with_wkvar,
735                                    subset_ty.vars().clone(),
736                                ))
737                            }
738                            _ => arg.fold_with(self),
739                        }
740                    })
741                    .collect();
742                BaseTy::Adt(adt_def.clone(), new_args)
743            }
744            // For these specific btys, we will recur and add wkvars
745            BaseTy::Ref(..) | BaseTy::Tuple(..) | BaseTy::Array(..) | BaseTy::Slice(..) => {
746                bty.super_fold_with(self)
747            }
748            // By default we will not recur on the bty to add wkvars
749            _ => bty.clone(),
750        }
751    }
752
753    fn fold_expr(&mut self, expr: &Expr) -> Expr {
754        expr.clone()
755    }
756
757    fn fold_sort(&mut self, sort: &rty::Sort) -> rty::Sort {
758        sort.clone()
759    }
760}
761
762/// NOTE(CK):
763///   * Skips params (we don't presently handle polymorphism, though even if we did,
764///     I'm not sure that we need to pass params to the weak kvars).
765///   * Skips locs because we can't encode those.
766///   * Skips aliases: a projection can normalize to a unit sort at a call-site,
767///     while the weak kvar's declaration would still encode it as an integer.
768///   * Skips unit + unit adts because they otherwise get encoded as a 0 tuple
769///     to fixpoint because we use them in the args to a weak kvar, which
770///     we don't want to do.
771fn is_weak_kvar_sort(sort: &rty::Sort) -> bool {
772    !sort.is_param()
773        && !matches!(sort, rty::Sort::Alias(..))
774        && !sort.is_loc()
775        && !sort.is_unit()
776        && sort.is_unit_adt().is_none()
777}
778
779fn make_vars_and_sorts_from_bound_vars<'a, I, II>(vars: I) -> Vec<(rty::Var, rty::Sort)>
780where
781    I: IntoIterator<IntoIter = II>,
782    II: DoubleEndedIterator<Item = &'a rty::BoundVariableKind>,
783{
784    vars.into_iter()
785        .enumerate()
786        .filter_map(|(i, var_kind)| {
787            if let rty::BoundVariableKind::Refine(sort, _, reft_kind) = var_kind
788                && is_weak_kvar_sort(sort)
789            {
790                let bound_reft = rty::BoundReft { var: rty::BoundVar::from(i), kind: *reft_kind };
791                Some((rty::Var::Bound(INNERMOST, bound_reft), sort.clone()))
792            } else {
793                None
794            }
795        })
796        .collect_vec()
797}
798
799// TODO: Use a Vec<Vec<_>> solution, per Nico.
800// This is a sort of annoying rearchitecture, but nothing impossible.
801fn shift_in_vars(vars: &mut [(rty::Var, rty::Sort)]) {
802    for (var, _) in vars.iter_mut() {
803        *var = var.shift_in(1);
804    }
805}
806
807fn shift_out_vars(vars: &mut [(rty::Var, rty::Sort)]) {
808    for (var, _) in vars.iter_mut() {
809        *var = var.shift_out(1);
810    }
811}
812
813// TODO: Don't make a weak kvar if the self_args is empty if there's a weak kvar
814//        that's been created before it with a superset of its params.
815fn make_weak_kvar(
816    wkvar_map: &mut WeakKvarMap,
817    def_id: DefId,
818    kvid: &mut rty::KVid,
819    self_args: Vec<(rty::Var, rty::Sort)>,
820    params: Vec<(rty::Var, rty::Sort)>,
821) -> Option<rty::WKVar> {
822    if self_args.is_empty() && params.is_empty() {
823        return None;
824    }
825    let num_self_args = self_args.len();
826    let (args, sorts): (Vec<rty::Var>, Vec<rty::Sort>) =
827        self_args.into_iter().chain(params).unzip();
828    let arg_exprs = args.into_iter().map(rty::Expr::var).collect();
829    // We don't have any solutions because these weak kvars are being generated
830    // (solutions only come from user annotations).
831    wkvar_map.insert(kvid.as_u32(), WeakKvarInfo { solutions: vec![], sorts });
832    let ret = rty::WKVar {
833        wkvid: rty::WKVid::new(def_id, *kvid),
834        self_args: num_self_args,
835        args: arg_exprs,
836    };
837    *kvid += 1;
838    Some(ret)
839}