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

flux_infer/
lean_format.rs
