flux_infer/
lean_format.rs

1use core::fmt;
2use std::{fmt::Write, iter, sync::LazyLock};
3
4use flux_common::{
5    bug,
6    dbg::{self, as_subscript},
7};
8use flux_middle::{
9    def_id::FluxDefId,
10    global_env::GlobalEnv,
11    rty::{KVid, PrettyMap, PrettyVar},
12};
13use itertools::Itertools;
14use liquid_fixpoint::{FixpointFmt, Identifier, Quantifier, ThyFunc};
15use rustc_data_structures::{fx::FxIndexSet, unord::UnordMap};
16use rustc_hir::{
17    def::DefKind,
18    def_id::DefId,
19    definitions::{DefPathData, DefPathDataName},
20};
21use rustc_middle::ty::{
22    TyCtxt,
23    print::{PrintTraitRefExt, with_forced_trimmed_paths},
24};
25use rustc_span::symbol::kw;
26
27use crate::fixpoint_encoding::{
28    ClosedSolution, InterpretedConst,
29    fixpoint::{
30        self, AdtId, BinOp, BinRel, Constant, Constraint, DataDecl, DataField, DataSort, Expr,
31        FunDef, FunSort, GlobalVar, KVarDecl, LocalVar, Pred, Sort, SortCtor, SortDecl, Var,
32    },
33};
34
35pub struct LeanCtxt<'a, 'genv, 'tcx> {
36    pub genv: GlobalEnv<'genv, 'tcx>,
37    pub pretty_var_map: &'a PrettyMap<LocalVar>,
38    pub adt_map: &'a FxIndexSet<DefId>,
39    pub opaque_adt_map: &'a [(FluxDefId, SortDecl)],
40    /// Maps the per-run `GlobalVar` index of a primop constant to its stable Lean name.
41    pub primop_var_map: &'a UnordMap<GlobalVar, String>,
42    pub hide_sort_vars: bool,
43}
44
45pub struct WithLeanCtxt<'a, 'b, 'genv, 'tcx, T> {
46    pub item: T,
47    pub cx: &'a LeanCtxt<'b, 'genv, 'tcx>,
48}
49
50impl<'a, 'b, 'genv, 'tcx, T: LeanFmt> fmt::Display for WithLeanCtxt<'a, 'b, 'genv, 'tcx, T> {
51    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
52        self.item.lean_fmt(f, self.cx)
53    }
54}
55
56pub trait LeanFmt {
57    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result;
58}
59
60pub struct LeanKConstraint<'a> {
61    pub theorem_name: &'a str,
62    pub kvars: &'a [KVarDecl],
63    /// Quantified above `kvars`' existentials so a kvar's solution can depend on them.
64    pub const_binds: &'a [fixpoint::Bind],
65    pub constr: &'a Constraint,
66    pub should_fail: bool,
67}
68
69struct LeanThyFunc<'a>(&'a ThyFunc);
70
71impl LeanFmt for SortDecl {
72    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
73        self.name.lean_fmt(f, cx)?;
74        write!(
75            f,
76            " {} : Type",
77            (0..(self.vars))
78                .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
79                .format(" ")
80        )
81    }
82}
83
84impl LeanFmt for InterpretedConst {
85    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
86        write!(
87            f,
88            "abbrev {} : {} := {}",
89            WithLeanCtxt { item: &self.0.name, cx },
90            WithLeanCtxt { item: &self.0.sort, cx },
91            WithLeanCtxt { item: &self.1, cx }
92        )
93    }
94}
95
96// TODO(lean-localize-imports) this seems wrong, but related to lack of storing `VariantIdx` in the `DataProj`
97impl LeanFmt for DataField {
98    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
99        write!(
100            f,
101            "({} : {})",
102            sanitize_name(&self.name.display().to_string()),
103            WithLeanCtxt { item: &self.sort, cx }
104        )
105    }
106}
107
108impl LeanFmt for DataSort {
109    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> std::fmt::Result {
110        match self {
111            DataSort::User(opaque_id) => {
112                let (def_id, _) = cx.opaque_adt_map.get(opaque_id.as_usize()).unwrap();
113                write!(f, "{}", snake_case_to_pascal_case(def_id.name().as_str()))
114            }
115            DataSort::Tuple(n) => write!(f, "Tupleₓ{}", dbg::as_subscript(n)),
116            DataSort::Adt(adt_id) => {
117                let def_id = cx.adt_map.get_index(adt_id.as_usize()).unwrap();
118                write!(f, "{}", def_id_to_pascal_case(def_id, &cx.genv.tcx()))
119            }
120        }
121    }
122}
123
124impl LeanFmt for DataDecl {
125    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
126        if self.ctors.len() == 1 {
127            writeln!(f, "@[ext]")?;
128            write!(f, "structure ")?;
129            self.name.lean_fmt(f, cx)?;
130            writeln!(
131                f,
132                " {} where",
133                (0..self.vars)
134                    .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
135                    .format(" ")
136            )?;
137            let ctor = &self.ctors[0];
138            if let fixpoint::Var::TupleCtor { .. } = &ctor.name {
139                writeln!(f, "  mk ::")?;
140                for (idx, field) in ctor.fields.iter().enumerate() {
141                    writeln!(
142                        f,
143                        "    fld{} : {} ",
144                        as_subscript(idx),
145                        WithLeanCtxt { item: &field.sort, cx }
146                    )?;
147                }
148            } else if let fixpoint::Var::DataCtor(adt_id, _) = &ctor.name {
149                writeln!(
150                    f,
151                    "  mk{}{} ::",
152                    WithLeanCtxt { item: &DataSort::Adt(*adt_id), cx },
153                    as_subscript(0)
154                )?;
155                for (idx, field) in ctor.fields.iter().enumerate() {
156                    writeln!(
157                        f,
158                        "    {} : {} ",
159                        WithLeanCtxt { item: LeanField(*adt_id, idx.try_into().unwrap()), cx },
160                        WithLeanCtxt { item: &field.sort, cx }
161                    )?;
162                }
163            } else {
164                bug!("unexpected ctor {ctor:?} in datadecl");
165            };
166        } else {
167            write!(f, "inductive ")?;
168            self.name.lean_fmt(f, cx)?;
169            writeln!(
170                f,
171                " {} where",
172                (0..self.vars)
173                    .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
174                    .format(" ")
175            )?;
176            for data_ctor in &self.ctors {
177                let fixpoint::Var::DataCtor(adt_id, variant_id) = &data_ctor.name else {
178                    bug!("unexpected ctor {data_ctor:?} in datadecl")
179                };
180                write!(f, "| ")?;
181                write!(
182                    f,
183                    " mk{}{} ",
184                    WithLeanCtxt { item: &DataSort::Adt(*adt_id), cx },
185                    as_subscript(variant_id.as_usize()),
186                )?;
187                // data_ctor.name.lean_fmt(f, cx)?;
188                for field in &data_ctor.fields {
189                    write!(f, " ")?;
190                    field.lean_fmt(f, cx)?;
191                }
192                writeln!(f)?;
193            }
194        }
195        Ok(())
196    }
197}
198
199impl<'a> fmt::Display for LeanThyFunc<'a> {
200    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
201        match self.0 {
202            ThyFunc::IntToBv8 => write!(f, "BitVec.ofInt 8"),
203            ThyFunc::IntToBv32 => write!(f, "BitVec.ofInt 32"),
204            ThyFunc::IntToBv64 => write!(f, "BitVec.ofInt 64"),
205            ThyFunc::IntToBv128 => write!(f, "BitVec.ofInt 128"),
206            ThyFunc::Bv8ToInt | ThyFunc::Bv32ToInt | ThyFunc::Bv64ToInt | ThyFunc::Bv128ToInt => {
207                write!(f, "BitVec.toNat")
208            }
209            ThyFunc::BvAdd => write!(f, "BitVec.add"),
210            ThyFunc::BvSub => write!(f, "BitVec.sub"),
211            ThyFunc::BvMul => write!(f, "BitVec.mul"),
212            ThyFunc::BvNeg => write!(f, "BitVec.neg"),
213            ThyFunc::BvSdiv => write!(f, "BitVec.sdiv"),
214            ThyFunc::BvSrem => write!(f, "BitVec.srem"),
215            ThyFunc::BvUdiv => write!(f, "BitVec.udiv"),
216            ThyFunc::BvAnd => write!(f, "BitVec.and"),
217            ThyFunc::BvOr => write!(f, "BitVec.or"),
218            ThyFunc::BvXor => write!(f, "BitVec.xor"),
219            ThyFunc::BvNot => write!(f, "BitVec.not"),
220            ThyFunc::BvSle => write!(f, "BitVec.sle"),
221            ThyFunc::BvSlt => write!(f, "BitVec.slt"),
222            ThyFunc::BvUle => write!(f, "BitVec.ule"),
223            ThyFunc::BvUlt => write!(f, "BitVec.ult"),
224            ThyFunc::BvAshr => write!(f, "BitVec_sshiftRight"),
225            ThyFunc::BvLshr => write!(f, "BitVec_ushiftRight"),
226            ThyFunc::BvShl => write!(f, "BitVec_shiftLeft"),
227            ThyFunc::BvSignExtend(size) => write!(f, "BitVec_signExtend {}", size),
228            ThyFunc::BvZeroExtend(size) => write!(f, "BitVec_zeroExtend {}", size),
229            ThyFunc::BvUrem => write!(f, "BitVec.umod"),
230            ThyFunc::BvSge => write!(f, "BitVec_sge"),
231            ThyFunc::BvSgt => write!(f, "BitVec_sgt"),
232            ThyFunc::BvUge => write!(f, "BitVec_uge"),
233            ThyFunc::BvUgt => write!(f, "BitVec_ugt"),
234            ThyFunc::MapDefault => write!(f, "SmtMap_default"),
235            ThyFunc::MapSelect => write!(f, "SmtMap_select"),
236            ThyFunc::MapStore => write!(f, "SmtMap_store"),
237            func => panic!("Unsupported theory function {}", func),
238        }
239    }
240}
241
242impl<T: LeanFmt> LeanFmt for &T {
243    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
244        (*self).lean_fmt(f, cx)
245    }
246}
247
248struct LeanAdt(AdtId);
249struct LeanDataProj(AdtId, u32);
250struct LeanField(AdtId, u32);
251
252impl LeanFmt for LeanField {
253    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
254        let adt_id = self.0;
255        if let Some(def_id) = cx.adt_map.get_index(adt_id.as_usize())
256            && let Ok(adt_sort_def) = cx.genv.adt_sort_def_of(*def_id)
257        {
258            write!(
259                f,
260                "{}",
261                sanitize_name(
262                    adt_sort_def.struct_variant().field_names()[self.1 as usize].as_str()
263                )
264            )
265        } else {
266            write!(f, "fld{}", as_subscript(self.1 as usize))
267        }
268    }
269}
270
271impl LeanFmt for LeanAdt {
272    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
273        let def_id = cx.adt_map.get_index(self.0.as_usize()).unwrap();
274        write!(f, "{}", def_id_to_pascal_case(def_id, &cx.genv.tcx()))
275    }
276}
277
278impl LeanFmt for LeanDataProj {
279    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
280        write!(
281            f,
282            "{}.{}",
283            WithLeanCtxt { item: LeanAdt(self.0), cx },
284            WithLeanCtxt { item: LeanField(self.0, self.1), cx }
285        )
286    }
287}
288
289impl LeanFmt for Var {
290    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
291        match self {
292            Var::Global(_gvar, def_id) => {
293                let path = stable_def_path(cx.genv.tcx(), def_id.parent());
294                if path.is_empty() {
295                    write!(f, "{}", def_id.name())
296                } else {
297                    write!(f, "{}_{}", sanitize_name(&path), def_id.name())
298                }
299            }
300            Var::Const(gvar, None) => {
301                if let Some(stable) = cx.primop_var_map.get(gvar) {
302                    write!(f, "{stable}")
303                } else {
304                    write!(f, "{}", sanitize_name(&self.display().to_string()))
305                }
306            }
307            Var::Const(_, Some(did)) => {
308                let path = stable_def_path(cx.genv.tcx(), *did);
309                write!(f, "{}", sanitize_name(&path))
310            }
311            Var::DataCtor(adt_id, idx) => {
312                write!(
313                    f,
314                    "{}.mk{}{}",
315                    WithLeanCtxt { item: LeanAdt(*adt_id), cx },
316                    WithLeanCtxt { item: &DataSort::Adt(*adt_id), cx },
317                    as_subscript(idx.as_usize())
318                )
319            }
320            Var::DataProj { adt_id, field } => LeanDataProj(*adt_id, *field).lean_fmt(f, cx),
321            Var::TupleCtor { arity } => {
322                write!(f, "{}.mk", WithLeanCtxt { item: &DataSort::Tuple(*arity), cx })
323            }
324            Var::TupleProj { arity, field } => {
325                write!(
326                    f,
327                    "{}.fld{}",
328                    WithLeanCtxt { item: &DataSort::Tuple(*arity), cx },
329                    as_subscript(*field as usize)
330                )
331            }
332            Var::Local(local_var) => {
333                write!(f, "{}", cx.pretty_var_map.get(&PrettyVar::Local(*local_var)))
334            }
335            Var::Param(param) => {
336                write!(f, "{}", cx.pretty_var_map.get(&PrettyVar::Param(*param)))
337            }
338            _ => {
339                write!(f, "{}", sanitize_name(&self.display().to_string()))
340            }
341        }
342    }
343}
344
345impl LeanFmt for Sort {
346    fn lean_fmt(&self, f: &mut std::fmt::Formatter, cx: &LeanCtxt) -> std::fmt::Result {
347        match self {
348            Sort::Int => write!(f, "Int"),
349            Sort::Bool => write!(f, "Prop"),
350            Sort::Real => write!(f, "Real"),
351            Sort::Str => write!(f, "String"),
352            Sort::Func(f_sort) => {
353                write!(
354                    f,
355                    "({} -> {})",
356                    WithLeanCtxt { item: &f_sort[0], cx },
357                    WithLeanCtxt { item: &f_sort[1], cx }
358                )
359            }
360            Sort::App(sort_ctor, args) => {
361                match sort_ctor {
362                    SortCtor::Data(sort) => {
363                        if args.is_empty() {
364                            sort.lean_fmt(f, cx)
365                        } else {
366                            write!(
367                                f,
368                                "({} {})",
369                                WithLeanCtxt { item: sort, cx },
370                                args.iter()
371                                    .map(|arg| { WithLeanCtxt { item: arg, cx } })
372                                    .format(" ")
373                            )
374                        }
375                    }
376                    SortCtor::Map => {
377                        write!(
378                            f,
379                            "(SmtMap {} {})",
380                            WithLeanCtxt { item: &args[0], cx },
381                            WithLeanCtxt { item: &args[1], cx }
382                        )
383                    }
384                    _ => todo!(),
385                }
386            }
387            Sort::BitVec(bv_size) => {
388                match bv_size.as_ref() {
389                    Sort::BvSize(size) => write!(f, "BitVec {}", size),
390                    s => {
391                        panic!(
392                            "encountered sort {} where bitvec size was expected",
393                            WithLeanCtxt { item: s, cx }
394                        )
395                    }
396                }
397            }
398            Sort::Abs(v, sort) => {
399                write!(
400                    f,
401                    "{{t{v} : Type}} -> [Inhabited t{v}] -> {}",
402                    WithLeanCtxt { item: sort.as_ref(), cx }
403                )
404            }
405            Sort::Var(v) => {
406                if cx.hide_sort_vars {
407                    write!(f, "_")
408                } else {
409                    write!(f, "t{v}")
410                }
411            }
412            Sort::BvSize(size) => {
413                panic!("sort BvSize({size}) should only occur as an argument to BitVec")
414            }
415        }
416    }
417}
418
419impl LeanFmt for Expr {
420    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
421        match self {
422            Expr::Var(v) => v.lean_fmt(f, cx),
423            Expr::Constant(c) => {
424                match c {
425                    Constant::Numeral(n) => write!(f, "{n}",),
426                    Constant::Boolean(b) => write!(f, "{}", if *b { "True" } else { "False" }),
427                    Constant::String(s) => write!(f, "{}", s.display()),
428                    Constant::Real(n) => write!(f, "{}", n.display()),
429                    Constant::BitVec(bv, size) => write!(f, "{}#{}", bv, size),
430                }
431            }
432            Expr::BinaryOp(bin_op, args) => {
433                let bin_op_str = match bin_op {
434                    BinOp::Add => "+",
435                    BinOp::Sub => "-",
436                    BinOp::Mul => "*",
437                    BinOp::Div => "/",
438                    BinOp::Mod => "%",
439                };
440                write!(f, "(")?;
441                args[0].lean_fmt(f, cx)?;
442                write!(f, " {} ", bin_op_str)?;
443                args[1].lean_fmt(f, cx)?;
444                write!(f, ")")
445            }
446            Expr::Atom(bin_rel, args) => {
447                let bin_rel_str = match bin_rel {
448                    BinRel::Eq => "=",
449                    BinRel::Ne => "≠",
450                    BinRel::Le => "≤",
451                    BinRel::Lt => "<",
452                    BinRel::Ge => "≥",
453                    BinRel::Gt => ">",
454                };
455                write!(f, "(")?;
456                args[0].lean_fmt(f, cx)?;
457                write!(f, " {} ", bin_rel_str)?;
458                args[1].lean_fmt(f, cx)?;
459                write!(f, ")")
460            }
461            Expr::App(function, sort_args, args, out_sort) => {
462                if out_sort.is_some() {
463                    write!(f, "(")?;
464                }
465                write!(f, "(")?;
466                function.as_ref().lean_fmt(f, cx)?;
467                if let Some(sort_args) = sort_args {
468                    let cx = &LeanCtxt { hide_sort_vars: true, ..*cx };
469                    for (i, s_arg) in sort_args.iter().enumerate() {
470                        if matches!(s_arg, Sort::BvSize(..)) {
471                            continue;
472                        }
473                        write!(f, " (t{i} := {})", WithLeanCtxt { item: s_arg, cx })?;
474                    }
475                }
476                for arg in args {
477                    write!(f, " ")?;
478                    arg.lean_fmt(f, cx)?;
479                }
480                write!(f, ")")?;
481                if let Some(out_sort) = out_sort {
482                    write!(f, " : (")?;
483                    out_sort.lean_fmt(f, cx)?;
484                    write!(f, "))")?;
485                }
486                Ok(())
487            }
488            Expr::And(exprs) => {
489                write!(f, "(")?;
490                for (i, expr) in exprs.iter().enumerate() {
491                    if i > 0 {
492                        write!(f, " ∧ ")?;
493                    }
494                    expr.lean_fmt(f, cx)?;
495                }
496                write!(f, ")")
497            }
498            Expr::Or(exprs) => {
499                write!(f, "(")?;
500                for (i, expr) in exprs.iter().enumerate() {
501                    if i > 0 {
502                        write!(f, " ∨ ")?;
503                    }
504                    expr.lean_fmt(f, cx)?;
505                }
506                write!(f, ")")
507            }
508            Expr::Neg(inner) => {
509                write!(f, "(-")?;
510                inner.as_ref().lean_fmt(f, cx)?;
511                write!(f, ")")
512            }
513            Expr::IfThenElse(ite) => {
514                let [condition, if_true, if_false] = ite.as_ref();
515                write!(f, "(if ")?;
516                condition.lean_fmt(f, cx)?;
517                write!(f, " then ")?;
518                if_true.lean_fmt(f, cx)?;
519                write!(f, " else ")?;
520                if_false.lean_fmt(f, cx)?;
521                write!(f, ")")
522            }
523            Expr::Not(inner) => {
524                write!(f, "(")?;
525                write!(f, "¬")?;
526                inner.as_ref().lean_fmt(f, cx)?;
527                write!(f, ")")
528            }
529            Expr::Imp(implication) => {
530                let [lhs, rhs] = implication.as_ref();
531                write!(f, "(")?;
532                lhs.lean_fmt(f, cx)?;
533                write!(f, " -> ")?;
534                rhs.lean_fmt(f, cx)?;
535                write!(f, ")")
536            }
537            Expr::Iff(equiv) => {
538                let [lhs, rhs] = equiv.as_ref();
539                write!(f, "(")?;
540                lhs.lean_fmt(f, cx)?;
541                write!(f, " <-> ")?;
542                rhs.lean_fmt(f, cx)?;
543                write!(f, ")")
544            }
545            Expr::Let(binder, exprs) => {
546                let [def, body] = exprs.as_ref();
547                write!(f, "(let ")?;
548                binder.lean_fmt(f, cx)?;
549                write!(f, " := ")?;
550                def.lean_fmt(f, cx)?;
551                write!(f, "; ")?;
552                body.lean_fmt(f, cx)?;
553                write!(f, ")")
554            }
555            Expr::ThyFunc(thy_func) => {
556                write!(f, "{}", LeanThyFunc(thy_func))
557            }
558            Expr::IsCtor(..) => {
559                todo!("not yet implemented: datatypes in lean")
560            }
561            Expr::Quantifier(q, bind, expr) => {
562                match q {
563                    Quantifier::Exists => write!(f, "(∃ ")?,
564                    Quantifier::Forall => write!(f, "(∀ ")?,
565                };
566                for (var, sort) in bind {
567                    write!(f, "(")?;
568                    var.lean_fmt(f, cx)?;
569                    write!(f, " : {})", WithLeanCtxt { item: sort, cx })?;
570                }
571                write!(f, ", ")?;
572                expr.lean_fmt(f, cx)?;
573                write!(f, ")")?;
574                Ok(())
575            }
576            Expr::WKVar(_) => {
577                todo!("not yet implemented")
578            }
579        }
580    }
581}
582
583impl LeanFmt for FunDef {
584    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
585        let FunDef { name, sort, comment: _, body } = self;
586        write!(f, "noncomputable def ")?;
587        name.lean_fmt(f, cx)?;
588        if let Some(body) = body {
589            for i in 0..sort.params {
590                write!(f, " {{t{i} : Type}} [Inhabited t{i}]")?;
591            }
592            for (arg, arg_sort) in iter::zip(&body.args, &sort.inputs) {
593                write!(f, " (")?;
594                arg.lean_fmt(f, cx)?;
595                write!(f, " : {})", WithLeanCtxt { item: arg_sort, cx })?;
596            }
597            writeln!(f, " : {} :=", WithLeanCtxt { item: &sort.output, cx })?;
598            write!(f, "  ")?;
599            body.expr.lean_fmt(f, cx)?;
600        } else {
601            write!(f, " : {} := sorry", WithLeanCtxt { item: sort, cx })?;
602        }
603        writeln!(f)
604    }
605}
606
607impl LeanFmt for FunSort {
608    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
609        for i in 0..self.params {
610            write!(f, "{{t{i} : Type}} -> [Inhabited t{i}] -> ")?;
611        }
612        if self.inputs.is_empty() {
613            write!(f, "{}", WithLeanCtxt { item: &self.output, cx })
614        } else {
615            write!(
616                f,
617                "{} -> {}",
618                self.inputs.iter().format_with(" -> ", |sort, f| {
619                    f(&format_args!("{}", WithLeanCtxt { item: sort, cx }))
620                }),
621                WithLeanCtxt { item: &self.output, cx }
622            )
623        }
624    }
625}
626
627impl LeanFmt for [Pred] {
628    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
629        if self.is_empty() {
630            return write!(f, "True");
631        }
632
633        write!(f, "(")?;
634        for (i, pred) in self.iter().enumerate() {
635            if i > 0 {
636                write!(f, " ∧ ")?;
637            }
638            pred.lean_fmt(f, cx)?;
639        }
640        write!(f, ")")?;
641        Ok(())
642    }
643}
644
645impl LeanFmt for Pred {
646    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
647        match self {
648            Pred::KVar(kvid, args) => {
649                write!(f, "({}", sanitize_name(&kvid.display().to_string()))?;
650                for arg in args {
651                    write!(f, " ")?;
652                    arg.lean_fmt(f, cx)?;
653                }
654                write!(f, ")")
655            }
656            Pred::Expr(expr) => expr.lean_fmt(f, cx),
657        }
658    }
659}
660
661impl LeanFmt for KVarDecl {
662    fn lean_fmt(&self, f: &mut std::fmt::Formatter, cx: &LeanCtxt) -> std::fmt::Result {
663        write!(f, "∃ {} : ", sanitize_name(&self.kvid.display().to_string()))?;
664        for (i, sort) in self.sorts.iter().enumerate() {
665            write!(f, "(a{i} : {}) -> ", WithLeanCtxt { item: sort, cx })?;
666        }
667        write!(f, "Prop")
668    }
669}
670
671impl LeanFmt for (&KVid, &ClosedSolution) {
672    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
673        let (kvid, (implicit, (explicit, inner))) = self;
674        write!(f, "def k{} ", kvid.as_usize())?;
675        for (arg, sort) in implicit.iter().chain(explicit) {
676            write!(f, "(")?;
677            arg.lean_fmt(f, cx)?;
678            write!(f, " : {}) ", WithLeanCtxt { item: sort, cx })?;
679        }
680        writeln!(f, ": Prop :=")?;
681        write!(f, "  ")?;
682        inner.lean_fmt(f, cx)?;
683        Ok(())
684    }
685}
686
687impl<'a> LeanFmt for LeanKConstraint<'a> {
688    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
689        let theorem_name = self.theorem_name.replace(".", "_");
690
691        write!(f, "\n\ndef {theorem_name} := ")?;
692
693        if self.should_fail {
694            write!(f, "¬")?;
695        }
696
697        for bind in self.const_binds {
698            write!(f, "∀ (")?;
699            bind.name.lean_fmt(f, cx)?;
700            write!(f, " : {}), ", WithLeanCtxt { item: &bind.sort, cx })?;
701        }
702
703        if self.kvars.is_empty() {
704            self.constr.lean_fmt(f, cx)
705        } else {
706            write!(
707                f,
708                "{}, ",
709                self.kvars
710                    .iter()
711                    .map(|kvar| { WithLeanCtxt { item: kvar, cx } })
712                    .format(", ")
713            )?;
714            self.constr.lean_fmt(f, cx)
715        }
716    }
717}
718
719impl LeanFmt for Constraint {
720    fn lean_fmt(&self, f: &mut fmt::Formatter, cx: &LeanCtxt) -> fmt::Result {
721        let mut fmt_cx = ConstraintFormatter::default();
722        fmt_cx.incr();
723        fmt_cx.newline(f)?;
724        self.fmt_nested(f, cx, &mut fmt_cx)?;
725        fmt_cx.decr();
726        Ok(())
727    }
728}
729
730impl FormatNested for Constraint {
731    fn fmt_nested(
732        &self,
733        f: &mut fmt::Formatter,
734        lean_cx: &LeanCtxt,
735        fmt_cx: &mut ConstraintFormatter,
736    ) -> fmt::Result {
737        match self {
738            Constraint::ForAll(bind, inner) => {
739                let trivial_pred =
740                    bind.preds.iter().all(Pred::is_trivially_true) || bind.preds.is_empty();
741                let trivial_bind = bind.name.display().to_string().starts_with("_");
742                if !trivial_bind {
743                    write!(f, "∀ (")?;
744                    bind.name.lean_fmt(f, lean_cx)?;
745                    write!(f, " : {}),", WithLeanCtxt { item: &bind.sort, cx: lean_cx })?;
746                    fmt_cx.incr();
747                    fmt_cx.newline(f)?;
748                }
749                if !trivial_pred {
750                    bind.preds.lean_fmt(f, lean_cx)?;
751                    write!(f, " ->")?;
752                    fmt_cx.incr();
753                    fmt_cx.newline(f)?;
754                }
755                inner.fmt_nested(f, lean_cx, fmt_cx)?;
756                if !trivial_pred {
757                    fmt_cx.decr();
758                }
759                if !trivial_bind {
760                    fmt_cx.decr();
761                }
762                Ok(())
763            }
764            Constraint::Conj(constraints) => {
765                let n = constraints.len();
766                for (i, constraint) in constraints.iter().enumerate() {
767                    write!(f, "(")?;
768                    constraint.fmt_nested(f, lean_cx, fmt_cx)?;
769                    write!(f, ")")?;
770                    if i < n - 1 {
771                        write!(f, " ∧")?;
772                    }
773                    fmt_cx.newline(f)?;
774                }
775                Ok(())
776            }
777            Constraint::Pred(head, _) => head.lean_fmt(f, lean_cx),
778        }
779    }
780}
781
782pub trait FormatNested {
783    fn fmt_nested(
784        &self,
785        f: &mut fmt::Formatter,
786        lean_cx: &LeanCtxt,
787        fmt_cx: &mut ConstraintFormatter,
788    ) -> fmt::Result;
789}
790
791#[derive(Default)]
792pub struct ConstraintFormatter {
793    level: u32,
794}
795
796impl ConstraintFormatter {
797    pub fn incr(&mut self) {
798        self.level += 1;
799    }
800
801    pub fn decr(&mut self) {
802        self.level -= 1;
803    }
804
805    pub fn newline(&self, f: &mut fmt::Formatter) -> fmt::Result {
806        f.write_char('\n')?;
807        self.padding(f)
808    }
809
810    pub fn padding(&self, f: &mut fmt::Formatter) -> fmt::Result {
811        for _ in 0..self.level {
812            f.write_str(" ")?;
813        }
814        Ok(())
815    }
816}
817
818static IMPL_RE: LazyLock<regex::Regex> =
819    std::sync::LazyLock::new(|| regex::Regex::new(r"\{impl#(\d+)\}").unwrap());
820
821fn is_reserved(name: &str) -> bool {
822    matches!(name, "end" | "def")
823}
824
825fn sanitize_name(name: &str) -> String {
826    if is_reserved(name) {
827        format!("{name}_")
828    } else {
829        IMPL_RE
830            .replace_all(name, "impl_${1}")
831            .replace("-", "_")
832            .replace("$", "_")
833    }
834}
835
836pub fn def_id_to_pascal_case(def_id: &DefId, tcx: &TyCtxt) -> String {
837    let snake = stable_def_path(*tcx, *def_id).replace("-", "_");
838    snake_case_to_pascal_case(&snake)
839}
840
841/// Like [`rustc_hir::definitions::DefPath::to_filename_friendly_no_crate`], but names impl
842/// blocks by their header (e.g., `RawValIter_T_as_Iterator`) instead of by their position in
843/// the enclosing module (`{impl#10}`). Lean names, and hence the names of user proof files,
844/// then stay the same when impl blocks are added, removed or reordered.
845///
846/// Components that would otherwise not be valid in a Lean identifier are rendered as follows:
847/// * A `const _` (e.g., the one `defs!` expands to, which declares flux items) is dropped if it
848///   is the first `const _` in its parent, and rendered as `_n` for the `n`-th one after that,
849///   so items declared by the first (usually the only) `defs!` keep the name of the enclosing
850///   item.
851/// * An anonymous component (e.g., `{closure#0}`) is rendered as `closure_0`.
852fn stable_def_path(tcx: TyCtxt, def_id: DefId) -> String {
853    // `DefPath::data` has one component per ancestor, excluding the crate root.
854    let mut ancestors = iter::successors(Some(def_id), |id| tcx.opt_parent(*id)).collect_vec();
855    ancestors.pop();
856    ancestors.reverse();
857    iter::zip(ancestors, tcx.def_path(def_id).data)
858        .filter_map(|(id, component)| {
859            match (component.data, component.data.name()) {
860                (DefPathData::Impl, _) => Some(stable_impl_name(tcx, id)),
861                (DefPathData::ValueNs(kw::Underscore), _) => {
862                    (component.disambiguator != 0).then(|| format!("_{}", component.disambiguator))
863                }
864                (_, DefPathDataName::Anon { namespace }) => {
865                    Some(format!("{namespace}_{}", component.disambiguator))
866                }
867                (_, DefPathDataName::Named(_)) => Some(component.as_sym(true).to_string()),
868            }
869        })
870        .join("-")
871}
872
873/// A name for an impl block derived from its header: `SelfTy_as_Trait` for trait impls and
874/// `SelfTy` for inherent impls. In the rare case where several impls in the same parent get
875/// the same name, e.g., because their generic parameters are named differently, they are
876/// suffixed with their index among those impls.
877fn stable_impl_name(tcx: TyCtxt, impl_id: DefId) -> String {
878    let name = impl_header_name(tcx, impl_id);
879    let parent = tcx.opt_parent(impl_id);
880    let clashing = tcx
881        .hir_crate_items(())
882        .definitions()
883        .map(|id| id.to_def_id())
884        .filter(|id| {
885            matches!(tcx.def_kind(*id), DefKind::Impl { .. })
886                && tcx.opt_parent(*id) == parent
887                && impl_header_name(tcx, *id) == name
888        })
889        .sorted_by_key(|id| tcx.def_key(*id).disambiguated_data.disambiguator)
890        .collect_vec();
891    if clashing.len() > 1
892        && let Some(idx) = clashing.iter().position(|id| *id == impl_id)
893    {
894        format!("{name}_{idx}")
895    } else {
896        name
897    }
898}
899
900fn impl_header_name(tcx: TyCtxt, impl_id: DefId) -> String {
901    let self_ty = tcx
902        .type_of(impl_id)
903        .instantiate_identity()
904        .skip_normalization();
905    let header = with_forced_trimmed_paths!(match tcx.def_kind(impl_id) {
906        DefKind::Impl { of_trait: true } => {
907            let trait_ref = tcx
908                .impl_trait_ref(impl_id)
909                .instantiate_identity()
910                .skip_normalization();
911            format!("{self_ty} as {}", trait_ref.print_only_trait_path())
912        }
913        _ => format!("{self_ty}"),
914    });
915    // Keep only identifier characters, e.g., `&'a [T] as Foo<u8>` becomes `a_T_as_Foo_u8`.
916    header
917        .split(|c: char| !c.is_alphanumeric())
918        .filter(|word| !word.is_empty())
919        .join("_")
920}
921
922pub fn snake_case_to_pascal_case(snake: &str) -> String {
923    snake
924        .split('_')
925        .filter(|s| !s.is_empty()) // skip empty segments (handles double underscores)
926        .map(|word| {
927            let mut chars = word.chars();
928            match chars.next() {
929                Some(first) => first.to_ascii_uppercase().to_string() + chars.as_str(),
930                None => String::new(),
931            }
932        })
933        .collect::<String>()
934}
flux_infer/lean_format.rs

flux_infer/
lean_format.rs
