flux_infer/
lean_format.rs

1use core::fmt;
2
3use flux_middle::global_env::GlobalEnv;
4use itertools::Itertools;
5use liquid_fixpoint::{FixpointFmt, Identifier, ThyFunc};
6
7use crate::fixpoint_encoding::fixpoint::{
8    BinOp, BinRel, ConstDecl, Constant, Constraint, DataDecl, DataField, DataSort, Expr, FunDef,
9    KVarDecl, Pred, Sort, SortCtor, SortDecl, Var,
10};
11
12struct LeanSort<'a>(&'a Sort);
13struct LeanKVarDecl<'a>(&'a KVarDecl);
14pub struct LeanKConstraint<'a, 'genv, 'tcx>(
15    pub &'a [KVarDecl],
16    pub &'a Constraint,
17    pub GlobalEnv<'genv, 'tcx>,
18);
19pub struct LeanFunDef<'a, 'genv, 'tcx>(pub &'a FunDef, pub GlobalEnv<'genv, 'tcx>);
20pub struct LeanSortDecl<'a, 'genv, 'tcx>(pub &'a SortDecl, pub GlobalEnv<'genv, 'tcx>);
21pub struct LeanDataDecl<'a, 'genv, 'tcx>(pub &'a DataDecl, pub GlobalEnv<'genv, 'tcx>);
22pub struct LeanConstDecl<'a, 'genv, 'tcx>(pub &'a ConstDecl, pub GlobalEnv<'genv, 'tcx>);
23pub struct LeanSortVar<'a>(pub &'a DataSort);
24struct LeanDataField<'a>(&'a DataField);
25struct LeanConstraint<'a, 'genv, 'tcx>(&'a Constraint, GlobalEnv<'genv, 'tcx>);
26struct LeanPred<'a, 'genv, 'tcx>(&'a Pred, GlobalEnv<'genv, 'tcx>);
27struct LeanExpr<'a, 'genv, 'tcx>(&'a Expr, GlobalEnv<'genv, 'tcx>);
28pub struct LeanVar<'a, 'genv, 'tcx>(pub &'a Var, pub GlobalEnv<'genv, 'tcx>);
29struct LeanThyFunc<'a>(&'a ThyFunc);
30
31impl<'a, 'genv, 'tcx> fmt::Display for LeanSortDecl<'a, 'genv, 'tcx> {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(
34            f,
35            "{} {} : Type",
36            LeanSortVar(&self.0.name),
37            (0..(self.0.vars))
38                .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
39                .format(" ")
40        )
41    }
42}
43
44impl<'a, 'genv, 'tcx> fmt::Display for LeanConstDecl<'a, 'genv, 'tcx> {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        write!(f, "{} : {}", LeanVar(&self.0.name, self.1), LeanSort(&self.0.sort))
47    }
48}
49
50impl<'a> fmt::Display for LeanDataField<'a> {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        write!(
53            f,
54            "({} : {})",
55            self.0.name.display().to_string().replace("$", "_"),
56            LeanSort(&self.0.sort)
57        )
58    }
59}
60
61impl<'a> fmt::Display for LeanSortVar<'a> {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        match self.0 {
64            DataSort::User(def_id) => write!(f, "{}", def_id.name()),
65            _ => write!(f, "{}", self.0.display().to_string().replace("$", "_")),
66        }
67    }
68}
69
70impl<'a, 'genv, 'tcx> fmt::Display for LeanDataDecl<'a, 'genv, 'tcx> {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        if self.0.ctors.len() == 1 {
73            writeln!(f, "@[ext]")?;
74            writeln!(
75                f,
76                "structure {} {} where",
77                LeanSortVar(&self.0.name),
78                (0..self.0.vars)
79                    .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
80                    .format(" ")
81            )?;
82            writeln!(f, "  {}::", self.0.ctors[0].name.display().to_string().replace("$", "_"),)?;
83            for field in &self.0.ctors[0].fields {
84                writeln!(f, "  {}", LeanDataField(field))?;
85            }
86        } else {
87            writeln!(
88                f,
89                "inductive {} {} where",
90                LeanSortVar(&self.0.name),
91                (0..self.0.vars)
92                    .map(|i| format!("(t{i} : Type) [Inhabited t{i}]"))
93                    .format(" ")
94            )?;
95            for data_ctor in &self.0.ctors {
96                writeln!(
97                    f,
98                    "| {} {}",
99                    LeanVar(&data_ctor.name, self.1),
100                    data_ctor.fields.iter().map(LeanDataField).format(" ")
101                )?;
102            }
103        }
104        Ok(())
105    }
106}
107
108impl<'a> fmt::Display for LeanThyFunc<'a> {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self.0 {
111            ThyFunc::IntToBv8 => write!(f, "BitVec.ofInt 8"),
112            ThyFunc::IntToBv32 => write!(f, "BitVec.ofInt 32"),
113            ThyFunc::IntToBv64 => write!(f, "BitVec.ofInt 64"),
114            ThyFunc::Bv8ToInt | ThyFunc::Bv32ToInt | ThyFunc::Bv64ToInt => {
115                write!(f, "BitVec.toNat")
116            }
117            ThyFunc::BvAdd => write!(f, "BitVec.add"),
118            ThyFunc::BvSub => write!(f, "BitVec.sub"),
119            ThyFunc::BvMul => write!(f, "BitVec.mul"),
120            ThyFunc::BvNeg => write!(f, "BitVec.neg"),
121            ThyFunc::BvSdiv => write!(f, "BitVec.sdiv"),
122            ThyFunc::BvSrem => write!(f, "BitVec.srem"),
123            ThyFunc::BvUdiv => write!(f, "BitVec.udiv"),
124            ThyFunc::BvAnd => write!(f, "BitVec.and"),
125            ThyFunc::BvOr => write!(f, "BitVec.or"),
126            ThyFunc::BvXor => write!(f, "BitVec.xor"),
127            ThyFunc::BvNot => write!(f, "BitVec.not"),
128            ThyFunc::BvSle => write!(f, "BitVec.sle"),
129            ThyFunc::BvSlt => write!(f, "BitVec.slt"),
130            ThyFunc::BvUle => write!(f, "BitVec.ule"),
131            ThyFunc::BvUlt => write!(f, "BitVec.ult"),
132            ThyFunc::BvAshr => write!(f, "BitVec_sshiftRight"),
133            ThyFunc::BvLshr => write!(f, "BitVec_ushiftRight"),
134            ThyFunc::BvShl => write!(f, "BitVec_shiftLeft"),
135            ThyFunc::BvSignExtend(size) => write!(f, "BitVec.signExtend {}", size),
136            ThyFunc::BvZeroExtend(size) => write!(f, "BitVec.zeroExtend {}", size),
137            ThyFunc::BvUrem => write!(f, "BitVec.umod"),
138            ThyFunc::BvSge => write!(f, "BitVec_sge"),
139            ThyFunc::BvSgt => write!(f, "BitVec_sgt"),
140            ThyFunc::BvUge => write!(f, "BitVec_uge"),
141            ThyFunc::BvUgt => write!(f, "BitVec_ugt"),
142            ThyFunc::MapDefault => write!(f, "SmtMap_default"),
143            ThyFunc::MapSelect => write!(f, "SmtMap_select"),
144            ThyFunc::MapStore => write!(f, "SmtMap_store"),
145            func => panic!("Unsupported theory function {}", func),
146        }
147    }
148}
149
150impl<'a, 'genv, 'tcx> fmt::Display for LeanVar<'a, 'genv, 'tcx> {
151    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
152        match self.0 {
153            Var::Global(_gvar, Some(def_id)) => {
154                let path = self
155                    .1
156                    .tcx()
157                    .def_path(def_id.parent())
158                    .to_filename_friendly_no_crate()
159                    .replace("-", "_");
160                if path.is_empty() {
161                    write!(f, "{}", def_id.name())
162                } else {
163                    write!(f, "{path}_{}", def_id.name())
164                }
165            }
166            Var::DataCtor(adt_id, _) | Var::DataProj { adt_id, field: _ } => {
167                write!(
168                    f,
169                    "{}.{}",
170                    LeanSortVar(&DataSort::Adt(*adt_id)),
171                    self.0.display().to_string().replace("$", "_")
172                )
173            }
174            _ => {
175                write!(f, "{}", self.0.display().to_string().replace("$", "_"))
176            }
177        }
178    }
179}
180
181impl<'a> fmt::Display for LeanSort<'a> {
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        match self.0 {
184            Sort::Int => write!(f, "Int"),
185            Sort::Bool => write!(f, "Prop"),
186            Sort::Real => write!(f, "Real"),
187            Sort::Str => write!(f, "String"),
188            Sort::Func(f_sort) => {
189                write!(f, "({} -> {})", LeanSort(&f_sort[0]), LeanSort(&f_sort[1]))
190            }
191            Sort::App(sort_ctor, args) => {
192                match sort_ctor {
193                    SortCtor::Data(sort) => {
194                        if args.is_empty() {
195                            write!(f, "{}", LeanSortVar(sort))
196                        } else {
197                            write!(
198                                f,
199                                "({} {})",
200                                LeanSortVar(sort),
201                                args.iter().map(LeanSort).format(" ")
202                            )
203                        }
204                    }
205                    SortCtor::Map => {
206                        write!(f, "(SmtMap {} {})", LeanSort(&args[0]), LeanSort(&args[1]))
207                    }
208                    _ => todo!(),
209                }
210            }
211            Sort::BitVec(bv_size) => {
212                match bv_size.as_ref() {
213                    Sort::BvSize(size) => write!(f, "BitVec {}", size),
214                    s => panic!("encountered sort {} where bitvec size was expected", LeanSort(s)),
215                }
216            }
217            Sort::Abs(v, sort) => {
218                write!(f, "{{t{v} : Type}} -> [Inhabited t{v}] -> {}", LeanSort(sort.as_ref()))
219            }
220            Sort::Var(v) => write!(f, "t{v}"),
221            s => todo!("{:?}", s),
222        }
223    }
224}
225
226impl<'a, 'genv, 'tcx> fmt::Display for LeanExpr<'a, 'genv, 'tcx> {
227    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
228        match self.0 {
229            Expr::Var(v) => write!(f, "{}", LeanVar(v, self.1)),
230            Expr::Constant(c) => {
231                match c {
232                    Constant::Numeral(n) => write!(f, "{n}",),
233                    Constant::Boolean(b) => write!(f, "{}", if *b { "True" } else { "False" }),
234                    Constant::String(s) => write!(f, "{}", s.display()),
235                    Constant::Real(n) => write!(f, "{n}.0"),
236                    Constant::BitVec(bv, size) => write!(f, "{}#{}", bv, size),
237                }
238            }
239            Expr::BinaryOp(bin_op, args) => {
240                let bin_op_str = match bin_op {
241                    BinOp::Add => "+",
242                    BinOp::Sub => "-",
243                    BinOp::Mul => "*",
244                    BinOp::Div => "/",
245                    BinOp::Mod => "%",
246                };
247                write!(
248                    f,
249                    "({} {} {})",
250                    LeanExpr(&args[0], self.1),
251                    bin_op_str,
252                    LeanExpr(&args[1], self.1)
253                )
254            }
255            Expr::Atom(bin_rel, args) => {
256                let bin_rel_str = match bin_rel {
257                    BinRel::Eq => "=",
258                    BinRel::Ne => "≠",
259                    BinRel::Le => "≤",
260                    BinRel::Lt => "<",
261                    BinRel::Ge => "≥",
262                    BinRel::Gt => ">",
263                };
264                write!(
265                    f,
266                    "({} {} {})",
267                    LeanExpr(&args[0], self.1),
268                    bin_rel_str,
269                    LeanExpr(&args[1], self.1)
270                )
271            }
272            Expr::App(function, sort_args, args) => {
273                if let Some(sort_args) = sort_args {
274                    write!(
275                        f,
276                        "({} {} {})",
277                        LeanExpr(function.as_ref(), self.1),
278                        sort_args
279                            .iter()
280                            .enumerate()
281                            .map(|(i, s_arg)| format!("(t{i} := {})", LeanSort(s_arg)))
282                            .format(" "),
283                        args.iter().map(|arg| LeanExpr(arg, self.1)).format(" ")
284                    )
285                } else {
286                    write!(
287                        f,
288                        "({} {})",
289                        LeanExpr(function.as_ref(), self.1),
290                        args.iter().map(|arg| LeanExpr(arg, self.1)).format(" ")
291                    )
292                }
293            }
294            Expr::And(exprs) => {
295                write!(
296                    f,
297                    "({})",
298                    exprs
299                        .iter()
300                        .map(|expr| LeanExpr(expr, self.1))
301                        .format(" && ")
302                )
303            }
304            Expr::Or(exprs) => {
305                write!(
306                    f,
307                    "({})",
308                    exprs
309                        .iter()
310                        .map(|expr| LeanExpr(expr, self.1))
311                        .format(" || ")
312                )
313            }
314            Expr::Neg(inner) => {
315                write!(f, "(-{})", LeanExpr(inner.as_ref(), self.1))
316            }
317            Expr::IfThenElse(ite) => {
318                let [condition, if_true, if_false] = ite.as_ref();
319                write!(
320                    f,
321                    "(if {} then {} else {})",
322                    LeanExpr(condition, self.1),
323                    LeanExpr(if_true, self.1),
324                    LeanExpr(if_false, self.1)
325                )
326            }
327            Expr::Not(inner) => {
328                write!(f, "(¬{})", LeanExpr(inner.as_ref(), self.1))
329            }
330            Expr::Imp(implication) => {
331                let [lhs, rhs] = implication.as_ref();
332                write!(f, "({} -> {})", LeanExpr(lhs, self.1), LeanExpr(rhs, self.1))
333            }
334            Expr::Iff(equiv) => {
335                let [lhs, rhs] = equiv.as_ref();
336                write!(f, "({} <-> {})", LeanExpr(lhs, self.1), LeanExpr(rhs, self.1))
337            }
338            Expr::Let(binder, exprs) => {
339                let [def, body] = exprs.as_ref();
340                write!(
341                    f,
342                    "(let {} := {}; {})",
343                    LeanVar(binder, self.1),
344                    LeanExpr(def, self.1),
345                    LeanExpr(body, self.1)
346                )
347            }
348            Expr::ThyFunc(thy_func) => {
349                write!(f, "{}", LeanThyFunc(thy_func))
350            }
351            Expr::IsCtor(..) => {
352                todo!("not yet implemented: datatypes in lean")
353            }
354            Expr::Exists(..) => {
355                todo!("not yet implemented: exists in lean")
356            }
357            Expr::BoundVar(_) => {
358                unreachable!("bound vars should only be present in fixpoint output")
359            }
360        }
361    }
362}
363
364impl<'a, 'genv, 'tcx> fmt::Display for LeanFunDef<'a, 'genv, 'tcx> {
365    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
366        let LeanFunDef(FunDef { name, args, out, comment: _, body }, _) = self;
367        writeln!(
368            f,
369            "def {} {} : {} :=",
370            LeanVar(name, self.1),
371            args.iter()
372                .map(|(arg, arg_sort)| {
373                    format!("({} : {})", LeanVar(arg, self.1), LeanSort(arg_sort))
374                })
375                .collect::<Vec<_>>()
376                .join(" "),
377            LeanSort(out)
378        )?;
379        writeln!(f, "  {}", LeanExpr(body, self.1))
380    }
381}
382
383impl<'a, 'genv, 'tcx> fmt::Display for LeanPred<'a, 'genv, 'tcx> {
384    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
385        match self.0 {
386            Pred::Expr(expr) => write!(f, "{}", LeanExpr(expr, self.1)),
387            Pred::And(preds) => {
388                write!(
389                    f,
390                    "({})",
391                    preds
392                        .iter()
393                        .map(|pred| LeanPred(pred, self.1))
394                        .format(" ∧ ")
395                )
396            }
397            Pred::KVar(kvid, args) => {
398                write!(
399                    f,
400                    "({} {})",
401                    kvid.display().to_string().replace("$", "_"),
402                    args.iter().map(|var| LeanVar(var, self.1)).format(" ")
403                )
404            }
405        }
406    }
407}
408
409impl<'a> fmt::Display for LeanKVarDecl<'a> {
410    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
411        let sorts = self
412            .0
413            .sorts
414            .iter()
415            .enumerate()
416            .map(|(i, sort)| format!("(a{i} : {})", LeanSort(sort)))
417            .format(" -> ");
418        write!(f, "∃ {} : {} -> Prop", self.0.kvid.display().to_string().replace("$", "_"), sorts)
419    }
420}
421
422impl<'a, 'genv, 'tcx> fmt::Display for LeanKConstraint<'a, 'genv, 'tcx> {
423    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
424        if self.0.is_empty() {
425            write!(f, "{}", LeanConstraint(self.1, self.2))
426        } else {
427            write!(
428                f,
429                "{}, {}",
430                self.0.iter().map(LeanKVarDecl).format(", "),
431                LeanConstraint(self.1, self.2)
432            )
433        }
434    }
435}
436
437impl<'a, 'genv, 'tcx> fmt::Display for LeanConstraint<'a, 'genv, 'tcx> {
438    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
439        match self.0 {
440            Constraint::ForAll(bind, inner) => {
441                if bind.pred.is_trivially_true() {
442                    write!(
443                        f,
444                        "(∀ ({} : {}), {})",
445                        LeanVar(&bind.name, self.1),
446                        LeanSort(&bind.sort),
447                        LeanConstraint(inner, self.1)
448                    )
449                } else {
450                    write!(
451                        f,
452                        "(∀ ({} : {}), ({} -> {}))",
453                        LeanVar(&bind.name, self.1),
454                        LeanSort(&bind.sort),
455                        LeanPred(&bind.pred, self.1),
456                        LeanConstraint(inner, self.1)
457                    )
458                }
459            }
460            Constraint::Conj(constraints) => {
461                write!(
462                    f,
463                    "({})",
464                    constraints
465                        .iter()
466                        .map(|constraint| LeanConstraint(constraint, self.1))
467                        .format(" ∧ ")
468                )
469            }
470            Constraint::Pred(pred, _) => {
471                write!(f, "{}", LeanPred(pred, self.1))
472            }
473        }
474    }
475}