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