flux_refineck/
type_env.rs

1mod place_ty;
2
3use std::{iter, ops::ControlFlow};
4
5use flux_common::{bug, dbg::debug_assert_eq3, tracked_span_bug, tracked_span_dbg_assert_eq};
6use flux_infer::{
7    fixpoint_encoding::KVarEncoding,
8    infer::{ConstrReason, InferCtxt, InferCtxtAt, InferCtxtRoot, InferResult},
9    refine_tree::Scope,
10};
11use flux_macros::DebugAsJson;
12use flux_middle::{
13    PlaceExt as _,
14    global_env::GlobalEnv,
15    pretty::{PrettyCx, PrettyNested},
16    queries::QueryResult,
17    rty::{
18        BaseTy, Binder, BoundReftKind, Ctor, Ensures, Expr, ExprKind, FnSig, GenericArg, HoleKind,
19        INNERMOST, Lambda, List, Loc, Mutability, Path, PtrKind, Region, SortCtor, SubsetTy, Ty,
20        TyKind, VariantIdx,
21        canonicalize::{Hoister, LocalHoister},
22        fold::{FallibleTypeFolder, TypeFoldable, TypeVisitable, TypeVisitor},
23        region_matching::{rty_match_regions, ty_match_regions},
24    },
25};
26use flux_rustc_bridge::{
27    self,
28    mir::{BasicBlock, Body, Local, LocalDecl, LocalDecls, Place, PlaceElem},
29    ty,
30};
31use itertools::{Itertools, izip};
32use rustc_data_structures::unord::UnordMap;
33use rustc_index::{IndexSlice, IndexVec};
34use rustc_middle::{mir::RETURN_PLACE, ty::TyCtxt};
35use rustc_span::{Span, Symbol};
36use rustc_type_ir::BoundVar;
37use serde::Serialize;
38
39use self::place_ty::{LocKind, PlacesTree};
40use super::rty::Sort;
41
42#[derive(Clone, Default)]
43pub struct TypeEnv<'a> {
44    bindings: PlacesTree,
45    local_decls: &'a LocalDecls,
46}
47
48pub struct BasicBlockEnvShape {
49    scope: Scope,
50    bindings: PlacesTree,
51}
52
53pub struct BasicBlockEnv {
54    data: Binder<BasicBlockEnvData>,
55    scope: Scope,
56}
57
58#[derive(Debug)]
59struct BasicBlockEnvData {
60    constrs: List<Expr>,
61    bindings: PlacesTree,
62}
63
64impl<'a> TypeEnv<'a> {
65    pub fn new(infcx: &mut InferCtxt, body: &'a Body, fn_sig: &FnSig) -> TypeEnv<'a> {
66        let mut env = TypeEnv { bindings: PlacesTree::default(), local_decls: &body.local_decls };
67
68        for requires in fn_sig.requires() {
69            infcx.assume_pred(requires);
70        }
71
72        for (local, ty) in body.args_iter().zip(fn_sig.inputs()) {
73            let ty = infcx.unpack(ty);
74            infcx.assume_invariants(&ty);
75            env.alloc_with_ty(local, ty);
76        }
77
78        for local in body.vars_and_temps_iter() {
79            env.alloc(local);
80        }
81
82        env.alloc(RETURN_PLACE);
83        env
84    }
85
86    pub fn empty() -> TypeEnv<'a> {
87        TypeEnv { bindings: PlacesTree::default(), local_decls: IndexSlice::empty() }
88    }
89
90    fn alloc_with_ty(&mut self, local: Local, ty: Ty) {
91        let ty = ty_match_regions(&ty, &self.local_decls[local].ty);
92        self.bindings.insert(local.into(), LocKind::Local, ty);
93    }
94
95    fn alloc(&mut self, local: Local) {
96        self.bindings
97            .insert(local.into(), LocKind::Local, Ty::uninit());
98    }
99
100    pub(crate) fn into_infer(self, scope: Scope) -> BasicBlockEnvShape {
101        BasicBlockEnvShape::new(scope, self)
102    }
103
104    pub(crate) fn lookup_rust_ty(&self, genv: GlobalEnv, place: &Place) -> QueryResult<ty::Ty> {
105        Ok(place.ty(genv, self.local_decls)?.ty)
106    }
107
108    pub(crate) fn lookup_place(
109        &mut self,
110        infcx: &mut InferCtxtAt,
111        place: &Place,
112    ) -> InferResult<Ty> {
113        let span = infcx.span;
114        Ok(self.bindings.lookup_unfolding(infcx, place, span)?.ty)
115    }
116
117    pub(crate) fn get(&self, path: &Path) -> Ty {
118        self.bindings.get(path)
119    }
120
121    pub fn update_path(&mut self, path: &Path, new_ty: Ty, span: Span) {
122        self.bindings.lookup(path, span).update(new_ty);
123    }
124
125    /// When checking a borrow in the right hand side of an assignment `x = &'?n p`, we use the
126    /// annotated region `'?n` in the type of the result. This region will only be used temporarily
127    /// and then replaced by the region in the type of `x` after the assignment. See [`TypeEnv::assign`]
128    pub(crate) fn borrow(
129        &mut self,
130        infcx: &mut InferCtxtAt,
131        re: Region,
132        mutbl: Mutability,
133        place: &Place,
134    ) -> InferResult<Ty> {
135        let span = infcx.span;
136        let result = self.bindings.lookup_unfolding(infcx, place, span)?;
137        if result.is_strg && mutbl == Mutability::Mut {
138            Ok(Ty::ptr(PtrKind::Mut(re), result.path()))
139        } else {
140            // FIXME(nilehmann) we should block the place here. That would require a notion
141            // of shared vs mutable block types because sometimes blocked places from a shared
142            // reference never get unblocked and we should still allow reads through them.
143            Ok(Ty::mk_ref(re, result.ty, mutbl))
144        }
145    }
146
147    // FIXME(nilehmann) this is only used in a single place and we have it because [`TypeEnv`]
148    // doesn't expose a lookup without unfolding
149    pub(crate) fn ptr_to_ref_at_place(
150        &mut self,
151        infcx: &mut InferCtxtAt,
152        place: &Place,
153    ) -> InferResult {
154        let lookup = self.bindings.lookup(place, infcx.span);
155        let TyKind::Ptr(PtrKind::Mut(re), path) = lookup.ty.kind() else {
156            tracked_span_bug!("ptr_to_borrow called on non mutable pointer type")
157        };
158
159        let ref_ty =
160            self.ptr_to_ref(infcx, ConstrReason::Other, *re, path, PtrToRefBound::Infer)?;
161
162        self.bindings.lookup(place, infcx.span).update(ref_ty);
163
164        Ok(())
165    }
166
167    /// Convert a (strong) pointer to a mutable reference.
168    ///
169    /// This roughly implements the following inference rule:
170    /// ```text
171    ///                   t₁ <: t₂
172    /// -------------------------------------------------
173    /// Γ₁,ℓ:t1,Γ₂ ; ptr(mut, ℓ) => Γ₁,ℓ:†t₂,Γ₂ ; &mut t2
174    /// ```
175    /// That's it, we first get the current type `t₁` at location `ℓ` and check it is a subtype
176    /// of `t₂`. Then, we update the type of `ℓ` to `t₂` and block the place.
177    ///
178    /// The bound `t₂` can be either inferred ([`PtrToRefBound::Infer`]), explicitly provided
179    /// ([`PtrToRefBound::Ty`]), or made equal to `t₁` ([`PtrToRefBound::Identity`]).
180    ///
181    /// As an example, consider the environment `x: i32[a]` and the pointer `ptr(mut, x)`.
182    /// Converting the pointer to a mutable reference with an inferred bound produces the following
183    /// derivation (roughly):
184    ///
185    /// ```text
186    ///                    i32[a] <: i32{v: $k(v)}
187    /// ----------------------------------------------------------------
188    /// x: i32[a] ; ptr(mut, x) => x:†i32{v: $k(v)} ; &mut i32{v: $k(v)}
189    /// ```
190    pub(crate) fn ptr_to_ref(
191        &mut self,
192        infcx: &mut InferCtxtAt,
193        reason: ConstrReason,
194        re: Region,
195        path: &Path,
196        bound: PtrToRefBound,
197    ) -> InferResult<Ty> {
198        // ℓ: t1
199        let t1 = self.bindings.lookup(path, infcx.span).fold(infcx)?;
200
201        // t1 <: t2
202        let t2 = match bound {
203            PtrToRefBound::Ty(t2) => {
204                let t2 = rty_match_regions(&t2, &t1);
205                infcx.subtyping(&t1, &t2, reason)?;
206                t2
207            }
208            PtrToRefBound::Infer => {
209                let t2 = t1.with_holes().replace_holes(|sorts, kind| {
210                    debug_assert_eq!(kind, HoleKind::Pred);
211                    infcx.fresh_kvar(sorts, KVarEncoding::Conj)
212                });
213                infcx.subtyping(&t1, &t2, reason)?;
214                t2
215            }
216            PtrToRefBound::Identity => t1.clone(),
217        };
218
219        // ℓ: †t2
220        self.bindings
221            .lookup(path, infcx.span)
222            .block_with(t2.clone());
223
224        Ok(Ty::mk_ref(re, t2, Mutability::Mut))
225    }
226
227    pub(crate) fn fold_local_ptrs(&mut self, infcx: &mut InferCtxtAt) -> InferResult {
228        for (loc, bound, ty) in self.bindings.local_ptrs() {
229            infcx.subtyping(&ty, &bound, ConstrReason::FoldLocal)?;
230            self.bindings.remove_local(&loc);
231        }
232        Ok(())
233    }
234
235    /// Updates the type of `place` to `new_ty`. This may involve a *strong update* if we have
236    /// ownership of `place` or a *weak update* if it's behind a reference (which fires a subtyping
237    /// constraint)
238    ///
239    /// When strong updating, the process involves recovering the original regions (lifetimes) used
240    /// in the (unrefined) Rust type of `place` and then substituting these regions in `new_ty`. For
241    /// instance, if we are assigning a value of type `S<&'?10 i32{v: v > 0}>` to a variable `x`,
242    /// and the (unrefined) Rust type of `x` is `S<&'?5 i32>`, before the assignment, we identify a
243    /// substitution that maps the region `'?10` to `'?5`. After applying this substitution, the
244    /// type of the place `x` is updated accordingly. This ensures that the lifetimes in the
245    /// assigned type are consistent with those expected by the place's original type definition.
246    pub(crate) fn assign(
247        &mut self,
248        infcx: &mut InferCtxtAt,
249        place: &Place,
250        new_ty: Ty,
251    ) -> InferResult {
252        let rustc_ty = place.ty(infcx.genv, self.local_decls)?.ty;
253        let new_ty = ty_match_regions(&new_ty, &rustc_ty);
254        let span = infcx.span;
255        let result = self.bindings.lookup_unfolding(infcx, place, span)?;
256        if result.is_strg {
257            result.update(new_ty);
258        } else if !place.behind_raw_ptr(infcx.genv, self.local_decls)? {
259            infcx.subtyping(&new_ty, &result.ty, ConstrReason::Assign)?;
260        }
261        Ok(())
262    }
263
264    pub(crate) fn move_place(&mut self, infcx: &mut InferCtxtAt, place: &Place) -> InferResult<Ty> {
265        let span = infcx.span;
266        let result = self.bindings.lookup_unfolding(infcx, place, span)?;
267        if result.is_strg {
268            let uninit = Ty::uninit();
269            Ok(result.update(uninit))
270        } else {
271            // ignore the 'move' and trust rustc managed the move correctly
272            // https://github.com/flux-rs/flux/issues/725#issuecomment-2295065634
273            Ok(result.ty)
274        }
275    }
276
277    pub(crate) fn unpack(&mut self, infcx: &mut InferCtxt) {
278        self.bindings.fmap_mut(|ty| infcx.hoister(true).hoist(ty));
279    }
280
281    pub(crate) fn unblock(&mut self, infcx: &mut InferCtxt, place: &Place) {
282        self.bindings.unblock(infcx, place);
283    }
284
285    pub(crate) fn check_goto(
286        self,
287        infcx: &mut InferCtxtAt,
288        bb_env: &BasicBlockEnv,
289        target: BasicBlock,
290    ) -> InferResult {
291        infcx.ensure_resolved_evars(|infcx| {
292            let bb_env = bb_env
293                .data
294                .replace_bound_refts_with(|sort, mode, _| infcx.fresh_infer_var(sort, mode));
295
296            // Check constraints
297            for constr in &bb_env.constrs {
298                infcx.check_pred(constr, ConstrReason::Goto(target));
299            }
300
301            // Check subtyping
302            let bb_env = bb_env.bindings.flatten();
303            for (path, _, ty2) in bb_env {
304                let ty1 = self.bindings.get(&path);
305                infcx.subtyping(&ty1.unblocked(), &ty2.unblocked(), ConstrReason::Goto(target))?;
306            }
307            Ok(())
308        })
309    }
310
311    pub(crate) fn fold(&mut self, infcx: &mut InferCtxtAt, place: &Place) -> InferResult {
312        let span = infcx.span;
313        self.bindings.lookup(place, span).fold(infcx)?;
314        Ok(())
315    }
316
317    pub(crate) fn unfold_local_ptr(
318        &mut self,
319        infcx: &mut InferCtxt,
320        bound: &Ty,
321    ) -> InferResult<Loc> {
322        let loc = Loc::from(infcx.define_var(&Sort::Loc));
323        let ty = infcx.unpack(bound);
324        self.bindings
325            .insert(loc, LocKind::LocalPtr(bound.clone()), ty);
326        Ok(loc)
327    }
328
329    /// ```text
330    /// -----------------------------------
331    /// Γ ; &strg <ℓ: t> => Γ,ℓ: t ; ptr(ℓ)
332    /// ```
333    pub(crate) fn unfold_strg_ref(
334        &mut self,
335        infcx: &mut InferCtxt,
336        path: &Path,
337        ty: &Ty,
338    ) -> InferResult<Loc> {
339        if let Some(loc) = path.to_loc() {
340            let ty = infcx.unpack(ty);
341            self.bindings.insert(loc, LocKind::Universal, ty);
342            Ok(loc)
343        } else {
344            bug!("unfold_strg_ref: unexpected path {path:?}")
345        }
346    }
347
348    pub(crate) fn unfold(
349        &mut self,
350        infcx: &mut InferCtxt,
351        place: &Place,
352        span: Span,
353    ) -> InferResult {
354        self.bindings.unfold(infcx, place, span)
355    }
356
357    pub(crate) fn downcast(
358        &mut self,
359        infcx: &mut InferCtxtAt,
360        place: &Place,
361        variant_idx: VariantIdx,
362    ) -> InferResult {
363        let mut down_place = place.clone();
364        let span = infcx.span;
365        down_place
366            .projection
367            .push(PlaceElem::Downcast(None, variant_idx));
368        self.bindings.unfold(infcx, &down_place, span)?;
369        Ok(())
370    }
371
372    pub fn fully_resolve_evars(&mut self, infcx: &InferCtxt) {
373        self.bindings.fmap_mut(|ty| infcx.fully_resolve_evars(ty));
374    }
375
376    pub(crate) fn assume_ensures(
377        &mut self,
378        infcx: &mut InferCtxt,
379        ensures: &[Ensures],
380        span: Span,
381    ) {
382        for ensure in ensures {
383            match ensure {
384                Ensures::Type(path, updated_ty) => {
385                    let updated_ty = infcx.unpack(updated_ty);
386                    infcx.assume_invariants(&updated_ty);
387                    self.update_path(path, updated_ty, span);
388                }
389                Ensures::Pred(e) => infcx.assume_pred(e),
390            }
391        }
392    }
393
394    pub(crate) fn check_ensures(
395        &mut self,
396        at: &mut InferCtxtAt,
397        ensures: &[Ensures],
398        reason: ConstrReason,
399    ) -> InferResult {
400        for constraint in ensures {
401            match constraint {
402                Ensures::Type(path, ty) => {
403                    let actual_ty = self.get(path);
404                    at.subtyping(&actual_ty, ty, reason)?;
405                }
406                Ensures::Pred(e) => {
407                    at.check_pred(e, ConstrReason::Ret);
408                }
409            }
410        }
411        Ok(())
412    }
413}
414
415pub(crate) enum PtrToRefBound {
416    Ty(Ty),
417    Infer,
418    Identity,
419}
420
421impl flux_infer::infer::LocEnv for TypeEnv<'_> {
422    fn ptr_to_ref(
423        &mut self,
424        infcx: &mut InferCtxtAt,
425        reason: ConstrReason,
426        re: Region,
427        path: &Path,
428        bound: Ty,
429    ) -> InferResult<Ty> {
430        self.ptr_to_ref(infcx, reason, re, path, PtrToRefBound::Ty(bound))
431    }
432
433    fn get(&self, path: &Path) -> Ty {
434        self.get(path)
435    }
436
437    fn unfold_strg_ref(&mut self, infcx: &mut InferCtxt, path: &Path, ty: &Ty) -> InferResult<Loc> {
438        self.unfold_strg_ref(infcx, path, ty)
439    }
440}
441
442impl BasicBlockEnvShape {
443    pub fn enter<'a>(&self, local_decls: &'a LocalDecls) -> TypeEnv<'a> {
444        TypeEnv { bindings: self.bindings.clone(), local_decls }
445    }
446
447    fn new(scope: Scope, env: TypeEnv) -> BasicBlockEnvShape {
448        let mut bindings = env.bindings;
449        bindings.fmap_mut(|ty| BasicBlockEnvShape::pack_ty(&scope, ty));
450        BasicBlockEnvShape { scope, bindings }
451    }
452
453    fn pack_ty(scope: &Scope, ty: &Ty) -> Ty {
454        match ty.kind() {
455            TyKind::Indexed(bty, idxs) => {
456                let bty = BasicBlockEnvShape::pack_bty(scope, bty);
457                if scope.has_free_vars(idxs) {
458                    Ty::exists_with_constr(bty, Expr::hole(HoleKind::Pred))
459                } else {
460                    Ty::indexed(bty, idxs.clone())
461                }
462            }
463            TyKind::Downcast(adt, args, ty, variant, fields) => {
464                debug_assert!(!scope.has_free_vars(args));
465                debug_assert!(!scope.has_free_vars(ty));
466                let fields = fields.iter().map(|ty| Self::pack_ty(scope, ty)).collect();
467                Ty::downcast(adt.clone(), args.clone(), ty.clone(), *variant, fields)
468            }
469            TyKind::Blocked(ty) => Ty::blocked(BasicBlockEnvShape::pack_ty(scope, ty)),
470            // FIXME(nilehmann) [`TyKind::Exists`] could also contain free variables.
471            TyKind::Exists(_)
472            | TyKind::Discr(..)
473            | TyKind::Ptr(..)
474            | TyKind::Uninit
475            | TyKind::Param(_)
476            | TyKind::Constr(_, _) => ty.clone(),
477            TyKind::Infer(_) => bug!("unexpected hole whecn checking function body"),
478            TyKind::StrgRef(..) => bug!("unexpected strong reference when checking function body"),
479        }
480    }
481
482    fn pack_bty(scope: &Scope, bty: &BaseTy) -> BaseTy {
483        match bty {
484            BaseTy::Adt(adt_def, args) => {
485                let args = List::from_vec(
486                    args.iter()
487                        .map(|arg| Self::pack_generic_arg(scope, arg))
488                        .collect(),
489                );
490                BaseTy::adt(adt_def.clone(), args)
491            }
492            BaseTy::FnDef(def_id, args) => {
493                let args = List::from_vec(
494                    args.iter()
495                        .map(|arg| Self::pack_generic_arg(scope, arg))
496                        .collect(),
497                );
498                BaseTy::fn_def(*def_id, args)
499            }
500            BaseTy::Tuple(tys) => {
501                let tys = tys
502                    .iter()
503                    .map(|ty| BasicBlockEnvShape::pack_ty(scope, ty))
504                    .collect();
505                BaseTy::Tuple(tys)
506            }
507            BaseTy::Slice(ty) => BaseTy::Slice(Self::pack_ty(scope, ty)),
508            BaseTy::Ref(r, ty, mutbl) => BaseTy::Ref(*r, Self::pack_ty(scope, ty), *mutbl),
509            BaseTy::Array(ty, c) => BaseTy::Array(Self::pack_ty(scope, ty), c.clone()),
510            BaseTy::Int(_)
511            | BaseTy::Param(_)
512            | BaseTy::Uint(_)
513            | BaseTy::Bool
514            | BaseTy::Float(_)
515            | BaseTy::Str
516            | BaseTy::RawPtr(_, _)
517            | BaseTy::RawPtrMetadata(_)
518            | BaseTy::Char
519            | BaseTy::Never
520            | BaseTy::Closure(..)
521            | BaseTy::Dynamic(..)
522            | BaseTy::Alias(..)
523            | BaseTy::FnPtr(..)
524            | BaseTy::Foreign(..)
525            | BaseTy::Coroutine(..) => {
526                if scope.has_free_vars(bty) {
527                    tracked_span_bug!("unexpected type with free vars")
528                } else {
529                    bty.clone()
530                }
531            }
532            BaseTy::Infer(..) => {
533                tracked_span_bug!("unexpected infer type")
534            }
535        }
536    }
537
538    fn pack_generic_arg(scope: &Scope, arg: &GenericArg) -> GenericArg {
539        match arg {
540            GenericArg::Ty(ty) => GenericArg::Ty(Self::pack_ty(scope, ty)),
541            GenericArg::Base(arg) => {
542                assert!(!scope.has_free_vars(arg));
543                GenericArg::Base(arg.clone())
544            }
545            GenericArg::Lifetime(re) => GenericArg::Lifetime(*re),
546            GenericArg::Const(c) => GenericArg::Const(c.clone()),
547        }
548    }
549
550    fn update(&mut self, path: &Path, ty: Ty, span: Span) {
551        self.bindings.lookup(path, span).update(ty);
552    }
553
554    /// join(self, genv, other) consumes the bindings in other, to "update"
555    /// `self` in place, and returns `true` if there was an actual change
556    /// or `false` indicating no change (i.e., a fixpoint was reached).
557    pub(crate) fn join(&mut self, other: TypeEnv, span: Span) -> bool {
558        let paths = self.bindings.paths();
559
560        // Join types
561        let mut modified = false;
562        for path in &paths {
563            let ty1 = self.bindings.get(path);
564            let ty2 = other.bindings.get(path);
565            let ty = self.join_ty(&ty1, &ty2);
566            modified |= ty1 != ty;
567            self.update(path, ty, span);
568        }
569
570        modified
571    }
572
573    fn join_ty(&self, ty1: &Ty, ty2: &Ty) -> Ty {
574        match (ty1.kind(), ty2.kind()) {
575            (TyKind::Blocked(ty1), _) => Ty::blocked(self.join_ty(ty1, &ty2.unblocked())),
576            (_, TyKind::Blocked(ty2)) => Ty::blocked(self.join_ty(&ty1.unblocked(), ty2)),
577            (TyKind::Uninit, _) | (_, TyKind::Uninit) => Ty::uninit(),
578            (TyKind::Exists(ty1), _) => self.join_ty(ty1.as_ref().skip_binder(), ty2),
579            (_, TyKind::Exists(ty2)) => self.join_ty(ty1, ty2.as_ref().skip_binder()),
580            (TyKind::Constr(_, ty1), _) => self.join_ty(ty1, ty2),
581            (_, TyKind::Constr(_, ty2)) => self.join_ty(ty1, ty2),
582            (TyKind::Indexed(bty1, idx1), TyKind::Indexed(bty2, idx2)) => {
583                let bty = self.join_bty(bty1, bty2);
584                let mut sorts = vec![];
585                let idx = self.join_idx(idx1, idx2, &bty.sort(), &mut sorts);
586                if sorts.is_empty() {
587                    Ty::indexed(bty, idx)
588                } else {
589                    let ty = Ty::constr(Expr::hole(HoleKind::Pred), Ty::indexed(bty, idx));
590                    Ty::exists(Binder::bind_with_sorts(ty, &sorts))
591                }
592            }
593            (TyKind::Ptr(rk1, path1), TyKind::Ptr(rk2, path2)) => {
594                debug_assert_eq!(rk1, rk2);
595                debug_assert_eq!(path1, path2);
596                Ty::ptr(*rk1, path1.clone())
597            }
598            (TyKind::Param(param_ty1), TyKind::Param(param_ty2)) => {
599                debug_assert_eq!(param_ty1, param_ty2);
600                Ty::param(*param_ty1)
601            }
602            (
603                TyKind::Downcast(adt1, args1, ty1, variant1, fields1),
604                TyKind::Downcast(adt2, args2, ty2, variant2, fields2),
605            ) => {
606                debug_assert_eq!(adt1, adt2);
607                debug_assert_eq!(args1, args2);
608                debug_assert!(ty1 == ty2 && !self.scope.has_free_vars(ty2));
609                debug_assert_eq!(variant1, variant2);
610                debug_assert_eq!(fields1.len(), fields2.len());
611                let fields = iter::zip(fields1, fields2)
612                    .map(|(ty1, ty2)| self.join_ty(ty1, ty2))
613                    .collect();
614                Ty::downcast(adt1.clone(), args1.clone(), ty1.clone(), *variant1, fields)
615            }
616            _ => tracked_span_bug!("unexpected types: `{ty1:?}` - `{ty2:?}`"),
617        }
618    }
619
620    fn join_idx(&self, e1: &Expr, e2: &Expr, sort: &Sort, bound_sorts: &mut Vec<Sort>) -> Expr {
621        match (e1.kind(), e2.kind(), sort) {
622            (ExprKind::Tuple(es1), ExprKind::Tuple(es2), Sort::Tuple(sorts)) => {
623                debug_assert_eq3!(es1.len(), es2.len(), sorts.len());
624                Expr::tuple(
625                    izip!(es1, es2, sorts)
626                        .map(|(e1, e2, sort)| self.join_idx(e1, e2, sort, bound_sorts))
627                        .collect(),
628                )
629            }
630            (
631                ExprKind::Ctor(Ctor::Struct(_), flds1),
632                ExprKind::Ctor(Ctor::Struct(_), flds2),
633                Sort::App(SortCtor::Adt(sort_def), args),
634            ) => {
635                let sorts = sort_def.struct_variant().field_sorts(args);
636                debug_assert_eq3!(flds1.len(), flds2.len(), sorts.len());
637
638                Expr::ctor_struct(
639                    sort_def.did(),
640                    izip!(flds1, flds2, &sorts)
641                        .map(|(f1, f2, sort)| self.join_idx(f1, f2, sort, bound_sorts))
642                        .collect(),
643                )
644            }
645            _ => {
646                let has_free_vars2 = self.scope.has_free_vars(e2);
647                let has_escaping_vars1 = e1.has_escaping_bvars();
648                let has_escaping_vars2 = e2.has_escaping_bvars();
649                if !has_free_vars2 && !has_escaping_vars1 && !has_escaping_vars2 && e1 == e2 {
650                    e1.clone()
651                } else if sort.is_pred() {
652                    // FIXME(nilehmann) we shouldn't special case predicates here. Instead, we
653                    // should differentiate between generics and indices.
654                    let fsort = sort.expect_func().expect_mono();
655                    Expr::abs(Lambda::bind_with_fsort(Expr::hole(HoleKind::Pred), fsort))
656                } else {
657                    bound_sorts.push(sort.clone());
658                    Expr::bvar(
659                        INNERMOST,
660                        BoundVar::from_usize(bound_sorts.len() - 1),
661                        BoundReftKind::Annon,
662                    )
663                }
664            }
665        }
666    }
667
668    fn join_bty(&self, bty1: &BaseTy, bty2: &BaseTy) -> BaseTy {
669        match (bty1, bty2) {
670            (BaseTy::Adt(def1, args1), BaseTy::Adt(def2, args2)) => {
671                tracked_span_dbg_assert_eq!(def1.did(), def2.did());
672                let args = iter::zip(args1, args2)
673                    .map(|(arg1, arg2)| self.join_generic_arg(arg1, arg2))
674                    .collect();
675                BaseTy::adt(def1.clone(), List::from_vec(args))
676            }
677            (BaseTy::Tuple(fields1), BaseTy::Tuple(fields2)) => {
678                let fields = iter::zip(fields1, fields2)
679                    .map(|(ty1, ty2)| self.join_ty(ty1, ty2))
680                    .collect();
681                BaseTy::Tuple(fields)
682            }
683            (BaseTy::Alias(kind1, alias_ty1), BaseTy::Alias(kind2, alias_ty2)) => {
684                tracked_span_dbg_assert_eq!(kind1, kind2);
685                tracked_span_dbg_assert_eq!(alias_ty1, alias_ty2);
686                BaseTy::Alias(*kind1, alias_ty1.clone())
687            }
688            (BaseTy::Ref(r1, ty1, mutbl1), BaseTy::Ref(r2, ty2, mutbl2)) => {
689                tracked_span_dbg_assert_eq!(r1, r2);
690                tracked_span_dbg_assert_eq!(mutbl1, mutbl2);
691                BaseTy::Ref(*r1, self.join_ty(ty1, ty2), *mutbl1)
692            }
693            (BaseTy::Array(ty1, len1), BaseTy::Array(ty2, len2)) => {
694                tracked_span_dbg_assert_eq!(len1, len2);
695                BaseTy::Array(self.join_ty(ty1, ty2), len1.clone())
696            }
697            (BaseTy::Slice(ty1), BaseTy::Slice(ty2)) => BaseTy::Slice(self.join_ty(ty1, ty2)),
698            _ => {
699                tracked_span_dbg_assert_eq!(bty1, bty2);
700                bty1.clone()
701            }
702        }
703    }
704
705    fn join_generic_arg(&self, arg1: &GenericArg, arg2: &GenericArg) -> GenericArg {
706        match (arg1, arg2) {
707            (GenericArg::Ty(ty1), GenericArg::Ty(ty2)) => GenericArg::Ty(self.join_ty(ty1, ty2)),
708            (GenericArg::Base(ctor1), GenericArg::Base(ctor2)) => {
709                let sty1 = ctor1.as_ref().skip_binder();
710                let sty2 = ctor2.as_ref().skip_binder();
711                debug_assert!(sty1.idx.is_nu());
712                debug_assert!(sty2.idx.is_nu());
713
714                let bty = self.join_bty(&sty1.bty, &sty2.bty);
715                let pred = if self.scope.has_free_vars(&sty2.pred) || sty1.pred != sty2.pred {
716                    Expr::hole(HoleKind::Pred)
717                } else {
718                    sty1.pred.clone()
719                };
720                let sort = bty.sort();
721                let ctor = Binder::bind_with_sort(SubsetTy::new(bty, Expr::nu(), pred), sort);
722                GenericArg::Base(ctor)
723            }
724            (GenericArg::Lifetime(re1), GenericArg::Lifetime(_re2)) => {
725                // TODO(nilehmann) loop_abstract_refinement.rs is triggering this assertion to fail
726                // wee should fix it.
727                // debug_assert_eq!(re1, _re2);
728                GenericArg::Lifetime(*re1)
729            }
730            (GenericArg::Const(c1), GenericArg::Const(c2)) => {
731                debug_assert_eq!(c1, c2);
732                GenericArg::Const(c1.clone())
733            }
734            _ => tracked_span_bug!("unexpected generic args: `{arg1:?}` - `{arg2:?}`"),
735        }
736    }
737
738    pub fn into_bb_env(self, infcx: &mut InferCtxtRoot) -> BasicBlockEnv {
739        let mut delegate = LocalHoister::default();
740        let mut hoister = Hoister::with_delegate(&mut delegate).transparent();
741
742        let mut bindings = self.bindings;
743        bindings.fmap_mut(|ty| hoister.hoist(ty));
744
745        BasicBlockEnv {
746            // We are relying on all the types in `bindings` not having escaping bvars, otherwise
747            // we would have to shift them in since we are creating a new binder.
748            data: delegate.bind(|vars, preds| {
749                // Replace all holes with a single fresh kvar on all parameters
750                let mut constrs = preds
751                    .into_iter()
752                    .filter(|pred| !matches!(pred.kind(), ExprKind::Hole(HoleKind::Pred)))
753                    .collect_vec();
754                let kvar = infcx.fresh_kvar_in_scope(
755                    std::slice::from_ref(&vars),
756                    &self.scope,
757                    KVarEncoding::Conj,
758                );
759                constrs.push(kvar);
760
761                // Replace remaining holes by fresh kvars
762                let mut kvar_gen = |binders: &[_], kind| {
763                    debug_assert_eq!(kind, HoleKind::Pred);
764                    let binders = std::iter::once(vars.clone())
765                        .chain(binders.iter().cloned())
766                        .collect_vec();
767                    infcx.fresh_kvar_in_scope(&binders, &self.scope, KVarEncoding::Conj)
768                };
769                bindings.fmap_mut(|binding| binding.replace_holes(&mut kvar_gen));
770
771                BasicBlockEnvData { constrs: constrs.into(), bindings }
772            }),
773            scope: self.scope,
774        }
775    }
776}
777
778impl TypeVisitable for BasicBlockEnvData {
779    fn visit_with<V: TypeVisitor>(&self, _visitor: &mut V) -> ControlFlow<V::BreakTy> {
780        unimplemented!()
781    }
782}
783
784impl TypeFoldable for BasicBlockEnvData {
785    fn try_fold_with<F: FallibleTypeFolder>(
786        &self,
787        folder: &mut F,
788    ) -> std::result::Result<Self, F::Error> {
789        Ok(BasicBlockEnvData {
790            constrs: self.constrs.try_fold_with(folder)?,
791            bindings: self.bindings.try_fold_with(folder)?,
792        })
793    }
794}
795
796impl BasicBlockEnv {
797    pub(crate) fn enter<'a>(
798        &self,
799        infcx: &mut InferCtxt,
800        local_decls: &'a LocalDecls,
801    ) -> TypeEnv<'a> {
802        let data = self
803            .data
804            .replace_bound_refts_with(|sort, _, _| Expr::fvar(infcx.define_var(sort)));
805        for constr in &data.constrs {
806            infcx.assume_pred(constr);
807        }
808        TypeEnv { bindings: data.bindings, local_decls }
809    }
810
811    pub(crate) fn scope(&self) -> &Scope {
812        &self.scope
813    }
814}
815
816mod pretty {
817    use std::fmt;
818
819    use flux_middle::pretty::*;
820
821    use super::*;
822
823    impl Pretty for TypeEnv<'_> {
824        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
825            w!(cx, f, "{:?}", &self.bindings)
826        }
827
828        fn default_cx(tcx: TyCtxt) -> PrettyCx {
829            PlacesTree::default_cx(tcx)
830        }
831    }
832
833    impl Pretty for BasicBlockEnvShape {
834        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
835            w!(cx, f, "{:?} {:?}", &self.scope, &self.bindings)
836        }
837
838        fn default_cx(tcx: TyCtxt) -> PrettyCx {
839            PlacesTree::default_cx(tcx)
840        }
841    }
842
843    impl Pretty for BasicBlockEnv {
844        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
845            w!(cx, f, "{:?} ", &self.scope)?;
846
847            let vars = self.data.vars();
848            cx.with_bound_vars(vars, || {
849                if !vars.is_empty() {
850                    cx.fmt_bound_vars(true, "for<", vars, "> ", f)?;
851                }
852                let data = self.data.as_ref().skip_binder();
853                if !data.constrs.is_empty() {
854                    w!(
855                        cx,
856                        f,
857                        "{:?} ⇒ ",
858                        join!(", ", data.constrs.iter().filter(|pred| !pred.is_trivially_true()))
859                    )?;
860                }
861                w!(cx, f, "{:?}", &data.bindings)
862            })
863        }
864
865        fn default_cx(tcx: TyCtxt) -> PrettyCx {
866            PlacesTree::default_cx(tcx)
867        }
868    }
869
870    impl_debug_with_default_cx! {
871        TypeEnv<'_> => "type_env",
872        BasicBlockEnvShape => "basic_block_env_shape",
873        BasicBlockEnv => "basic_block_env"
874    }
875}
876
877/// A very explicit representation of [`TypeEnv`] for debugging/tracing/serialization ONLY.
878#[derive(Serialize, DebugAsJson)]
879pub struct TypeEnvTrace(Vec<TypeEnvBind>);
880
881#[derive(Serialize)]
882struct TypeEnvBind {
883    local: LocInfo,
884    name: Option<String>,
885    kind: String,
886    ty: String,
887    span: Option<SpanTrace>,
888}
889
890#[derive(Serialize)]
891enum LocInfo {
892    Local(String),
893    Var(String),
894}
895
896fn loc_info(loc: &Loc) -> LocInfo {
897    match loc {
898        Loc::Local(local) => LocInfo::Local(format!("{local:?}")),
899        Loc::Var(var) => LocInfo::Var(format!("{var:?}")),
900    }
901}
902
903fn loc_name(local_names: &UnordMap<Local, Symbol>, loc: &Loc) -> Option<String> {
904    if let Loc::Local(local) = loc {
905        let name = local_names.get(local)?;
906        return Some(format!("{name}"));
907    }
908    None
909}
910
911fn loc_span(
912    genv: GlobalEnv,
913    local_decls: &IndexVec<Local, LocalDecl>,
914    loc: &Loc,
915) -> Option<SpanTrace> {
916    if let Loc::Local(local) = loc {
917        return local_decls
918            .get(*local)
919            .map(|local_decl| SpanTrace::new(genv, local_decl.source_info.span));
920    }
921    None
922}
923
924impl TypeEnvTrace {
925    pub fn new(
926        genv: GlobalEnv,
927        local_names: &UnordMap<Local, Symbol>,
928        local_decls: &IndexVec<Local, LocalDecl>,
929        env: &TypeEnv,
930    ) -> Self {
931        let mut bindings = vec![];
932        let cx = PrettyCx::default(genv).hide_regions(true);
933        env.bindings
934            .iter()
935            .filter(|(_, binding)| !binding.ty.is_uninit())
936            .sorted_by(|(loc1, _), (loc2, _)| loc1.cmp(loc2))
937            .for_each(|(loc, binding)| {
938                let name = loc_name(local_names, loc);
939                let local = loc_info(loc);
940                let kind = format!("{:?}", binding.kind);
941                let ty = binding.ty.nested_string(&cx);
942                let span = loc_span(genv, local_decls, loc);
943                bindings.push(TypeEnvBind { name, local, kind, ty, span });
944            });
945
946        TypeEnvTrace(bindings)
947    }
948}
949
950#[derive(Serialize, DebugAsJson)]
951pub struct SpanTrace {
952    file: Option<String>,
953    start_line: usize,
954    start_col: usize,
955    end_line: usize,
956    end_col: usize,
957}
958
959impl SpanTrace {
960    fn span_file(tcx: TyCtxt, span: Span) -> Option<String> {
961        let sm = tcx.sess.source_map();
962        let current_dir = &tcx.sess.opts.working_dir;
963        let current_dir = current_dir.local_path()?;
964        if let rustc_span::FileName::Real(file_name) = sm.span_to_filename(span) {
965            let file_path = file_name.local_path()?;
966            let full_path = current_dir.join(file_path);
967            Some(full_path.display().to_string())
968        } else {
969            None
970        }
971    }
972    pub fn new(genv: GlobalEnv, span: Span) -> Self {
973        let tcx = genv.tcx();
974        let sm = tcx.sess.source_map();
975        let (_, start_line, start_col, end_line, end_col) = sm.span_to_location_info(span);
976        let file = SpanTrace::span_file(tcx, span);
977        SpanTrace { file, start_line, start_col, end_line, end_col }
978    }
979}
flux_refineck/type_env.rs

flux_refineck/
type_env.rs
