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

flux_infer/
lean_format.rs
