flux_infer/
lean_encoding.rs

1use std::{
2    fs::{self, OpenOptions},
3    io::{self, Write},
4    path::{Path, PathBuf},
5    process::{Command, Stdio},
6};
7
8use flux_common::{
9    bug,
10    dbg::{self, SpanTrace},
11    result::ResultExt,
12};
13use flux_config as config;
14use flux_middle::{
15    def_id::{FluxDefId, MaybeExternId},
16    global_env::GlobalEnv,
17    queries::QueryErr,
18    rty::{BinOp, BvSize, PrettyMap, Sort, local_deps},
19};
20use itertools::Itertools;
21use rustc_data_structures::fx::FxIndexSet;
22use rustc_hash::FxHashMap;
23use rustc_hir::def_id::DefId;
24use rustc_span::ErrorGuaranteed;
25
26use crate::{
27    fixpoint_encoding::{ConstDeps, InterpretedConst, KVarSolutions, SortDeps, fixpoint},
28    lean_format::{self, LeanCtxt, WithLeanCtxt, def_id_to_pascal_case, snake_case_to_pascal_case},
29};
30
31fn sort_name_fragment(sort: &Sort) -> String {
32    match sort {
33        Sort::Int => "Int".to_string(),
34        Sort::Bool => "Bool".to_string(),
35        Sort::Real => "Real".to_string(),
36        Sort::BitVec(BvSize::Fixed(n)) => format!("Bv{n}"),
37        Sort::BitVec(BvSize::Param(p)) => format!("BvParam{}", p.as_u32()),
38        Sort::BitVec(BvSize::Infer(_)) => "BvInfer".to_string(),
39        _ => {
40            format!("{sort:?}")
41                .chars()
42                .filter(|c| c.is_alphanumeric())
43                .collect()
44        }
45    }
46}
47
48fn prim_op_lean_name(op: &BinOp) -> String {
49    match op {
50        BinOp::Add(s) => format!("PrimOpAdd{}", sort_name_fragment(s)),
51        BinOp::Sub(s) => format!("PrimOpSub{}", sort_name_fragment(s)),
52        BinOp::Mul(s) => format!("PrimOpMul{}", sort_name_fragment(s)),
53        BinOp::Div(s) => format!("PrimOpDiv{}", sort_name_fragment(s)),
54        BinOp::Mod(s) => format!("PrimOpMod{}", sort_name_fragment(s)),
55        BinOp::BitAnd(s) => format!("PrimOpBitAnd{}", sort_name_fragment(s)),
56        BinOp::BitOr(s) => format!("PrimOpBitOr{}", sort_name_fragment(s)),
57        BinOp::BitXor(s) => format!("PrimOpBitXor{}", sort_name_fragment(s)),
58        BinOp::BitShl(s) => format!("PrimOpBitShl{}", sort_name_fragment(s)),
59        BinOp::BitShr(s) => format!("PrimOpBitShr{}", sort_name_fragment(s)),
60        BinOp::Gt(s) => format!("PrimOpGt{}", sort_name_fragment(s)),
61        BinOp::Ge(s) => format!("PrimOpGe{}", sort_name_fragment(s)),
62        BinOp::Lt(s) => format!("PrimOpLt{}", sort_name_fragment(s)),
63        BinOp::Le(s) => format!("PrimOpLe{}", sort_name_fragment(s)),
64        BinOp::Eq => "PrimOpEq".to_string(),
65        BinOp::Ne => "PrimOpNe".to_string(),
66        BinOp::And => "PrimOpAnd".to_string(),
67        BinOp::Or => "PrimOpOr".to_string(),
68        BinOp::Iff => "PrimOpIff".to_string(),
69        BinOp::Imp => "PrimOpImp".to_string(),
70    }
71}
72
73/// Helper macro to create Vec<String> from string-like values
74macro_rules! string_vec {
75    ($($s:expr),* $(,)?) => {
76        vec![$($s.to_string()),*]
77    };
78}
79
80fn vc_name(genv: GlobalEnv, def_id: DefId) -> String {
81    def_id_to_pascal_case(&def_id, &genv.tcx())
82}
83
84fn proof_name(genv: GlobalEnv, def_id: DefId) -> String {
85    format!("{}_proof", vc_name(genv, def_id))
86}
87
88fn project() -> String {
89    config::lean_project().to_string()
90}
91
92fn namespaced<F>(f: &mut fs::File, w: F) -> io::Result<()>
93where
94    F: Fn(&mut fs::File) -> io::Result<()>,
95{
96    let namespace = "F";
97    writeln!(f, "\nnamespace {namespace}\n")?;
98    w(f)?;
99    writeln!(f, "\nend {namespace}")
100}
101
102// Via Gemini: https://gemini.google.com/share/9027e898b136
103/// Renames all files and directories from 'src' to 'dst'
104fn rename_dir_contents(src: &Path, dst: &Path) -> io::Result<()> {
105    // 1. Ensure the destination root exists
106    if !dst.exists() {
107        fs::create_dir_all(dst)?;
108    }
109
110    for entry in fs::read_dir(src)? {
111        let entry = entry?;
112        let file_type = entry.file_type()?;
113        let src_path = entry.path();
114        let dst_path = dst.join(entry.file_name());
115
116        if file_type.is_dir() {
117            // 2. If it's a directory, recurse
118            rename_dir_contents(&src_path, &dst_path)?;
119        } else {
120            // 3. If it's a file, rename (overwrites if exists)
121            // On Windows, this will fail if the target file is "in use"
122            // TODO: how to FORCE overwrite on Windows?
123            fs::rename(&src_path, &dst_path)?;
124        }
125    }
126    Ok(())
127}
128
129fn constant_deps(expr: &fixpoint::Expr, acc: &mut FxIndexSet<fixpoint::Var>) {
130    match expr {
131        fixpoint::Expr::Var(v) if matches!(v, fixpoint::Var::Const(..)) => {
132            acc.insert(*v);
133        }
134        fixpoint::Expr::App(func, _, args, _) => {
135            constant_deps(func, acc);
136            args.iter().for_each(|expr| constant_deps(expr, acc));
137        }
138        fixpoint::Expr::And(inner) | fixpoint::Expr::Or(inner) => {
139            inner.iter().for_each(|expr| constant_deps(expr, acc));
140        }
141        fixpoint::Expr::Atom(_, inner)
142        | fixpoint::Expr::BinaryOp(_, inner)
143        | fixpoint::Expr::Imp(inner)
144        | fixpoint::Expr::Iff(inner)
145        | fixpoint::Expr::Let(_, inner) => {
146            inner.iter().for_each(|expr| constant_deps(expr, acc));
147        }
148        fixpoint::Expr::IfThenElse(inner) => {
149            inner.iter().for_each(|expr| constant_deps(expr, acc));
150        }
151        fixpoint::Expr::Exists(_, inner)
152        | fixpoint::Expr::Neg(inner)
153        | fixpoint::Expr::Not(inner)
154        | fixpoint::Expr::IsCtor(_, inner) => {
155            constant_deps(inner, acc);
156        }
157        fixpoint::Expr::Var(..) | fixpoint::Expr::Constant(..) | fixpoint::Expr::ThyFunc(..) => {}
158    }
159}
160
161pub fn finalize(genv: GlobalEnv) -> io::Result<()> {
162    let project = project();
163    let src = genv.temp_dir().path().join(&project);
164    let dst = final_project_path(genv);
165    if src.exists() { rename_dir_contents(&src, &dst) } else { Ok(()) }
166}
167
168fn final_project_path(genv: GlobalEnv) -> PathBuf {
169    genv.lean_parent_dir().join(project())
170}
171
172fn project_path(genv: GlobalEnv, kind: FileKind) -> PathBuf {
173    let project = project();
174    match kind {
175        FileKind::Flux => genv.temp_dir().path().join(project),
176        FileKind::User => final_project_path(genv),
177    }
178}
179
180fn run_proof(genv: GlobalEnv, def_id: DefId) -> io::Result<()> {
181    let proof_path = LeanFile::Proof(def_id).path(genv, true);
182    let out = Command::new("lake")
183        .arg("--quiet")
184        .arg("--log-level=error")
185        .arg("lean")
186        .arg(proof_path)
187        .stdout(Stdio::piped())
188        .stderr(Stdio::piped())
189        .current_dir(project_path(genv, FileKind::User))
190        .spawn()?
191        .wait_with_output()?;
192    if out.stderr.is_empty() && out.stdout.is_empty() {
193        Ok(())
194    } else {
195        let stderr =
196            std::str::from_utf8(&out.stderr).unwrap_or("Lean exited with a non-zero return code");
197        Err(io::Error::other(stderr))
198    }
199}
200
201fn run_check(genv: GlobalEnv, def_id: DefId) -> io::Result<()> {
202    let checking_path = LeanFile::Checking(def_id).path(genv, true);
203    let status = Command::new("lake")
204        .arg("--quiet")
205        .arg("--log-level=error")
206        .arg("lean")
207        .arg(checking_path)
208        .current_dir(project_path(genv, FileKind::User))
209        .spawn()?
210        .wait()?;
211    if status.success() {
212        Ok(())
213    } else {
214        Err(io::Error::other("Lean existed with a non-zero exit code"))
215    }
216}
217
218fn run_lean(genv: GlobalEnv, def_id: DefId) -> io::Result<()> {
219    run_proof(genv, def_id)?;
220    run_check(genv, def_id)?;
221    Ok(())
222}
223
224pub fn check_proof(genv: GlobalEnv, def_id: DefId) -> Result<(), ErrorGuaranteed> {
225    run_lean(genv, def_id)
226        .map_err(|_| {
227            let name = genv.tcx().def_path(def_id).to_string_no_crate_verbose();
228            let msg = format!("failed to check external proof for `crate{name}`");
229            let span = genv.tcx().def_span(def_id);
230            QueryErr::Emitted(genv.sess().dcx().handle().struct_span_err(span, msg).emit())
231        })
232        .emit(&genv)?;
233    Ok(())
234}
235
236/// Create a file at the given path, creating any missing parent directories.
237fn create_file_with_dirs<P: AsRef<Path>>(path: P) -> io::Result<Option<fs::File>> {
238    let path = path.as_ref();
239    if let Some(parent) = path.parent() {
240        fs::create_dir_all(parent)?;
241    }
242    match OpenOptions::new().write(true).create_new(true).open(path) {
243        Ok(file) => Ok(Some(file)),
244        Err(e) if e.kind() == io::ErrorKind::AlreadyExists => Ok(None),
245        Err(e) => Err(e),
246    }
247}
248
249/// Create or truncate a file at the given path, creating any missing parent directories.
250fn create_or_truncate_file_with_dirs<P: AsRef<Path>>(path: P) -> io::Result<fs::File> {
251    let path = path.as_ref();
252    if let Some(parent) = path.parent() {
253        fs::create_dir_all(parent)?;
254    }
255    OpenOptions::new()
256        .write(true)
257        .create(true)
258        .truncate(true)
259        .open(path)
260}
261
262#[derive(Eq, PartialEq, Hash, Debug, Clone)]
263pub enum FileKind {
264    /// Files that are only written by Flux
265    Flux,
266    /// Files that are modified by users
267    User,
268}
269
270/// Different kinds of Lean files
271#[derive(Eq, PartialEq, Hash, Debug, Clone)]
272pub enum LeanFile {
273    /// "Root" of the lean project, importing all the generated checking files
274    Basic,
275    /// "builtin" definitions
276    Fluxlib,
277    /// (human) opaque flux sorts, to be defined by the user in Lean
278    OpaqueSort(String),
279    /// (machine) sorts generated from flux definitions
280    Struct(String),
281    /// (human) opaque flux functions, to be defined by the user in Lean
282    OpaqueFun(String),
283    /// (machine) functions generated from flux definitions
284    Fun(String),
285    /// (machine) opaque constant encoded as a Lean axiom
286    OpaqueConst(String),
287    /// (machine) propositions holding the flux VCs
288    Vc(DefId),
289    /// (human) interactively written proofs of flux VCs
290    Proof(DefId),
291    /// (machine) files checking that a proof has the expected theorem type
292    Checking(DefId),
293}
294
295impl LeanFile {
296    fn kind(&self) -> FileKind {
297        match self {
298            LeanFile::Basic
299            | LeanFile::Fluxlib
300            | LeanFile::Vc(_)
301            | LeanFile::Checking(_)
302            | LeanFile::Fun(_)
303            | LeanFile::OpaqueConst(_)
304            | LeanFile::Struct(_) => FileKind::Flux,
305            LeanFile::OpaqueSort(_) | LeanFile::OpaqueFun(_) | LeanFile::Proof(_) => FileKind::User,
306        }
307    }
308
309    fn segments(&self, genv: GlobalEnv) -> Vec<String> {
310        let project_name = snake_case_to_pascal_case(&project());
311        match self {
312            LeanFile::Basic => {
313                string_vec![project_name, "Basic"]
314            }
315            LeanFile::Fluxlib => {
316                string_vec![project_name, "Flux", "Prelude"]
317            }
318            LeanFile::OpaqueSort(name) => {
319                // let name = self.datasort_name(sort);
320                string_vec![project_name, "User", "Struct", name]
321            }
322            LeanFile::Struct(name) => {
323                // let name = self.datasort_name(sort);
324                string_vec![project_name, "Flux", "Struct", name]
325            }
326            LeanFile::OpaqueFun(name) => {
327                // let name = self.var_name(name);
328                string_vec![project_name, "User", "Fun", name]
329            }
330            LeanFile::Fun(name) => {
331                // let name = self.var_name(name);
332                string_vec![project_name, "Flux", "Fun", name]
333            }
334            LeanFile::OpaqueConst(name) => {
335                string_vec![project_name, "Flux", "Const", name]
336            }
337            LeanFile::Vc(def_id) => {
338                let name = vc_name(genv, *def_id);
339                string_vec![project_name, "Flux", "VC", name]
340            }
341            LeanFile::Proof(def_id) => {
342                let name = format!("{}Proof", vc_name(genv, *def_id));
343                string_vec![project_name, "User", "Proof", name]
344            }
345            LeanFile::Checking(def_id) => {
346                let name = vc_name(genv, *def_id);
347                string_vec![project_name, "Flux", "Checking", name]
348            }
349        }
350    }
351
352    /// All paths should be generated here
353    fn path(&self, genv: GlobalEnv, force_final: bool) -> PathBuf {
354        let mut path =
355            if force_final { final_project_path(genv) } else { project_path(genv, self.kind()) };
356        for segment in self.segments(genv) {
357            path = path.join(segment);
358        }
359        path.set_extension("lean");
360        path
361    }
362
363    pub fn import(&self, genv: GlobalEnv) -> String {
364        format!("import {}", self.segments(genv).join("."))
365    }
366}
367
368pub struct LeanEncoder<'genv, 'tcx> {
369    genv: GlobalEnv<'genv, 'tcx>,
370    def_id: MaybeExternId,
371    pretty_var_map: PrettyMap<fixpoint::LocalVar>,
372    sort_deps: SortDeps,
373    fun_deps: Vec<fixpoint::FunDef>,
374    constants: ConstDeps,
375    kvar_solutions: KVarSolutions,
376    kvar_decls: Vec<fixpoint::KVarDecl>,
377    constraint: fixpoint::Constraint,
378    sort_files: FxHashMap<fixpoint::DataSort, LeanFile>,
379    fun_files: FxHashMap<FluxDefId, LeanFile>,
380    const_files: FxHashMap<fixpoint::Var, LeanFile>,
381    primop_var_map: FxHashMap<fixpoint::GlobalVar, String>,
382}
383
384impl<'genv, 'tcx> LeanEncoder<'genv, 'tcx> {
385    fn lean_cx(&self) -> LeanCtxt<'_, 'genv, 'tcx> {
386        LeanCtxt {
387            genv: self.genv,
388            pretty_var_map: &self.pretty_var_map,
389            adt_map: &self.sort_deps.adt_map,
390            opaque_adt_map: &self.sort_deps.opaque_sorts,
391            kvar_solutions: &self.kvar_solutions,
392            primop_var_map: &self.primop_var_map,
393        }
394    }
395
396    fn datasort_name(&self, sort: &fixpoint::DataSort) -> String {
397        let name = format!("{}", WithLeanCtxt { item: sort, cx: &self.lean_cx() });
398        snake_case_to_pascal_case(&name)
399    }
400
401    fn lean_file_for_fun(&self, fun: &fixpoint::FunDef) -> LeanFile {
402        let name = self.var_name(&fun.name);
403        if fun.body.is_some() { LeanFile::Fun(name) } else { LeanFile::OpaqueFun(name) }
404    }
405
406    fn lean_file_for_interpreted_const(&self, const_: &InterpretedConst) -> LeanFile {
407        let name = self.var_name(&const_.0.name);
408        LeanFile::Fun(name)
409    }
410
411    fn var_name(&self, var: &fixpoint::Var) -> String {
412        let name = format!("{}", WithLeanCtxt { item: var, cx: &self.lean_cx() });
413        snake_case_to_pascal_case(&name)
414    }
415
416    fn open_classical(&self) -> &str {
417        "open Classical"
418    }
419
420    fn new(
421        genv: GlobalEnv<'genv, 'tcx>,
422        def_id: MaybeExternId,
423        pretty_var_map: PrettyMap<fixpoint::LocalVar>,
424        sort_deps: SortDeps,
425        fun_deps: Vec<fixpoint::FunDef>,
426        constants: ConstDeps,
427        kvar_solutions: KVarSolutions,
428        kvar_decls: Vec<fixpoint::KVarDecl>,
429        constraint: fixpoint::Constraint,
430    ) -> io::Result<Self> {
431        let primop_var_map: FxHashMap<fixpoint::GlobalVar, String> = constants
432            .opaque
433            .iter()
434            .filter_map(|(decl, op)| {
435                if let fixpoint::Var::Const(gvar, None) = decl.name {
436                    Some((gvar, prim_op_lean_name(op)))
437                } else {
438                    None
439                }
440            })
441            .collect();
442        let mut encoder = Self {
443            genv,
444            def_id,
445            pretty_var_map,
446            sort_deps,
447            fun_deps,
448            constants,
449            kvar_decls,
450            kvar_solutions,
451            constraint,
452            fun_files: FxHashMap::default(),
453            sort_files: FxHashMap::default(),
454            const_files: FxHashMap::default(),
455            primop_var_map,
456        };
457        encoder.fun_files = encoder.fun_files();
458        encoder.sort_files = encoder.sort_files();
459        encoder.const_files = encoder.const_files();
460        Ok(encoder)
461    }
462
463    fn run(&self) -> io::Result<()> {
464        self.generate_lake_project_if_not_present()?;
465        self.generate_lib_if_absent()?;
466        self.generate_vc_file()?;
467        self.generate_proof_if_absent()?;
468        self.generate_checking_file()?;
469        Ok(())
470    }
471
472    fn fun_files(&self) -> FxHashMap<FluxDefId, LeanFile> {
473        let mut res = FxHashMap::default();
474        for fun_def in &self.fun_deps {
475            let fixpoint::Var::Global(_, did) = fun_def.name else {
476                bug!("expected global var with id")
477            };
478            let name = self.var_name(&fun_def.name);
479            let file = LeanFile::Fun(name);
480            res.insert(did, file);
481        }
482        res
483    }
484
485    fn sort_files(&self) -> FxHashMap<fixpoint::DataSort, LeanFile> {
486        let mut res = FxHashMap::default();
487        for (_, sort) in &self.sort_deps.opaque_sorts {
488            let data_sort = sort.name.clone();
489            let name = self.datasort_name(&sort.name);
490            let file = LeanFile::OpaqueSort(name);
491            res.insert(data_sort, file);
492        }
493        for data_decl in &self.sort_deps.data_decls {
494            let data_sort = data_decl.name.clone();
495            let name = self.datasort_name(&data_decl.name);
496            let file = LeanFile::Struct(name);
497            res.insert(data_sort, file);
498        }
499        res
500    }
501
502    fn const_files(&self) -> FxHashMap<fixpoint::Var, LeanFile> {
503        let mut res = FxHashMap::default();
504        for (decl, _) in &self.constants.interpreted {
505            res.insert(decl.name, LeanFile::Fun(self.var_name(&decl.name)));
506        }
507        for (decl, op) in &self.constants.opaque {
508            res.insert(decl.name, LeanFile::OpaqueConst(prim_op_lean_name(op)));
509        }
510        res
511    }
512
513    fn generate_lake_project_if_not_present(&self) -> io::Result<()> {
514        let path = project_path(self.genv, FileKind::User);
515        if !path.exists() {
516            Command::new("lake")
517                .current_dir(self.genv.lean_parent_dir())
518                .arg("+v4.28.0")
519                .arg("new")
520                .arg(project())
521                .arg("lib")
522                .spawn()
523                .and_then(|mut child| child.wait())
524                .map(|_| ())?;
525        }
526        Ok(())
527    }
528
529    fn generate_opaque_sort_file_if_not_present(
530        &self,
531        sort: &fixpoint::SortDecl,
532    ) -> io::Result<()> {
533        let name = self.datasort_name(&sort.name);
534        let file = &LeanFile::OpaqueSort(name);
535
536        let path = file.path(self.genv, false);
537        if let Some(mut file) = create_file_with_dirs(path)? {
538            writeln!(file, "{}", &LeanFile::Fluxlib.import(self.genv))?;
539            writeln!(file, "{}", self.open_classical())?;
540            namespaced(&mut file, |f| {
541                writeln!(f, "def {} := sorry", WithLeanCtxt { item: sort, cx: &self.lean_cx() })
542            })?;
543            file.sync_all()?;
544        }
545        Ok(())
546    }
547
548    fn data_decl_dependencies(&self, data_decl: &fixpoint::DataDecl) -> Vec<&LeanFile> {
549        let name = &data_decl.name;
550        let mut acc = vec![];
551        data_decl.deps(&mut acc);
552        acc.into_iter()
553            .map(|data_sort| {
554                self.sort_files.get(&data_sort).unwrap_or_else(|| {
555                    panic!(
556                        "Missing sort file for dependency {:?} of data decl {:?}",
557                        data_sort, name
558                    )
559                })
560            })
561            .unique()
562            .collect()
563    }
564
565    fn generate_struct_file_if_not_present(
566        &self,
567        data_decl: &fixpoint::DataDecl,
568    ) -> io::Result<()> {
569        let name = self.datasort_name(&data_decl.name);
570        let file = &LeanFile::Struct(name);
571        let path = file.path(self.genv, false);
572        // No need to regenerate if created in this session; but otherwise regenerate as struct may have changed
573        if let Some(mut file) = create_file_with_dirs(path)? {
574            // import prelude
575            writeln!(file, "{}", &LeanFile::Fluxlib.import(self.genv))?;
576            // import sort dependencies
577            for dep in self.data_decl_dependencies(data_decl) {
578                writeln!(file, "{}", dep.import(self.genv))?;
579            }
580            writeln!(file, "{}", self.open_classical())?;
581
582            // write data decl
583            namespaced(&mut file, |f| {
584                writeln!(f, "{}", WithLeanCtxt { item: data_decl, cx: &self.lean_cx() })
585            })?;
586            file.sync_all()?;
587        }
588        Ok(())
589    }
590
591    fn sort_file(&self, sort: &fixpoint::DataSort) -> &LeanFile {
592        self.sort_files
593            .get(sort)
594            .unwrap_or_else(|| panic!("Missing sort file for sort {:?}", sort))
595    }
596
597    fn fun_file(&self, did: &FluxDefId) -> &LeanFile {
598        self.fun_files
599            .get(did)
600            .unwrap_or_else(|| panic!("Missing fun file for fun {:?}", did))
601    }
602
603    fn const_file(&self, name: &fixpoint::Var) -> &LeanFile {
604        self.const_files
605            .get(name)
606            .unwrap_or_else(|| panic!("Missing const file for const {name:?}"))
607    }
608
609    fn fun_def_dependencies(&self, did: FluxDefId, fun_def: &fixpoint::FunDef) -> Vec<&LeanFile> {
610        let mut res = vec![];
611
612        // 1. Collect the sort dependencies
613        let mut sorts = vec![];
614        fun_def.sort.deps(&mut sorts);
615        for data_sort in sorts {
616            res.push(self.sort_file(&data_sort));
617        }
618
619        // 2. Collect the fun dependencies
620        if !self.genv.normalized_info(did).uif {
621            let body = self.genv.inlined_body(did);
622            for dep_id in local_deps(&body) {
623                res.push(self.fun_file(&dep_id.to_def_id()));
624            }
625        }
626
627        let mut deps = FxIndexSet::default();
628        if let Some(body) = &fun_def.body {
629            constant_deps(&body.expr, &mut deps);
630        }
631        for (decl, _) in &self.constants.interpreted {
632            if deps.contains(&decl.name) {
633                res.push(self.const_file(&decl.name));
634            }
635        }
636        res
637    }
638
639    fn generate_fun_def_file_if_not_present(
640        &self,
641        did: FluxDefId,
642        fun_def: &fixpoint::FunDef,
643    ) -> io::Result<()> {
644        let path = self.lean_file_for_fun(fun_def).path(self.genv, false);
645        if let Some(mut file) = create_file_with_dirs(path)? {
646            // import prelude
647            writeln!(file, "{}", &LeanFile::Fluxlib.import(self.genv))?;
648            // import sort dependencies
649            for dep in self.fun_def_dependencies(did, fun_def) {
650                writeln!(file, "{}", dep.import(self.genv))?;
651            }
652            writeln!(file, "{}", self.open_classical())?;
653
654            // write fun def
655            namespaced(&mut file, |f| {
656                writeln!(f, "{}", WithLeanCtxt { item: fun_def, cx: &self.lean_cx() })
657            })?;
658            file.sync_all()?;
659        }
660        Ok(())
661    }
662
663    fn generate_interpreted_const_file_if_not_present(
664        &self,
665        interpreted_const: &InterpretedConst,
666    ) -> io::Result<()> {
667        let (const_decl, _) = interpreted_const;
668        let path = self
669            .lean_file_for_interpreted_const(interpreted_const)
670            .path(self.genv, false);
671        if let Some(mut file) = create_file_with_dirs(path)? {
672            // import prelude
673            writeln!(file, "{}", &LeanFile::Fluxlib.import(self.genv))?;
674
675            let mut sort_deps = vec![];
676            const_decl.sort.deps(&mut sort_deps);
677            for dep in sort_deps {
678                writeln!(file, "{}", self.sort_file(&dep).import(self.genv))?;
679            }
680
681            writeln!(file, "{}", self.open_classical())?;
682
683            namespaced(&mut file, |f| {
684                if let Some(comment) = &const_decl.comment {
685                    writeln!(f, "--{comment}")?;
686                }
687                writeln!(f, "{}", WithLeanCtxt { item: interpreted_const, cx: &self.lean_cx() })
688            })?;
689            file.sync_all()?;
690        }
691        Ok(())
692    }
693
694    fn generate_opaque_const_file(
695        &self,
696        const_decl: &fixpoint::ConstDecl,
697        op: &BinOp,
698    ) -> io::Result<()> {
699        let stable_name = prim_op_lean_name(op);
700        let file = LeanFile::OpaqueConst(stable_name.clone());
701        let path = file.path(self.genv, false);
702        let mut file = create_or_truncate_file_with_dirs(path)?;
703        writeln!(file, "{}", &LeanFile::Fluxlib.import(self.genv))?;
704
705        let mut sort_deps = vec![];
706        const_decl.sort.deps(&mut sort_deps);
707        for dep in sort_deps {
708            writeln!(file, "{}", self.sort_file(&dep).import(self.genv))?;
709        }
710
711        writeln!(file, "{}", self.open_classical())?;
712
713        namespaced(&mut file, |f| {
714            if let Some(comment) = &const_decl.comment {
715                writeln!(f, "--{comment}")?;
716            }
717            writeln!(
718                f,
719                "axiom {stable_name} : {}",
720                WithLeanCtxt { item: &const_decl.sort, cx: &self.lean_cx() }
721            )
722        })?;
723        file.sync_all()?;
724        Ok(())
725    }
726
727    fn generate_lib_if_absent(&self) -> io::Result<()> {
728        let path = LeanFile::Fluxlib.path(self.genv, false);
729        if let Some(mut file) = create_file_with_dirs(path)? {
730            writeln!(file, "-- FLUX LIBRARY [DO NOT MODIFY] --")?;
731            // TODO: Can't we write this from a single `write!` call?
732            writeln!(
733                file,
734                "abbrev BitVec_shiftLeft {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.shiftLeft x (s.toNat)"
735            )?;
736            writeln!(
737                file,
738                "abbrev BitVec_ushiftRight {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.ushiftRight x (s.toNat)"
739            )?;
740            writeln!(
741                file,
742                "abbrev BitVec_sshiftRight {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.sshiftRight x (s.toNat)"
743            )?;
744            writeln!(
745                file,
746                "abbrev BitVec_uge {{ n : Nat }} (x y : BitVec n) := (BitVec.ult x y).not"
747            )?;
748            writeln!(
749                file,
750                "abbrev BitVec_sge {{ n : Nat }} (x y : BitVec n) := (BitVec.slt x y).not"
751            )?;
752            writeln!(
753                file,
754                "abbrev BitVec_ugt {{ n : Nat }} (x y : BitVec n) := (BitVec.ule x y).not"
755            )?;
756            writeln!(
757                file,
758                "abbrev BitVec_sgt {{ n : Nat }} (x y : BitVec n) := (BitVec.sle x y).not"
759            )?;
760            writeln!(
761                file,
762                "abbrev BitVec_zeroExtend {{n : Nat}} (extra : Nat) (x : BitVec n) : BitVec (n + extra) := BitVec.zeroExtend (n + extra) x"
763            )?;
764            writeln!(
765                file,
766                "abbrev BitVec_signExtend {{n : Nat}} (extra : Nat) (x : BitVec n) : BitVec (n + extra) := BitVec.signExtend (n + extra) x"
767            )?;
768            writeln!(
769                file,
770                "abbrev SmtMap (t0 t1 : Type) [Inhabited t0] [BEq t0] [Inhabited t1] : Type := t0 -> t1"
771            )?;
772            writeln!(
773                file,
774                "abbrev SmtMap_default {{ t0 t1: Type }} (v : t1) [Inhabited t0] [BEq t0] [Inhabited t1] : SmtMap t0 t1 := fun _ => v"
775            )?;
776            writeln!(
777                file,
778                "abbrev SmtMap_store {{ t0 t1 : Type }} [Inhabited t0] [BEq t0] [Inhabited t1] (m : SmtMap t0 t1) (k : t0) (v : t1) : SmtMap t0 t1 :=\n  fun x => if x == k then v else m x"
779            )?;
780            writeln!(
781                file,
782                "abbrev SmtMap_select {{ t0 t1 : Type }} [Inhabited t0] [BEq t0] [Inhabited t1] (m : SmtMap t0 t1) (k : t0) := m k"
783            )?;
784        }
785        Ok(())
786    }
787
788    fn generate_vc_prelude(&self) -> io::Result<()> {
789        // 1. Generate lake project and lib file
790        self.generate_lib_if_absent()?;
791
792        // 2. Generate Opaque Struct Files
793        for (_, sort) in &self.sort_deps.opaque_sorts {
794            self.generate_opaque_sort_file_if_not_present(sort)?;
795        }
796        // 2. Generate Struct Files
797        for data_decl in &self.sort_deps.data_decls {
798            self.generate_struct_file_if_not_present(data_decl)?;
799        }
800        // 3. Generate Func Def Files
801        for fun_def in &self.fun_deps {
802            let fixpoint::Var::Global(_, did) = fun_def.name else {
803                bug!("expected global var with id")
804            };
805            self.generate_fun_def_file_if_not_present(did, fun_def)?;
806        }
807        // 4. Generate Const Decl Files
808        for const_decl in &self.constants.interpreted {
809            self.generate_interpreted_const_file_if_not_present(const_decl)?;
810        }
811        // 5. Generate Opaque Const Files (primop axioms)
812        for (const_decl, op) in &self.constants.opaque {
813            self.generate_opaque_const_file(const_decl, op)?;
814        }
815        Ok(())
816    }
817
818    fn generate_vc_imports(&self, file: &mut fs::File) -> io::Result<()> {
819        writeln!(file, "{}", LeanFile::Fluxlib.import(self.genv))?;
820
821        for (_, sort) in &self.sort_deps.opaque_sorts {
822            let name = self.datasort_name(&sort.name);
823            writeln!(file, "{}", LeanFile::OpaqueSort(name).import(self.genv))?;
824        }
825
826        for data_decl in &self.sort_deps.data_decls {
827            let name = self.datasort_name(&data_decl.name);
828            writeln!(file, "{}", LeanFile::Struct(name).import(self.genv))?;
829        }
830
831        for fun_def in &self.fun_deps {
832            writeln!(file, "{}", self.lean_file_for_fun(fun_def).import(self.genv))?;
833        }
834
835        for const_decl in &self.constants.interpreted {
836            writeln!(
837                file,
838                "{}",
839                self.lean_file_for_interpreted_const(const_decl)
840                    .import(self.genv)
841            )?;
842        }
843
844        for (_, op) in &self.constants.opaque {
845            writeln!(file, "{}", LeanFile::OpaqueConst(prim_op_lean_name(op)).import(self.genv))?;
846        }
847
848        Ok(())
849    }
850
851    fn generate_vc_file(&self) -> io::Result<()> {
852        // 1. Generate imports
853        self.generate_vc_prelude()?;
854
855        // 2. Create file and add imports
856        let def_id = self.def_id.resolved_id();
857        let path = LeanFile::Vc(def_id).path(self.genv, false);
858        if let Some(mut file) = create_file_with_dirs(path)? {
859            self.generate_vc_imports(&mut file)?;
860            writeln!(file, "{}", self.open_classical())?;
861
862            let vc_name = vc_name(self.genv, def_id);
863            // 3. Write the VC
864            namespaced(&mut file, |f| {
865                write!(
866                    f,
867                    "{}",
868                    WithLeanCtxt {
869                        item: lean_format::LeanKConstraint {
870                            theorem_name: &vc_name,
871                            kvars: &self.kvar_decls,
872                            constr: &self.constraint,
873                            should_fail: self
874                                .def_id
875                                .as_local()
876                                .map(|def_id| self.genv.should_fail(def_id))
877                                .unwrap_or(false)
878                        },
879                        cx: &self.lean_cx()
880                    }
881                )
882            })?;
883            file.sync_all()?;
884        }
885
886        Ok(())
887    }
888
889    fn generate_proof_if_absent(&self) -> io::Result<()> {
890        let def_id = self.def_id.resolved_id();
891        let vc_name = vc_name(self.genv, def_id);
892        let proof_name = proof_name(self.genv, def_id);
893        let path = LeanFile::Proof(def_id).path(self.genv, false);
894
895        if let Some(mut file) = create_file_with_dirs(path)? {
896            writeln!(file, "{}", LeanFile::Fluxlib.import(self.genv))?;
897            writeln!(file, "{}", LeanFile::Vc(def_id).import(self.genv))?;
898            writeln!(file, "{}", self.open_classical())?;
899            namespaced(&mut file, |f| {
900                writeln!(f, "def {proof_name} : {vc_name} := by")?;
901                writeln!(f, "  unfold {vc_name}")?;
902                writeln!(f, "  sorry")
903            })?;
904            file.sync_all()?;
905        }
906        Ok(())
907    }
908
909    fn generate_checking_file(&self) -> io::Result<()> {
910        let def_id = self.def_id.resolved_id();
911        let vc_name = vc_name(self.genv, def_id);
912        let proof_name = proof_name(self.genv, def_id);
913        let path = LeanFile::Checking(def_id).path(self.genv, false);
914
915        let mut file = create_or_truncate_file_with_dirs(path)?;
916        writeln!(file, "{}", LeanFile::Vc(def_id).import(self.genv))?;
917        writeln!(file, "{}", LeanFile::Proof(def_id).import(self.genv))?;
918        writeln!(file)?;
919        writeln!(file, "#check (F.{proof_name} : F.{vc_name})")?;
920        file.sync_all()?;
921        Ok(())
922    }
923
924    pub fn encode(
925        genv: GlobalEnv<'genv, 'tcx>,
926        def_id: MaybeExternId,
927        pretty_var_map: PrettyMap<fixpoint::LocalVar>,
928        sort_deps: SortDeps,
929        fun_deps: Vec<fixpoint::FunDef>,
930        constants: ConstDeps,
931        kvar_decls: Vec<fixpoint::KVarDecl>,
932        constraint: fixpoint::Constraint,
933        kvar_solutions: KVarSolutions,
934    ) -> io::Result<()> {
935        let encoder = Self::new(
936            genv,
937            def_id,
938            pretty_var_map,
939            sort_deps,
940            fun_deps,
941            constants,
942            kvar_solutions,
943            kvar_decls,
944            constraint,
945        )?;
946        encoder.run()?;
947        Ok(())
948    }
949}
950
951fn hyperlink_proof(genv: GlobalEnv, def_id: MaybeExternId) {
952    let proof_name = proof_name(genv, def_id.resolved_id());
953    let path = LeanFile::Proof(def_id.resolved_id()).path(genv, false);
954    if let Some(span) = genv.proven_externally(def_id.local_id()) {
955        let dst_span = SpanTrace::from_path(&path, 3, 5, proof_name.len());
956        dbg::hyperlink_json!(genv.tcx(), span, dst_span);
957    }
958}
959
960fn record_proof(genv: GlobalEnv, def_id: MaybeExternId) -> io::Result<()> {
961    let path = LeanFile::Basic.path(genv, false);
962
963    let mut file = match create_file_with_dirs(&path)? {
964        Some(mut file) => {
965            // First invocation: reset VCs
966            writeln!(file, "-- Flux Basic Imports [DO NOT MODIFY] --")?;
967            file
968        }
969        None => fs::OpenOptions::new().append(true).open(path)?,
970    };
971    writeln!(file, "{}", LeanFile::Checking(def_id.resolved_id()).import(genv))
972}
973
974/// We need to both hyperlink the proof (so users can easily jump to it)
975/// and record the checking file in `Basic.lean` (so that it gets checked by `lake build`),
976/// regardless of whether the proof was cached.
977pub fn log_proof(genv: GlobalEnv, def_id: MaybeExternId) -> Result<(), ErrorGuaranteed> {
978    hyperlink_proof(genv, def_id);
979    record_proof(genv, def_id)
980        .map_err(|_| {
981            let name = genv
982                .tcx()
983                .def_path(def_id.resolved_id())
984                .to_string_no_crate_verbose();
985            let msg = format!("failed to record proof for `{name}`");
986            let span = genv.tcx().def_span(def_id);
987            QueryErr::Emitted(genv.sess().dcx().handle().struct_span_err(span, msg).emit())
988        })
989        .emit(&genv)?;
990    Ok(())
991}