Skip to main content

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