flux_middle/rty/
expr.rs

1use std::{fmt, iter, ops::ControlFlow, sync::OnceLock};
2
3use flux_arc_interner::{Interned, List, impl_internable, impl_slice_internable};
4use flux_common::bug;
5use flux_macros::{TypeFoldable, TypeVisitable};
6use flux_rustc_bridge::{
7    ToRustc,
8    const_eval::{scalar_to_bits, scalar_to_int, scalar_to_uint},
9    ty::{Const, ConstKind, ValTree, VariantIdx},
10};
11use itertools::Itertools;
12use rustc_abi::{FIRST_VARIANT, FieldIdx};
13use rustc_data_structures::snapshot_map::SnapshotMap;
14use rustc_hir::def_id::DefId;
15use rustc_index::newtype_index;
16use rustc_macros::{Decodable, Encodable, TyDecodable, TyEncodable};
17use rustc_middle::{
18    mir::Local,
19    ty::{ParamConst, ScalarInt, TyCtxt},
20};
21use rustc_span::{Span, Symbol};
22use rustc_type_ir::{BoundVar, DebruijnIndex, INNERMOST};
23
24use super::{
25    BaseTy, Binder, BoundReftKind, BoundVariableKinds, FuncSort, GenericArgs, GenericArgsExt as _,
26    IntTy, Sort, UintTy,
27};
28use crate::{
29    big_int::BigInt,
30    def_id::FluxDefId,
31    fhir::{self, SpecFuncKind},
32    global_env::GlobalEnv,
33    pretty::*,
34    queries::QueryResult,
35    rty::{
36        BoundVariableKind,
37        fold::{
38            TypeFoldable, TypeFolder, TypeSuperFoldable, TypeSuperVisitable, TypeVisitable as _,
39            TypeVisitor,
40        },
41    },
42};
43
44/// A lambda abstraction with an elaborated output sort. We need the output sort of lambdas for
45/// encoding into fixpoint
46#[derive(Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable, TypeVisitable, TypeFoldable)]
47pub struct Lambda {
48    body: Binder<Expr>,
49    output: Sort,
50}
51
52impl Lambda {
53    pub fn bind_with_vars(body: Expr, inputs: BoundVariableKinds, output: Sort) -> Self {
54        debug_assert!(inputs.iter().all(BoundVariableKind::is_refine));
55        Self { body: Binder::bind_with_vars(body, inputs), output }
56    }
57
58    pub fn bind_with_fsort(body: Expr, fsort: FuncSort) -> Self {
59        Self { body: Binder::bind_with_sorts(body, fsort.inputs()), output: fsort.output().clone() }
60    }
61
62    pub fn apply(&self, args: &[Expr]) -> Expr {
63        self.body.replace_bound_refts(args)
64    }
65
66    pub fn vars(&self) -> &BoundVariableKinds {
67        self.body.vars()
68    }
69
70    pub fn output(&self) -> Sort {
71        self.output.clone()
72    }
73
74    pub fn fsort(&self) -> FuncSort {
75        let inputs_and_output = self
76            .vars()
77            .iter()
78            .map(|kind| kind.expect_sort().clone())
79            .chain(iter::once(self.output()))
80            .collect();
81        FuncSort { inputs_and_output }
82    }
83}
84
85#[derive(Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable, TypeVisitable, TypeFoldable)]
86pub struct AliasReft {
87    /// Id of the associated refinement in the trait
88    pub assoc_id: FluxDefId,
89    pub args: GenericArgs,
90}
91
92impl AliasReft {
93    pub fn to_rustc_trait_ref<'tcx>(&self, tcx: TyCtxt<'tcx>) -> rustc_middle::ty::TraitRef<'tcx> {
94        let trait_def_id = self.assoc_id.parent();
95        let args = self
96            .args
97            .to_rustc(tcx)
98            .truncate_to(tcx, tcx.generics_of(trait_def_id));
99        rustc_middle::ty::TraitRef::new(tcx, trait_def_id, args)
100    }
101}
102
103#[derive(Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable)]
104pub struct Expr {
105    kind: Interned<ExprKind>,
106    espan: Option<ESpan>,
107}
108
109impl Expr {
110    pub fn at_opt(self, espan: Option<ESpan>) -> Expr {
111        Expr { kind: self.kind, espan }
112    }
113
114    pub fn at(self, espan: ESpan) -> Expr {
115        self.at_opt(Some(espan))
116    }
117
118    pub fn at_base(self, base: ESpan) -> Expr {
119        if let Some(espan) = self.espan { self.at(espan.with_base(base)) } else { self }
120    }
121
122    pub fn span(&self) -> Option<ESpan> {
123        self.espan
124    }
125
126    pub fn tt() -> Expr {
127        static TRUE: OnceLock<Expr> = OnceLock::new();
128        TRUE.get_or_init(|| ExprKind::Constant(Constant::Bool(true)).intern())
129            .clone()
130    }
131
132    pub fn ff() -> Expr {
133        static FALSE: OnceLock<Expr> = OnceLock::new();
134        FALSE
135            .get_or_init(|| ExprKind::Constant(Constant::Bool(false)).intern())
136            .clone()
137    }
138
139    pub fn and_from_iter(exprs: impl IntoIterator<Item = Expr>) -> Expr {
140        exprs
141            .into_iter()
142            .reduce(|acc, e| Expr::binary_op(BinOp::And, acc, e))
143            .unwrap_or_else(Expr::tt)
144    }
145
146    pub fn or_from_iter(exprs: impl IntoIterator<Item = Expr>) -> Expr {
147        exprs
148            .into_iter()
149            .reduce(|acc, e| Expr::binary_op(BinOp::Or, acc, e))
150            .unwrap_or_else(Expr::ff)
151    }
152
153    pub fn and(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
154        Expr::and_from_iter([e1.into(), e2.into()])
155    }
156
157    pub fn or(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
158        Expr::or_from_iter([e1.into(), e2.into()])
159    }
160
161    pub fn zero() -> Expr {
162        static ZERO: OnceLock<Expr> = OnceLock::new();
163        ZERO.get_or_init(|| ExprKind::Constant(Constant::ZERO).intern())
164            .clone()
165    }
166
167    pub fn int_max(int_ty: IntTy) -> Expr {
168        let bit_width: u64 = int_ty
169            .bit_width()
170            .unwrap_or(flux_config::pointer_width().bits());
171        Expr::constant(Constant::int_max(bit_width.try_into().unwrap()))
172    }
173
174    pub fn int_min(int_ty: IntTy) -> Expr {
175        let bit_width: u64 = int_ty
176            .bit_width()
177            .unwrap_or(flux_config::pointer_width().bits());
178        Expr::constant(Constant::int_min(bit_width.try_into().unwrap()))
179    }
180
181    pub fn uint_max(uint_ty: UintTy) -> Expr {
182        let bit_width: u64 = uint_ty
183            .bit_width()
184            .unwrap_or(flux_config::pointer_width().bits());
185        Expr::constant(Constant::uint_max(bit_width.try_into().unwrap()))
186    }
187
188    pub fn nu() -> Expr {
189        Expr::bvar(INNERMOST, BoundVar::ZERO, BoundReftKind::Annon)
190    }
191
192    pub fn is_nu(&self) -> bool {
193        if let ExprKind::Var(Var::Bound(INNERMOST, var)) = self.kind()
194            && var.var == BoundVar::ZERO
195        {
196            true
197        } else {
198            false
199        }
200    }
201
202    pub fn unit() -> Expr {
203        Expr::tuple(List::empty())
204    }
205
206    pub fn var(var: Var) -> Expr {
207        ExprKind::Var(var).intern()
208    }
209
210    pub fn fvar(name: Name) -> Expr {
211        Var::Free(name).to_expr()
212    }
213
214    pub fn evar(evid: EVid) -> Expr {
215        Var::EVar(evid).to_expr()
216    }
217
218    pub fn bvar(debruijn: DebruijnIndex, var: BoundVar, kind: BoundReftKind) -> Expr {
219        Var::Bound(debruijn, BoundReft { var, kind }).to_expr()
220    }
221
222    pub fn early_param(index: u32, name: Symbol) -> Expr {
223        Var::EarlyParam(EarlyReftParam { index, name }).to_expr()
224    }
225
226    pub fn local(local: Local) -> Expr {
227        ExprKind::Local(local).intern()
228    }
229
230    pub fn constant(c: Constant) -> Expr {
231        ExprKind::Constant(c).intern()
232    }
233
234    pub fn const_def_id(c: DefId) -> Expr {
235        ExprKind::ConstDefId(c).intern()
236    }
237
238    pub fn const_generic(param: ParamConst) -> Expr {
239        ExprKind::Var(Var::ConstGeneric(param)).intern()
240    }
241
242    pub fn tuple(flds: List<Expr>) -> Expr {
243        ExprKind::Tuple(flds).intern()
244    }
245
246    pub fn ctor_struct(def_id: DefId, flds: List<Expr>) -> Expr {
247        ExprKind::Ctor(Ctor::Struct(def_id), flds).intern()
248    }
249
250    pub fn ctor_enum(def_id: DefId, idx: VariantIdx) -> Expr {
251        ExprKind::Ctor(Ctor::Enum(def_id, idx), List::empty()).intern()
252    }
253
254    pub fn ctor(ctor: Ctor, flds: List<Expr>) -> Expr {
255        ExprKind::Ctor(ctor, flds).intern()
256    }
257
258    pub fn from_bits(bty: &BaseTy, bits: u128) -> Expr {
259        // FIXME: We are assuming the higher bits are not set. check this assumption
260        match bty {
261            BaseTy::Int(_) => {
262                let bits = bits as i128;
263                ExprKind::Constant(Constant::from(bits)).intern()
264            }
265            BaseTy::Uint(_) => ExprKind::Constant(Constant::from(bits)).intern(),
266            BaseTy::Bool => ExprKind::Constant(Constant::Bool(bits != 0)).intern(),
267            BaseTy::Char => {
268                let c = char::from_u32(bits.try_into().unwrap()).unwrap();
269                ExprKind::Constant(Constant::Char(c)).intern()
270            }
271            _ => bug!(),
272        }
273    }
274
275    pub fn ite(p: impl Into<Expr>, e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
276        ExprKind::IfThenElse(p.into(), e1.into(), e2.into()).intern()
277    }
278
279    pub fn abs(lam: Lambda) -> Expr {
280        ExprKind::Abs(lam).intern()
281    }
282
283    pub fn let_(init: Expr, body: Binder<Expr>) -> Expr {
284        ExprKind::Let(init, body).intern()
285    }
286
287    pub fn bounded_quant(kind: fhir::QuantKind, rng: fhir::Range, body: Binder<Expr>) -> Expr {
288        ExprKind::BoundedQuant(kind, rng, body).intern()
289    }
290
291    pub fn hole(kind: HoleKind) -> Expr {
292        ExprKind::Hole(kind).intern()
293    }
294
295    pub fn kvar(kvar: KVar) -> Expr {
296        ExprKind::KVar(kvar).intern()
297    }
298
299    pub fn alias(alias: AliasReft, args: List<Expr>) -> Expr {
300        ExprKind::Alias(alias, args).intern()
301    }
302
303    pub fn forall(expr: Binder<Expr>) -> Expr {
304        ExprKind::ForAll(expr).intern()
305    }
306
307    pub fn binary_op(op: BinOp, e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
308        ExprKind::BinaryOp(op, e1.into(), e2.into()).intern()
309    }
310
311    pub fn unit_struct(def_id: DefId) -> Expr {
312        Expr::ctor_struct(def_id, List::empty())
313    }
314
315    pub fn app(func: impl Into<Expr>, args: List<Expr>) -> Expr {
316        ExprKind::App(func.into(), args).intern()
317    }
318
319    pub fn global_func(kind: SpecFuncKind) -> Expr {
320        ExprKind::GlobalFunc(kind).intern()
321    }
322
323    pub fn eq(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
324        ExprKind::BinaryOp(BinOp::Eq, e1.into(), e2.into()).intern()
325    }
326
327    pub fn unary_op(op: UnOp, e: impl Into<Expr>) -> Expr {
328        ExprKind::UnaryOp(op, e.into()).intern()
329    }
330
331    pub fn ne(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
332        ExprKind::BinaryOp(BinOp::Ne, e1.into(), e2.into()).intern()
333    }
334
335    pub fn ge(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
336        ExprKind::BinaryOp(BinOp::Ge(Sort::Int), e1.into(), e2.into()).intern()
337    }
338
339    pub fn gt(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
340        ExprKind::BinaryOp(BinOp::Gt(Sort::Int), e1.into(), e2.into()).intern()
341    }
342
343    pub fn lt(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
344        ExprKind::BinaryOp(BinOp::Lt(Sort::Int), e1.into(), e2.into()).intern()
345    }
346
347    pub fn le(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
348        ExprKind::BinaryOp(BinOp::Le(Sort::Int), e1.into(), e2.into()).intern()
349    }
350
351    pub fn implies(e1: impl Into<Expr>, e2: impl Into<Expr>) -> Expr {
352        ExprKind::BinaryOp(BinOp::Imp, e1.into(), e2.into()).intern()
353    }
354
355    pub fn field_proj(e: impl Into<Expr>, proj: FieldProj) -> Expr {
356        ExprKind::FieldProj(e.into(), proj).intern()
357    }
358
359    pub fn field_projs(e: impl Into<Expr>, projs: &[FieldProj]) -> Expr {
360        projs.iter().copied().fold(e.into(), Expr::field_proj)
361    }
362
363    pub fn path_proj(base: Expr, field: FieldIdx) -> Expr {
364        ExprKind::PathProj(base, field).intern()
365    }
366
367    pub fn not(&self) -> Expr {
368        ExprKind::UnaryOp(UnOp::Not, self.clone()).intern()
369    }
370
371    pub fn neg(&self) -> Expr {
372        ExprKind::UnaryOp(UnOp::Neg, self.clone()).intern()
373    }
374
375    pub fn kind(&self) -> &ExprKind {
376        &self.kind
377    }
378
379    /// An expression is an *atom* if it is "self-delimiting", i.e., it has a clear boundary
380    /// when printed. This is used to avoid unnecessary parenthesis when pretty printing.
381    pub fn is_atom(&self) -> bool {
382        !matches!(self.kind(), ExprKind::Abs(..) | ExprKind::BinaryOp(..) | ExprKind::ForAll(..))
383    }
384
385    /// Simple syntactic check to see if the expression is a trivially true predicate. This is used
386    /// mostly for filtering predicates when pretty printing but also to simplify types in general.
387    pub fn is_trivially_true(&self) -> bool {
388        self.is_true()
389            || matches!(self.kind(), ExprKind::BinaryOp(BinOp::Eq | BinOp::Iff | BinOp::Imp, e1, e2) if e1.erase_spans() == e2.erase_spans())
390    }
391
392    /// Simple syntactic check to see if the expression is a trivially false predicate.
393    pub fn is_trivially_false(&self) -> bool {
394        self.is_false()
395    }
396
397    /// Whether the expression is *literally* the constant `true`.
398    fn is_true(&self) -> bool {
399        matches!(self.kind(), ExprKind::Constant(Constant::Bool(true)))
400    }
401
402    /// Whether the expression is *literally* the constant `false`.
403    fn is_false(&self) -> bool {
404        matches!(self.kind(), ExprKind::Constant(Constant::Bool(false)))
405    }
406
407    pub fn from_const(tcx: TyCtxt, c: &Const) -> Expr {
408        match &c.kind {
409            ConstKind::Param(param_const) => Expr::const_generic(*param_const),
410            ConstKind::Value(ty, ValTree::Leaf(scalar)) => {
411                Expr::constant(Constant::from_scalar_int(tcx, *scalar, ty).unwrap())
412            }
413            ConstKind::Value(_ty, ValTree::Branch(_)) => {
414                bug!("todo: ValTree::Branch {c:?}")
415            }
416            // We should have normalized away the unevaluated constants
417            ConstKind::Unevaluated(_) => bug!("unexpected `ConstKind::Unevaluated`"),
418
419            ConstKind::Infer(_) => bug!("unexpected `ConstKind::Infer`"),
420        }
421    }
422
423    pub fn is_binary_op(&self) -> bool {
424        matches!(self.kind(), ExprKind::BinaryOp(..))
425    }
426
427    fn const_op(op: &BinOp, c1: &Constant, c2: &Constant) -> Option<Constant> {
428        match op {
429            BinOp::Iff => c1.iff(c2),
430            BinOp::Imp => c1.imp(c2),
431            BinOp::Or => c1.or(c2),
432            BinOp::And => c1.and(c2),
433            BinOp::Gt(Sort::Int) => c1.gt(c2),
434            BinOp::Ge(Sort::Int) => c1.ge(c2),
435            BinOp::Lt(Sort::Int) => c2.gt(c1),
436            BinOp::Le(Sort::Int) => c2.ge(c1),
437            BinOp::Eq => Some(c1.eq(c2)),
438            BinOp::Ne => Some(c1.ne(c2)),
439            _ => None,
440        }
441    }
442
443    /// Simplify the expression by removing double negations, short-circuiting boolean connectives and
444    /// doing constant folding. Note that we also have [`TypeFoldable::normalize`] which applies beta
445    /// reductions for tuples and abstractions.
446    ///
447    /// Additionally replaces any occurrences of elements in assumed_preds with True.
448    pub fn simplify(&self, assumed_preds: &SnapshotMap<Expr, ()>) -> Expr {
449        struct Simplify<'a> {
450            assumed_preds: &'a SnapshotMap<Expr, ()>,
451        }
452
453        impl TypeFolder for Simplify<'_> {
454            fn fold_expr(&mut self, expr: &Expr) -> Expr {
455                if self.assumed_preds.get(&expr.erase_spans()).is_some() {
456                    return Expr::tt();
457                }
458                let span = expr.span();
459                match expr.kind() {
460                    ExprKind::BinaryOp(op, e1, e2) => {
461                        let e1 = e1.fold_with(self);
462                        let e2 = e2.fold_with(self);
463                        match (op, e1.kind(), e2.kind()) {
464                            (BinOp::And, ExprKind::Constant(Constant::Bool(false)), _) => {
465                                Expr::constant(Constant::Bool(false)).at_opt(e1.span())
466                            }
467                            (BinOp::And, _, ExprKind::Constant(Constant::Bool(false))) => {
468                                Expr::constant(Constant::Bool(false)).at_opt(e2.span())
469                            }
470                            (BinOp::And, ExprKind::Constant(Constant::Bool(true)), _) => e2,
471                            (BinOp::And, _, ExprKind::Constant(Constant::Bool(true))) => e1,
472                            (op, ExprKind::Constant(c1), ExprKind::Constant(c2)) => {
473                                if let Some(c) = Expr::const_op(op, c1, c2) {
474                                    Expr::constant(c).at_opt(span.or(e2.span()))
475                                } else {
476                                    Expr::binary_op(op.clone(), e1, e2).at_opt(span)
477                                }
478                            }
479                            _ => Expr::binary_op(op.clone(), e1, e2).at_opt(span),
480                        }
481                    }
482                    ExprKind::UnaryOp(UnOp::Not, e) => {
483                        let e = e.fold_with(self);
484                        match e.kind() {
485                            ExprKind::Constant(Constant::Bool(b)) => {
486                                Expr::constant(Constant::Bool(!b))
487                            }
488                            ExprKind::UnaryOp(UnOp::Not, e) => e.clone(),
489                            ExprKind::BinaryOp(BinOp::Eq, e1, e2) => {
490                                Expr::binary_op(BinOp::Ne, e1, e2).at_opt(span)
491                            }
492                            _ => Expr::unary_op(UnOp::Not, e).at_opt(span),
493                        }
494                    }
495                    ExprKind::IfThenElse(p, e1, e2) => {
496                        let p = p.fold_with(self);
497                        if p.is_trivially_true() {
498                            e1.fold_with(self).at_opt(span)
499                        } else if p.is_trivially_false() {
500                            e2.fold_with(self).at_opt(span)
501                        } else {
502                            Expr::ite(p, e1.fold_with(self), e2.fold_with(self)).at_opt(span)
503                        }
504                    }
505                    _ => expr.super_fold_with(self),
506                }
507            }
508        }
509        self.fold_with(&mut Simplify { assumed_preds })
510    }
511
512    pub fn to_loc(&self) -> Option<Loc> {
513        match self.kind() {
514            ExprKind::Local(local) => Some(Loc::Local(*local)),
515            ExprKind::Var(var) => Some(Loc::Var(*var)),
516            _ => None,
517        }
518    }
519
520    pub fn to_path(&self) -> Option<Path> {
521        let mut expr = self;
522        let mut proj = vec![];
523        while let ExprKind::PathProj(e, field) = expr.kind() {
524            proj.push(*field);
525            expr = e;
526        }
527        proj.reverse();
528        Some(Path::new(expr.to_loc()?, proj))
529    }
530
531    pub fn is_abs(&self) -> bool {
532        matches!(self.kind(), ExprKind::Abs(..))
533    }
534
535    /// Whether this is an aggregate expression with no fields.
536    pub fn is_unit(&self) -> bool {
537        matches!(self.kind(), ExprKind::Tuple(flds) if flds.is_empty())
538            || matches!(self.kind(), ExprKind::Ctor(Ctor::Struct(_), flds) if flds.is_empty())
539    }
540
541    pub fn eta_expand_abs(&self, inputs: &BoundVariableKinds, output: Sort) -> Lambda {
542        let args = (0..inputs.len())
543            .map(|idx| Expr::bvar(INNERMOST, BoundVar::from_usize(idx), BoundReftKind::Annon))
544            .collect();
545        let body = Expr::app(self, args);
546        Lambda::bind_with_vars(body, inputs.clone(), output)
547    }
548
549    /// Applies a field projection to an expression and optimistically try to beta reduce it
550    pub fn proj_and_reduce(&self, proj: FieldProj) -> Expr {
551        match self.kind() {
552            ExprKind::Tuple(flds) | ExprKind::Ctor(Ctor::Struct(_), flds) => {
553                flds[proj.field_idx() as usize].clone()
554            }
555            _ => Expr::field_proj(self.clone(), proj),
556        }
557    }
558
559    pub fn flatten_conjs(&self) -> Vec<&Expr> {
560        fn go<'a>(e: &'a Expr, vec: &mut Vec<&'a Expr>) {
561            if let ExprKind::BinaryOp(BinOp::And, e1, e2) = e.kind() {
562                go(e1, vec);
563                go(e2, vec);
564            } else {
565                vec.push(e);
566            }
567        }
568        let mut vec = vec![];
569        go(self, &mut vec);
570        vec
571    }
572
573    pub fn has_evars(&self) -> bool {
574        struct HasEvars;
575
576        impl TypeVisitor for HasEvars {
577            type BreakTy = ();
578            fn visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::BreakTy> {
579                if let ExprKind::Var(Var::EVar(_)) = expr.kind() {
580                    ControlFlow::Break(())
581                } else {
582                    expr.super_visit_with(self)
583                }
584            }
585        }
586
587        self.visit_with(&mut HasEvars).is_break()
588    }
589
590    pub fn erase_spans(&self) -> Expr {
591        struct SpanEraser;
592        impl TypeFolder for SpanEraser {
593            fn fold_expr(&mut self, e: &Expr) -> Expr {
594                e.super_fold_with(self).at_opt(None)
595            }
596        }
597        self.fold_with(&mut SpanEraser)
598    }
599}
600
601#[derive(Clone, Copy, PartialEq, Eq, Hash, TyEncodable, TyDecodable, Debug)]
602pub struct ESpan {
603    /// The top-level span information
604    pub span: Span,
605    /// The span for the (base) call-site for def-expanded spans
606    pub base: Option<Span>,
607}
608
609impl ESpan {
610    pub fn new(span: Span) -> Self {
611        Self { span, base: None }
612    }
613
614    pub fn with_base(&self, espan: ESpan) -> Self {
615        Self { span: self.span, base: Some(espan.span) }
616    }
617}
618
619#[derive(
620    Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable, Debug, TypeFoldable, TypeVisitable,
621)]
622pub enum BinOp {
623    Iff,
624    Imp,
625    Or,
626    And,
627    Eq,
628    Ne,
629    Gt(Sort),
630    Ge(Sort),
631    Lt(Sort),
632    Le(Sort),
633    Add(Sort),
634    Sub(Sort),
635    Mul(Sort),
636    Div(Sort),
637    Mod(Sort),
638    BitAnd,
639    BitOr,
640    BitShl,
641    BitShr,
642}
643
644#[derive(Clone, Copy, PartialEq, Eq, Hash, Encodable, Debug, Decodable)]
645pub enum UnOp {
646    Not,
647    Neg,
648}
649
650#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, TyEncodable, TyDecodable)]
651pub enum Ctor {
652    /// for indices represented as `struct` in the refinement logic (e.g. using `refined_by` annotations)
653    Struct(DefId),
654    /// for indices represented as  `enum` in the refinement logic (e.g. using `reflected` annotations)
655    Enum(DefId, VariantIdx),
656}
657
658impl Ctor {
659    pub fn def_id(&self) -> DefId {
660        match self {
661            Self::Struct(def_id) | Self::Enum(def_id, _) => *def_id,
662        }
663    }
664
665    pub fn variant_idx(&self) -> VariantIdx {
666        match self {
667            Self::Struct(_) => FIRST_VARIANT,
668            Self::Enum(_, variant_idx) => *variant_idx,
669        }
670    }
671
672    fn is_enum(&self) -> bool {
673        matches!(self, Self::Enum(..))
674    }
675}
676
677#[derive(Clone, PartialEq, Eq, Hash, TyEncodable, Debug, TyDecodable)]
678pub enum ExprKind {
679    Var(Var),
680    Local(Local),
681    Constant(Constant),
682    ConstDefId(DefId),
683    BinaryOp(BinOp, Expr, Expr),
684    GlobalFunc(SpecFuncKind),
685    UnaryOp(UnOp, Expr),
686    FieldProj(Expr, FieldProj),
687    /// A variant used in the logic to represent a variant of an ADT as a pair of the `DefId` and variant-index
688    Ctor(Ctor, List<Expr>),
689    Tuple(List<Expr>),
690    PathProj(Expr, FieldIdx),
691    IfThenElse(Expr, Expr, Expr),
692    KVar(KVar),
693    Alias(AliasReft, List<Expr>),
694    Let(Expr, Binder<Expr>),
695    /// Function application. The syntax allows arbitrary expressions in function position, but in
696    /// practice we are restricted by what's possible to encode in fixpoint. In a nutshell, we need
697    /// to make sure that expressions that can't be encoded are eliminated before we generate the
698    /// fixpoint constraint. Most notably, lambda abstractions have to be fully applied before
699    /// encoding into fixpoint (except when they appear as an index at the top-level).
700    App(Expr, List<Expr>),
701    /// Lambda abstractions. They are purely syntactic and we don't encode them in the logic. As such,
702    /// they have some syntactic restrictions that we must carefully maintain:
703    ///
704    /// 1. They can appear as an index at the top level.
705    /// 2. We can only substitute an abstraction for a variable in function position (or as an index).
706    ///    More generally, we need to partially evaluate expressions such that all abstractions in
707    ///    non-index position are eliminated before encoding into fixpoint. Right now, the
708    ///    implementation only evaluates abstractions that are immediately applied to arguments,
709    ///    thus the restriction.
710    Abs(Lambda),
711
712    /// Bounded quantifiers `exists i in 0..4 { pred(i) }` and `forall i in 0..4 { pred(i) }`.
713    BoundedQuant(fhir::QuantKind, fhir::Range, Binder<Expr>),
714    /// A hole is an expression that must be inferred either *semantically* by generating a kvar or
715    /// *syntactically* by generating an evar. Whether a hole can be inferred semantically or
716    /// syntactically depends on the position it appears: only holes appearing in predicate position
717    /// can be inferred with a kvar (provided it satisfies the fixpoint horn constraints) and only
718    /// holes used as an index (a position that fully determines their value) can be inferred with
719    /// an evar.
720    ///
721    /// Holes are implicitly defined in a scope, i.e., their solution could mention free and bound
722    /// variables in this scope. This must be considered when generating an inference variables for
723    /// them (either evar or kvar). In fact, the main reason we have holes is that we want to
724    /// decouple the places where we generate holes (where we don't want to worry about the scope),
725    /// and the places where we generate inference variable for them (where we do need to worry
726    /// about the scope).
727    Hole(HoleKind),
728    ForAll(Binder<Expr>),
729}
730
731impl ExprKind {
732    fn intern(self) -> Expr {
733        Expr { kind: Interned::new(self), espan: None }
734    }
735}
736
737#[derive(Clone, Copy, PartialEq, Eq, Hash, TyEncodable, TyDecodable, Debug)]
738pub enum AggregateKind {
739    Tuple(usize),
740    Adt(DefId),
741}
742
743impl AggregateKind {
744    pub fn to_proj(self, field: u32) -> FieldProj {
745        match self {
746            AggregateKind::Tuple(arity) => FieldProj::Tuple { arity, field },
747            AggregateKind::Adt(def_id) => FieldProj::Adt { def_id, field },
748        }
749    }
750}
751
752#[derive(Clone, Copy, PartialEq, Eq, Hash, TyEncodable, TyDecodable, Debug)]
753pub enum FieldProj {
754    Tuple { arity: usize, field: u32 },
755    Adt { def_id: DefId, field: u32 },
756}
757
758impl FieldProj {
759    pub fn arity(&self, genv: GlobalEnv) -> QueryResult<usize> {
760        match self {
761            FieldProj::Tuple { arity, .. } => Ok(*arity),
762            FieldProj::Adt { def_id, .. } => {
763                Ok(genv.adt_sort_def_of(*def_id)?.struct_variant().fields())
764            }
765        }
766    }
767
768    pub fn field_idx(&self) -> u32 {
769        match self {
770            FieldProj::Tuple { field, .. } | FieldProj::Adt { field, .. } => *field,
771        }
772    }
773}
774
775/// The position where a [hole] appears. This determines how it will be inferred. This is related
776/// to, but not the same as, an [`InferMode`].
777///
778/// [`InferMode`]: super::InferMode
779/// [hole]: ExprKind::Hole
780#[derive(
781    Debug, Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable, TypeFoldable, TypeVisitable,
782)]
783pub enum HoleKind {
784    /// A hole in predicate position (e.g., the predicate in a [`TyKind::Constr`]). It will be
785    /// inferred by generating a kvar.
786    ///
787    /// [`TyKind::Constr`]: super::TyKind::Constr
788    Pred,
789    /// A hole used as a refinement argument or index. It will be inferred by generating an evar.
790    /// The expression filling the hole must have the provided sort.
791    ///
792    /// NOTE(nilehmann) we used to require the `Sort` for generating the evar because we needed it
793    /// to eta-expand aggregate sorts. We've since removed this behavior but I'm keeping it here
794    /// just in case. We could remove in case it becomes too problematic.
795    Expr(Sort),
796}
797
798/// In theory a kvar is just an unknown predicate that can use some variables in scope. In practice,
799/// fixpoint makes a difference between the first and the rest of the arguments, the first one being
800/// the kvar's *self argument*. Fixpoint will only instantiate qualifiers that use the self argument.
801/// Flux generalizes the self argument to be a list. We call the rest of the arguments the *scope*.
802#[derive(Clone, PartialEq, Eq, Hash, TyEncodable, TyDecodable, TypeVisitable, TypeFoldable)]
803pub struct KVar {
804    pub kvid: KVid,
805    /// The number of arguments consider to be *self arguments*.
806    pub self_args: usize,
807    /// The list of *all* arguments with the self arguments at the beginning, i.e., the
808    /// list of self arguments followed by the scope.
809    pub args: List<Expr>,
810}
811
812#[derive(Clone, Debug, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Encodable, Decodable)]
813pub struct EarlyReftParam {
814    pub index: u32,
815    pub name: Symbol,
816}
817
818#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, Encodable, Decodable, Debug)]
819pub struct BoundReft {
820    pub var: BoundVar,
821    pub kind: BoundReftKind,
822}
823
824#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, TyEncodable, TyDecodable)]
825pub enum Var {
826    Free(Name),
827    Bound(DebruijnIndex, BoundReft),
828    EarlyParam(EarlyReftParam),
829    EVar(EVid),
830    ConstGeneric(ParamConst),
831}
832
833#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord, TyEncodable, TyDecodable)]
834pub struct Path {
835    pub loc: Loc,
836    projection: List<FieldIdx>,
837}
838
839#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, TyEncodable, TyDecodable)]
840pub enum Loc {
841    Local(Local),
842    Var(Var),
843}
844
845newtype_index! {
846    /// *E*xistential *v*ariable *id*
847    #[debug_format = "?{}e"]
848    #[orderable]
849    #[encodable]
850    pub struct EVid {}
851}
852
853newtype_index! {
854    #[debug_format = "$k{}"]
855    #[encodable]
856    pub struct KVid {}
857}
858
859newtype_index! {
860    #[debug_format = "a{}"]
861    #[orderable]
862    #[encodable]
863    pub struct Name {}
864}
865
866impl KVar {
867    pub fn new(kvid: KVid, self_args: usize, args: Vec<Expr>) -> Self {
868        KVar { kvid, self_args, args: List::from_vec(args) }
869    }
870
871    fn self_args(&self) -> &[Expr] {
872        &self.args[..self.self_args]
873    }
874
875    fn scope(&self) -> &[Expr] {
876        &self.args[self.self_args..]
877    }
878}
879
880impl Var {
881    pub fn to_expr(&self) -> Expr {
882        Expr::var(*self)
883    }
884}
885
886impl Path {
887    pub fn new(loc: Loc, projection: impl Into<List<FieldIdx>>) -> Path {
888        Path { loc, projection: projection.into() }
889    }
890
891    pub fn projection(&self) -> &[FieldIdx] {
892        &self.projection[..]
893    }
894
895    pub fn to_expr(&self) -> Expr {
896        self.projection
897            .iter()
898            .fold(self.loc.to_expr(), |e, f| Expr::path_proj(e, *f))
899    }
900
901    pub fn to_loc(&self) -> Option<Loc> {
902        if self.projection.is_empty() { Some(self.loc) } else { None }
903    }
904}
905
906impl Loc {
907    pub fn to_expr(&self) -> Expr {
908        match self {
909            Loc::Local(local) => Expr::local(*local),
910            Loc::Var(var) => Expr::var(*var),
911        }
912    }
913}
914
915macro_rules! impl_ops {
916    ($($op:ident: $method:ident),*) => {$(
917        impl<Rhs> std::ops::$op<Rhs> for Expr
918        where
919            Rhs: Into<Expr>,
920        {
921            type Output = Expr;
922
923            fn $method(self, rhs: Rhs) -> Self::Output {
924                let sort = crate::rty::Sort::Int;
925                Expr::binary_op(BinOp::$op(sort), self, rhs)
926            }
927        }
928
929        impl<Rhs> std::ops::$op<Rhs> for &Expr
930        where
931            Rhs: Into<Expr>,
932        {
933            type Output = Expr;
934
935            fn $method(self, rhs: Rhs) -> Self::Output {
936                let sort = crate::rty::Sort::Int;
937                Expr::binary_op(BinOp::$op(sort), self, rhs)
938            }
939        }
940    )*};
941}
942impl_ops!(Add: add, Sub: sub, Mul: mul, Div: div);
943
944impl From<i32> for Expr {
945    fn from(value: i32) -> Self {
946        Expr::constant(Constant::from(value))
947    }
948}
949
950impl From<&Expr> for Expr {
951    fn from(e: &Expr) -> Self {
952        e.clone()
953    }
954}
955
956impl From<Path> for Expr {
957    fn from(path: Path) -> Self {
958        path.to_expr()
959    }
960}
961
962impl From<Name> for Expr {
963    fn from(name: Name) -> Self {
964        Expr::fvar(name)
965    }
966}
967
968impl From<Var> for Expr {
969    fn from(var: Var) -> Self {
970        Expr::var(var)
971    }
972}
973
974impl From<Loc> for Path {
975    fn from(loc: Loc) -> Self {
976        Path::new(loc, vec![])
977    }
978}
979
980impl From<Name> for Loc {
981    fn from(name: Name) -> Self {
982        Loc::Var(Var::Free(name))
983    }
984}
985
986impl From<Local> for Loc {
987    fn from(local: Local) -> Self {
988        Loc::Local(local)
989    }
990}
991
992#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Encodable, Decodable)]
993pub struct Real(pub u128);
994
995impl liquid_fixpoint::FixpointFmt for Real {
996    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
997        write!(f, "{}.0", self.0)
998    }
999}
1000
1001#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Encodable, Decodable)]
1002pub enum Constant {
1003    Int(BigInt),
1004    Real(Real),
1005    Bool(bool),
1006    Str(Symbol),
1007    Char(char),
1008    BitVec(u128, u32),
1009}
1010
1011impl Constant {
1012    pub const ZERO: Constant = Constant::Int(BigInt::ZERO);
1013    pub const ONE: Constant = Constant::Int(BigInt::ONE);
1014    pub const TRUE: Constant = Constant::Bool(true);
1015
1016    fn to_bool(self) -> Option<bool> {
1017        match self {
1018            Constant::Bool(b) => Some(b),
1019            _ => None,
1020        }
1021    }
1022
1023    fn to_int(self) -> Option<BigInt> {
1024        match self {
1025            Constant::Int(n) => Some(n),
1026            _ => None,
1027        }
1028    }
1029
1030    pub fn iff(&self, other: &Constant) -> Option<Constant> {
1031        let b1 = self.to_bool()?;
1032        let b2 = other.to_bool()?;
1033        Some(Constant::Bool(b1 == b2))
1034    }
1035
1036    pub fn imp(&self, other: &Constant) -> Option<Constant> {
1037        let b1 = self.to_bool()?;
1038        let b2 = other.to_bool()?;
1039        Some(Constant::Bool(!b1 || b2))
1040    }
1041
1042    pub fn or(&self, other: &Constant) -> Option<Constant> {
1043        let b1 = self.to_bool()?;
1044        let b2 = other.to_bool()?;
1045        Some(Constant::Bool(b1 || b2))
1046    }
1047
1048    pub fn and(&self, other: &Constant) -> Option<Constant> {
1049        let b1 = self.to_bool()?;
1050        let b2 = other.to_bool()?;
1051        Some(Constant::Bool(b1 && b2))
1052    }
1053
1054    pub fn eq(&self, other: &Constant) -> Constant {
1055        Constant::Bool(*self == *other)
1056    }
1057
1058    pub fn ne(&self, other: &Constant) -> Constant {
1059        Constant::Bool(*self != *other)
1060    }
1061
1062    pub fn gt(&self, other: &Constant) -> Option<Constant> {
1063        let n1 = self.to_int()?;
1064        let n2 = other.to_int()?;
1065        Some(Constant::Bool(n1 > n2))
1066    }
1067
1068    pub fn ge(&self, other: &Constant) -> Option<Constant> {
1069        let n1 = self.to_int()?;
1070        let n2 = other.to_int()?;
1071        Some(Constant::Bool(n1 >= n2))
1072    }
1073
1074    pub fn from_scalar_int<'tcx, T>(tcx: TyCtxt<'tcx>, scalar: ScalarInt, t: &T) -> Option<Self>
1075    where
1076        T: ToRustc<'tcx, T = rustc_middle::ty::Ty<'tcx>>,
1077    {
1078        use rustc_middle::ty::TyKind;
1079        let ty = t.to_rustc(tcx);
1080        match ty.kind() {
1081            TyKind::Int(int_ty) => Some(Constant::from(scalar_to_int(tcx, scalar, *int_ty))),
1082            TyKind::Uint(uint_ty) => Some(Constant::from(scalar_to_uint(tcx, scalar, *uint_ty))),
1083            TyKind::Bool => {
1084                let b = scalar_to_bits(tcx, scalar, ty)?;
1085                Some(Constant::Bool(b != 0))
1086            }
1087            TyKind::Char => {
1088                let b = scalar_to_bits(tcx, scalar, ty)?;
1089                Some(Constant::Char(char::from_u32(b as u32)?))
1090            }
1091            _ => bug!(),
1092        }
1093    }
1094
1095    /// See [`BigInt::int_min`]
1096    pub fn int_min(bit_width: u32) -> Constant {
1097        Constant::Int(BigInt::int_min(bit_width))
1098    }
1099
1100    /// See [`BigInt::int_max`]
1101    pub fn int_max(bit_width: u32) -> Constant {
1102        Constant::Int(BigInt::int_max(bit_width))
1103    }
1104
1105    /// See [`BigInt::uint_max`]
1106    pub fn uint_max(bit_width: u32) -> Constant {
1107        Constant::Int(BigInt::uint_max(bit_width))
1108    }
1109}
1110
1111impl From<i32> for Constant {
1112    fn from(c: i32) -> Self {
1113        Constant::Int(c.into())
1114    }
1115}
1116
1117impl From<usize> for Constant {
1118    fn from(u: usize) -> Self {
1119        Constant::Int(u.into())
1120    }
1121}
1122
1123impl From<u128> for Constant {
1124    fn from(c: u128) -> Self {
1125        Constant::Int(c.into())
1126    }
1127}
1128
1129impl From<i128> for Constant {
1130    fn from(c: i128) -> Self {
1131        Constant::Int(c.into())
1132    }
1133}
1134
1135impl From<bool> for Constant {
1136    fn from(b: bool) -> Self {
1137        Constant::Bool(b)
1138    }
1139}
1140
1141impl From<Symbol> for Constant {
1142    fn from(s: Symbol) -> Self {
1143        Constant::Str(s)
1144    }
1145}
1146
1147impl From<char> for Constant {
1148    fn from(c: char) -> Self {
1149        Constant::Char(c)
1150    }
1151}
1152
1153impl_internable!(ExprKind);
1154impl_slice_internable!(Expr, KVar);
1155
1156#[derive(Debug)]
1157pub struct FieldBind<T> {
1158    pub name: Symbol,
1159    pub value: T,
1160}
1161
1162impl<T: Pretty> Pretty for FieldBind<T> {
1163    fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1164        w!(cx, f, "{}: {:?}", ^self.name, &self.value)
1165    }
1166}
1167
1168pub(crate) mod pretty {
1169    use flux_rustc_bridge::def_id_to_string;
1170
1171    use super::*;
1172    use crate::{name_of_thy_func, rty::pretty::nested_with_bound_vars};
1173
1174    #[derive(PartialEq, Eq, PartialOrd, Ord)]
1175    enum Precedence {
1176        Iff,
1177        Imp,
1178        Or,
1179        And,
1180        Cmp,
1181        Bitvec,
1182        AddSub,
1183        MulDiv,
1184    }
1185
1186    impl BinOp {
1187        fn precedence(&self) -> Precedence {
1188            match self {
1189                BinOp::Iff => Precedence::Iff,
1190                BinOp::Imp => Precedence::Imp,
1191                BinOp::Or => Precedence::Or,
1192                BinOp::And => Precedence::And,
1193                BinOp::Eq
1194                | BinOp::Ne
1195                | BinOp::Gt(_)
1196                | BinOp::Lt(_)
1197                | BinOp::Ge(_)
1198                | BinOp::Le(_) => Precedence::Cmp,
1199                BinOp::Add(_) | BinOp::Sub(_) => Precedence::AddSub,
1200                BinOp::Mul(_) | BinOp::Div(_) | BinOp::Mod(_) => Precedence::MulDiv,
1201                BinOp::BitAnd | BinOp::BitOr | BinOp::BitShl | BinOp::BitShr => Precedence::Bitvec,
1202            }
1203        }
1204    }
1205
1206    impl Precedence {
1207        pub fn is_associative(&self) -> bool {
1208            !matches!(self, Precedence::Imp | Precedence::Cmp)
1209        }
1210    }
1211
1212    pub fn should_parenthesize(op: &BinOp, child: &Expr) -> bool {
1213        if let ExprKind::BinaryOp(child_op, ..) = child.kind() {
1214            child_op.precedence() < op.precedence()
1215                || (child_op.precedence() == op.precedence() && !op.precedence().is_associative())
1216        } else {
1217            false
1218        }
1219    }
1220
1221    impl Pretty for Ctor {
1222        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1223            match self {
1224                Ctor::Struct(def_id) => {
1225                    w!(cx, f, "{:?}", def_id)
1226                }
1227                Ctor::Enum(def_id, variant_idx) => {
1228                    w!(cx, f, "{:?}::{:?}", def_id, ^variant_idx)
1229                }
1230            }
1231        }
1232    }
1233
1234    impl Pretty for fhir::QuantKind {
1235        fn fmt(&self, _cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1236            match self {
1237                fhir::QuantKind::Exists => w!(cx, f, "∃"),
1238                fhir::QuantKind::Forall => w!(cx, f, "∀"),
1239            }
1240        }
1241    }
1242
1243    impl Pretty for Expr {
1244        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1245            let e = if cx.simplify_exprs {
1246                self.simplify(&SnapshotMap::default())
1247            } else {
1248                self.clone()
1249            };
1250            match e.kind() {
1251                ExprKind::Var(var) => w!(cx, f, "{:?}", var),
1252                ExprKind::Local(local) => w!(cx, f, "{:?}", ^local),
1253                ExprKind::ConstDefId(did) => w!(cx, f, "{}", ^def_id_to_string(*did)),
1254                ExprKind::Constant(c) => w!(cx, f, "{:?}", c),
1255                ExprKind::BinaryOp(op, e1, e2) => {
1256                    if should_parenthesize(op, e1) {
1257                        w!(cx, f, "({:?})", e1)?;
1258                    } else {
1259                        w!(cx, f, "{:?}", e1)?;
1260                    }
1261                    if matches!(op, BinOp::Div(_)) {
1262                        w!(cx, f, "{:?}", op)?;
1263                    } else {
1264                        w!(cx, f, " {:?} ", op)?;
1265                    }
1266                    if should_parenthesize(op, e2) {
1267                        w!(cx, f, "({:?})", e2)?;
1268                    } else {
1269                        w!(cx, f, "{:?}", e2)?;
1270                    }
1271                    Ok(())
1272                }
1273                ExprKind::UnaryOp(op, e) => {
1274                    if e.is_atom() {
1275                        w!(cx, f, "{:?}{:?}", op, e)
1276                    } else {
1277                        w!(cx, f, "{:?}({:?})", op, e)
1278                    }
1279                }
1280                ExprKind::FieldProj(e, proj) => {
1281                    if e.is_atom() {
1282                        w!(cx, f, "{:?}.{}", e, ^fmt_field_proj(cx, *proj))
1283                    } else {
1284                        w!(cx, f, "({:?}).{}", e, ^fmt_field_proj(cx, *proj))
1285                    }
1286                }
1287                ExprKind::Tuple(flds) => {
1288                    if let [e] = &flds[..] {
1289                        w!(cx, f, "({:?},)", e)
1290                    } else {
1291                        w!(cx, f, "({:?})", join!(", ", flds))
1292                    }
1293                }
1294                ExprKind::Ctor(ctor, flds) => {
1295                    let def_id = ctor.def_id();
1296                    if let Some(adt_sort_def) = cx.adt_sort_def_of(def_id) {
1297                        let variant = adt_sort_def.variant(ctor.variant_idx()).field_names();
1298                        let fields = iter::zip(variant, flds)
1299                            .map(|(name, value)| FieldBind { name: *name, value: value.clone() })
1300                            .collect_vec();
1301                        match ctor {
1302                            Ctor::Struct(_) => {
1303                                w!(cx, f, "{:?} {{ {:?} }}", def_id, join!(", ", fields))
1304                            }
1305                            Ctor::Enum(_, idx) => {
1306                                if fields.is_empty() {
1307                                    w!(cx, f, "{:?}::{:?}", def_id, ^idx.index())
1308                                } else {
1309                                    w!(cx, f, "{:?}::{:?}({:?})", def_id, ^idx.index(), join!(", ", fields))
1310                                }
1311                            }
1312                        }
1313                    } else {
1314                        match ctor {
1315                            Ctor::Struct(_) => {
1316                                w!(cx, f, "{:?} {{ {:?} }}", def_id, join!(", ", flds))
1317                            }
1318                            Ctor::Enum(_, idx) => {
1319                                w!(cx, f, "{:?}::{:?} {{ {:?} }}", def_id, ^idx, join!(", ", flds))
1320                            }
1321                        }
1322                    }
1323                }
1324                ExprKind::PathProj(e, field) => {
1325                    if e.is_atom() {
1326                        w!(cx, f, "{:?}.{:?}", e, field)
1327                    } else {
1328                        w!(cx, f, "({:?}).{:?}", e, field)
1329                    }
1330                }
1331                ExprKind::App(func, args) => {
1332                    w!(cx, f, "{:?}({})",
1333                        parens!(func, !func.is_atom()),
1334                        ^args
1335                            .iter()
1336                            .format_with(", ", |arg, f| f(&format_args_cx!(cx, "{:?}", arg)))
1337                    )
1338                }
1339                ExprKind::IfThenElse(p, e1, e2) => {
1340                    w!(cx, f, "if {:?} {{ {:?} }} else {{ {:?} }}", p, e1, e2)
1341                }
1342                ExprKind::Hole(_) => {
1343                    w!(cx, f, "*")
1344                }
1345                ExprKind::KVar(kvar) => {
1346                    w!(cx, f, "{:?}", kvar)
1347                }
1348                ExprKind::Alias(alias, args) => {
1349                    w!(cx, f, "{:?}({:?})", alias, join!(", ", args))
1350                }
1351                ExprKind::Abs(lam) => {
1352                    w!(cx, f, "{:?}", lam)
1353                }
1354                ExprKind::GlobalFunc(SpecFuncKind::Def(did) | SpecFuncKind::Uif(did)) => {
1355                    w!(cx, f, "{}", ^did.name())
1356                }
1357                ExprKind::GlobalFunc(SpecFuncKind::Thy(itf)) => {
1358                    if let Some(name) = name_of_thy_func(*itf) {
1359                        w!(cx, f, "{}", ^name)
1360                    } else {
1361                        w!(cx, f, "<error>")
1362                    }
1363                }
1364                ExprKind::ForAll(expr) => {
1365                    let vars = expr.vars();
1366                    cx.with_bound_vars(vars, || {
1367                        if !vars.is_empty() {
1368                            cx.fmt_bound_vars(false, "∀", vars, ". ", f)?;
1369                        }
1370                        w!(cx, f, "{:?}", expr.skip_binder_ref())
1371                    })
1372                }
1373                ExprKind::BoundedQuant(kind, rng, body) => {
1374                    let vars = body.vars();
1375                    cx.with_bound_vars(vars, || {
1376                        w!(
1377                            cx,
1378                            f,
1379                            "{:?} {}..{} {{ {:?} }}",
1380                            kind,
1381                            ^rng.start,
1382                            ^rng.end,
1383                            body.skip_binder_ref()
1384                        )
1385                    })
1386                }
1387                ExprKind::Let(init, body) => {
1388                    let vars = body.vars();
1389                    cx.with_bound_vars(vars, || {
1390                        cx.fmt_bound_vars(false, "(let ", vars, " = ", f)?;
1391                        w!(cx, f, "{:?} in {:?})", init, body.skip_binder_ref())
1392                    })
1393                }
1394            }
1395        }
1396    }
1397
1398    fn fmt_field_proj(cx: &PrettyCx, proj: FieldProj) -> String {
1399        if let FieldProj::Adt { def_id, field } = proj
1400            && let Some(adt_sort_def) = cx.adt_sort_def_of(def_id)
1401        {
1402            format!("{}", adt_sort_def.struct_variant().field_names()[field as usize])
1403        } else {
1404            format!("{}", proj.field_idx())
1405        }
1406    }
1407
1408    impl Pretty for Constant {
1409        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1410            match self {
1411                Constant::Int(i) => w!(cx, f, "{i}"),
1412                Constant::BitVec(i, sz) => w!(cx, f, "bv({i}, {sz})"),
1413                Constant::Real(r) => w!(cx, f, "{}.0", ^r.0),
1414                Constant::Bool(b) => w!(cx, f, "{b}"),
1415                Constant::Str(sym) => w!(cx, f, "\"{sym}\""),
1416                Constant::Char(c) => w!(cx, f, "\'{c}\'"),
1417            }
1418        }
1419    }
1420
1421    impl Pretty for AliasReft {
1422        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1423            w!(cx, f, "<({:?}) as {:?}", &self.args[0], self.assoc_id.parent())?;
1424            let args = &self.args[1..];
1425            if !args.is_empty() {
1426                w!(cx, f, "<{:?}>", join!(", ", args))?;
1427            }
1428            w!(cx, f, ">::{}", ^self.assoc_id.name())
1429        }
1430    }
1431
1432    impl Pretty for Lambda {
1433        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1434            let vars = self.body.vars();
1435            cx.with_bound_vars(vars, || {
1436                cx.fmt_bound_vars(false, "λ", vars, ". ", f)?;
1437                w!(cx, f, "{:?}", self.body.as_ref().skip_binder())
1438            })
1439        }
1440    }
1441
1442    impl Pretty for Var {
1443        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1444            match self {
1445                Var::Bound(debruijn, var) => cx.fmt_bound_reft(*debruijn, *var, f),
1446                Var::EarlyParam(var) => w!(cx, f, "{}", ^var.name),
1447                Var::Free(name) => w!(cx, f, "{:?}", ^name),
1448                Var::EVar(evar) => w!(cx, f, "{:?}", ^evar),
1449                Var::ConstGeneric(param) => w!(cx, f, "{}", ^param.name),
1450            }
1451        }
1452    }
1453
1454    impl Pretty for KVar {
1455        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1456            w!(cx, f, "{:?}", ^self.kvid)?;
1457            match cx.kvar_args {
1458                KVarArgs::All => {
1459                    w!(
1460                        cx,
1461                        f,
1462                        "({:?})[{:?}]",
1463                        join!(", ", self.self_args()),
1464                        join!(", ", self.scope())
1465                    )?;
1466                }
1467                KVarArgs::SelfOnly => w!(cx, f, "({:?})", join!(", ", self.self_args()))?,
1468                KVarArgs::Hide => {}
1469            }
1470            Ok(())
1471        }
1472    }
1473
1474    impl Pretty for Path {
1475        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1476            w!(cx, f, "{:?}", &self.loc)?;
1477            for field in &self.projection {
1478                w!(cx, f, ".{}", ^u32::from(*field))?;
1479            }
1480            Ok(())
1481        }
1482    }
1483
1484    impl Pretty for Loc {
1485        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1486            match self {
1487                Loc::Local(local) => w!(cx, f, "{:?}", ^local),
1488                Loc::Var(var) => w!(cx, f, "{:?}", var),
1489            }
1490        }
1491    }
1492
1493    impl Pretty for BinOp {
1494        fn fmt(&self, _cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1495            match self {
1496                BinOp::Iff => w!(cx, f, "⇔"),
1497                BinOp::Imp => w!(cx, f, "⇒"),
1498                BinOp::Or => w!(cx, f, "∨"),
1499                BinOp::And => w!(cx, f, "∧"),
1500                BinOp::Eq => w!(cx, f, "="),
1501                BinOp::Ne => w!(cx, f, "≠"),
1502                BinOp::Gt(_) => w!(cx, f, ">"),
1503                BinOp::Ge(_) => w!(cx, f, "≥"),
1504                BinOp::Lt(_) => w!(cx, f, "<"),
1505                BinOp::Le(_) => w!(cx, f, "≤"),
1506                BinOp::Add(_) => w!(cx, f, "+"),
1507                BinOp::Sub(_) => w!(cx, f, "-"),
1508                BinOp::Mul(_) => w!(cx, f, "*"),
1509                BinOp::Div(_) => w!(cx, f, "/"),
1510                BinOp::Mod(_) => w!(cx, f, "mod"),
1511                BinOp::BitAnd => w!(cx, f, "&"),
1512                BinOp::BitOr => w!(cx, f, "|"),
1513                BinOp::BitShl => w!(cx, f, "<<"),
1514                BinOp::BitShr => w!(cx, f, ">>"),
1515            }
1516        }
1517    }
1518
1519    impl Pretty for UnOp {
1520        fn fmt(&self, _cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1521            match self {
1522                UnOp::Not => w!(cx, f, "¬"),
1523                UnOp::Neg => w!(cx, f, "-"),
1524            }
1525        }
1526    }
1527
1528    impl_debug_with_default_cx!(Expr, Loc, Path, Var, KVar, Lambda, AliasReft);
1529
1530    impl PrettyNested for Lambda {
1531        fn fmt_nested(&self, cx: &PrettyCx) -> Result<NestedString, fmt::Error> {
1532            nested_with_bound_vars(cx, "λ", self.body.vars(), None, |prefix| {
1533                let expr_d = self.body.skip_binder_ref().fmt_nested(cx)?;
1534                let text = format!("{}{}", prefix, expr_d.text);
1535                Ok(NestedString { text, children: expr_d.children, key: None })
1536            })
1537        }
1538    }
1539
1540    pub fn aggregate_nested(
1541        cx: &PrettyCx,
1542        ctor: &Ctor,
1543        flds: &[Expr],
1544        is_named: bool,
1545    ) -> Result<NestedString, fmt::Error> {
1546        let def_id = ctor.def_id();
1547        let mut text =
1548            if is_named && ctor.is_enum() { format_cx!(cx, "{:?}", ctor) } else { "".to_string() };
1549        if flds.is_empty() {
1550            // No fields, no index
1551            Ok(NestedString { text, children: None, key: None })
1552        } else if flds.len() == 1 {
1553            // Single field, inline index
1554            text += &format_cx!(cx, "{:?} ", flds[0].clone());
1555            Ok(NestedString { text, children: None, key: None })
1556        } else {
1557            let keys = if let Some(adt_sort_def) = cx.adt_sort_def_of(def_id) {
1558                adt_sort_def
1559                    .variant(ctor.variant_idx())
1560                    .field_names()
1561                    .iter()
1562                    .map(|name| format!("{}", name))
1563                    .collect_vec()
1564            } else {
1565                (0..flds.len()).map(|i| format!("arg{}", i)).collect_vec()
1566            };
1567            // Multiple fields, nested index
1568            text += "{..}";
1569            let mut children = vec![];
1570            for (key, fld) in iter::zip(keys, flds) {
1571                let fld_d = fld.fmt_nested(cx)?;
1572                children.push(NestedString { key: Some(key), ..fld_d });
1573            }
1574            Ok(NestedString { text, children: Some(children), key: None })
1575        }
1576    }
1577
1578    impl PrettyNested for Expr {
1579        fn fmt_nested(&self, cx: &PrettyCx) -> Result<NestedString, fmt::Error> {
1580            let e = if cx.simplify_exprs {
1581                self.simplify(&SnapshotMap::default())
1582            } else {
1583                self.clone()
1584            };
1585            match e.kind() {
1586                ExprKind::Var(..)
1587                | ExprKind::Local(..)
1588                | ExprKind::Constant(..)
1589                | ExprKind::ConstDefId(..)
1590                | ExprKind::Hole(..)
1591                | ExprKind::GlobalFunc(..)
1592                | ExprKind::KVar(..) => debug_nested(cx, &e),
1593
1594                ExprKind::IfThenElse(p, e1, e2) => {
1595                    let p_d = p.fmt_nested(cx)?;
1596                    let e1_d = e1.fmt_nested(cx)?;
1597                    let e2_d = e2.fmt_nested(cx)?;
1598                    let text = format!("(if {} then {} else {})", p_d.text, e1_d.text, e2_d.text);
1599                    let children = float_children(vec![p_d.children, e1_d.children, e2_d.children]);
1600                    Ok(NestedString { text, children, key: None })
1601                }
1602                ExprKind::BinaryOp(op, e1, e2) => {
1603                    let e1_d = e1.fmt_nested(cx)?;
1604                    let e2_d = e2.fmt_nested(cx)?;
1605                    let e1_text = if should_parenthesize(op, e1) {
1606                        format!("({})", e1_d.text)
1607                    } else {
1608                        e1_d.text
1609                    };
1610                    let e2_text = if should_parenthesize(op, e2) {
1611                        format!("({})", e2_d.text)
1612                    } else {
1613                        e2_d.text
1614                    };
1615                    let op_d = debug_nested(cx, op)?;
1616                    let op_text = if matches!(op, BinOp::Div(_)) {
1617                        op_d.text
1618                    } else {
1619                        format!(" {} ", op_d.text)
1620                    };
1621                    let text = format!("{}{}{}", e1_text, op_text, e2_text);
1622                    let children = float_children(vec![e1_d.children, e2_d.children]);
1623                    Ok(NestedString { text, children, key: None })
1624                }
1625                ExprKind::UnaryOp(op, e) => {
1626                    let e_d = e.fmt_nested(cx)?;
1627                    let op_d = debug_nested(cx, op)?;
1628                    let text = if e.is_atom() {
1629                        format!("{}{}", op_d.text, e_d.text)
1630                    } else {
1631                        format!("{}({})", op_d.text, e_d.text)
1632                    };
1633                    Ok(NestedString { text, children: e_d.children, key: None })
1634                }
1635                ExprKind::FieldProj(e, proj) => {
1636                    let e_d = e.fmt_nested(cx)?;
1637                    let text = if e.is_atom() {
1638                        format!("{}.{}", e_d.text, fmt_field_proj(cx, *proj))
1639                    } else {
1640                        format!("({}).{}", e_d.text, fmt_field_proj(cx, *proj))
1641                    };
1642                    Ok(NestedString { text, children: e_d.children, key: None })
1643                }
1644                ExprKind::Tuple(flds) => {
1645                    let mut texts = vec![];
1646                    let mut kidss = vec![];
1647                    for e in flds {
1648                        let e_d = e.fmt_nested(cx)?;
1649                        texts.push(e_d.text);
1650                        kidss.push(e_d.children);
1651                    }
1652                    let text = if let [e] = &texts[..] {
1653                        format!("({},)", e)
1654                    } else {
1655                        format!("({})", texts.join(", "))
1656                    };
1657                    let children = float_children(kidss);
1658                    Ok(NestedString { text, children, key: None })
1659                }
1660                ExprKind::Ctor(ctor, flds) => aggregate_nested(cx, ctor, flds, true),
1661                ExprKind::PathProj(e, field) => {
1662                    let e_d = e.fmt_nested(cx)?;
1663                    let text = if e.is_atom() {
1664                        format!("{}.{:?}", e_d.text, field)
1665                    } else {
1666                        format!("({}).{:?}", e_d.text, field)
1667                    };
1668                    Ok(NestedString { text, children: e_d.children, key: None })
1669                }
1670                ExprKind::Alias(alias, args) => {
1671                    let mut texts = vec![];
1672                    let mut kidss = vec![];
1673                    for arg in args {
1674                        let arg_d = arg.fmt_nested(cx)?;
1675                        texts.push(arg_d.text);
1676                        kidss.push(arg_d.children);
1677                    }
1678                    let text = format_cx!(cx, "{:?}({:?})", alias, texts.join(", "));
1679                    let children = float_children(kidss);
1680                    Ok(NestedString { text, children, key: None })
1681                }
1682                ExprKind::App(func, args) => {
1683                    let func_d = func.fmt_nested(cx)?;
1684                    let mut texts = vec![];
1685                    let mut kidss = vec![func_d.children];
1686                    for arg in args {
1687                        let arg_d = arg.fmt_nested(cx)?;
1688                        texts.push(arg_d.text);
1689                        kidss.push(arg_d.children);
1690                    }
1691                    let text = if func.is_atom() {
1692                        format!("{}({})", func_d.text, texts.join(", "))
1693                    } else {
1694                        format!("({})({})", func_d.text, texts.join(", "))
1695                    };
1696                    let children = float_children(kidss);
1697                    Ok(NestedString { text, children, key: None })
1698                }
1699                ExprKind::Abs(lambda) => lambda.fmt_nested(cx),
1700                ExprKind::Let(init, body) => {
1701                    // FIXME this is very wrong!
1702                    nested_with_bound_vars(cx, "let", body.vars(), None, |prefix| {
1703                        let body = body.skip_binder_ref().fmt_nested(cx)?;
1704                        let text = format!("{:?} {}{}", init, prefix, body.text);
1705                        Ok(NestedString { text, children: body.children, key: None })
1706                    })
1707                }
1708                ExprKind::BoundedQuant(kind, rng, body) => {
1709                    let left = match kind {
1710                        fhir::QuantKind::Forall => "∀",
1711                        fhir::QuantKind::Exists => "∃",
1712                    };
1713                    let right = Some(format!(" in {}..{}", rng.start, rng.end));
1714
1715                    nested_with_bound_vars(cx, left, body.vars(), right, |all_str| {
1716                        let expr_d = body.as_ref().skip_binder().fmt_nested(cx)?;
1717                        let text = format!("{}{}", all_str, expr_d.text);
1718                        Ok(NestedString { text, children: expr_d.children, key: None })
1719                    })
1720                }
1721                ExprKind::ForAll(expr) => {
1722                    nested_with_bound_vars(cx, "∀", expr.vars(), None, |all_str| {
1723                        let expr_d = expr.as_ref().skip_binder().fmt_nested(cx)?;
1724                        let text = format!("{}{}", all_str, expr_d.text);
1725                        Ok(NestedString { text, children: expr_d.children, key: None })
1726                    })
1727                }
1728            }
1729        }
1730    }
1731}
flux_middle/rty/expr.rs

flux_middle/rty/
expr.rs
