flux_infer/
lean_format.rs

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

flux_infer/
lean_format.rs
