flux_infer/
lean_format.rs

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

flux_infer/
lean_format.rs
