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

flux_infer/
lean_format.rs
