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::{PrettyMap, local_deps},
19};
20use itertools::Itertools;
21use rustc_hash::FxHashMap;
22use rustc_hir::def_id::DefId;
23use rustc_span::ErrorGuaranteed;
24
25use crate::{
26    fixpoint_encoding::{KVarSolutions, SortDeps, fixpoint},
27    lean_format::{
28        self, BoolMode, LeanCtxt, WithLeanCtxt, def_id_to_pascal_case, snake_case_to_pascal_case,
29    },
30};
31
32/// Helper macro to create Vec<String> from string-like values
33macro_rules! string_vec {
34    ($($s:expr),* $(,)?) => {
35        vec![$($s.to_string()),*]
36    };
37}
38
39fn vc_name(genv: GlobalEnv, def_id: DefId) -> String {
40    def_id_to_pascal_case(&def_id, &genv.tcx())
41}
42
43fn project() -> String {
44    config::lean_project().to_string()
45}
46
47// Via Gemini: https://gemini.google.com/share/9027e898b136
48/// Renames all files and directories from 'src' to 'dst'
49fn rename_dir_contents(src: &Path, dst: &Path) -> io::Result<()> {
50    // 1. Ensure the destination root exists
51    if !dst.exists() {
52        fs::create_dir_all(dst)?;
53    }
54
55    for entry in fs::read_dir(src)? {
56        let entry = entry?;
57        let file_type = entry.file_type()?;
58        let src_path = entry.path();
59        let dst_path = dst.join(entry.file_name());
60
61        if file_type.is_dir() {
62            // 2. If it's a directory, recurse
63            rename_dir_contents(&src_path, &dst_path)?;
64        } else {
65            // 3. If it's a file, rename (overwrites if exists)
66            // On Windows, this will fail if the target file is "in use"
67            // TODO: how to FORCE overwrite on Windows?
68            fs::rename(&src_path, &dst_path)?;
69        }
70    }
71    Ok(())
72}
73
74pub fn finalize(genv: GlobalEnv) -> io::Result<()> {
75    let project = project();
76    let src = genv.temp_dir().path().join(&project);
77    let dst = genv.lean_parent_dir().join(&project);
78    if src.exists() { rename_dir_contents(&src, &dst) } else { Ok(()) }
79}
80
81fn project_path(genv: GlobalEnv, kind: FileKind) -> PathBuf {
82    let project = project();
83    match kind {
84        FileKind::Flux => genv.temp_dir().path().join(project),
85        FileKind::User => genv.lean_parent_dir().join(project),
86    }
87}
88
89fn run_lean(genv: GlobalEnv, def_id: DefId) -> io::Result<()> {
90    let proof_path = LeanFile::Proof(def_id).path(genv);
91    let out = Command::new("lake")
92        .arg("--quiet")
93        .arg("--log-level=error")
94        .arg("lean")
95        .arg(proof_path)
96        .stdout(Stdio::piped())
97        .stderr(Stdio::piped())
98        .current_dir(project_path(genv, FileKind::User))
99        .spawn()?
100        .wait_with_output()?;
101    if out.stderr.is_empty() && out.stdout.is_empty() {
102        Ok(())
103    } else {
104        let stderr =
105            std::str::from_utf8(&out.stderr).unwrap_or("Lean exited with a non-zero return code");
106        Err(io::Error::other(stderr))
107    }
108}
109
110pub fn check_proof(genv: GlobalEnv, def_id: DefId) -> Result<(), ErrorGuaranteed> {
111    run_lean(genv, def_id)
112        .map_err(|_| {
113            let name = genv.tcx().def_path(def_id).to_string_no_crate_verbose();
114            let msg = format!("failed to check external proof for `crate{name}`");
115            let span = genv.tcx().def_span(def_id);
116            QueryErr::Emitted(genv.sess().dcx().handle().struct_span_err(span, msg).emit())
117        })
118        .emit(&genv)?;
119    Ok(())
120}
121
122/// Create a file at the given path, creating any missing parent directories.
123fn create_file_with_dirs<P: AsRef<Path>>(path: P) -> io::Result<Option<fs::File>> {
124    let path = path.as_ref();
125    if let Some(parent) = path.parent() {
126        fs::create_dir_all(parent)?;
127    }
128    match OpenOptions::new().write(true).create_new(true).open(path) {
129        Ok(file) => Ok(Some(file)),
130        Err(e) if e.kind() == io::ErrorKind::AlreadyExists => Ok(None),
131        Err(e) => Err(e),
132    }
133}
134
135#[derive(Eq, PartialEq, Hash, Debug, Clone)]
136pub enum FileKind {
137    /// Files that are only written by Flux
138    Flux,
139    /// Files that are modified by users
140    User,
141}
142
143/// Different kinds of Lean files
144#[derive(Eq, PartialEq, Hash, Debug, Clone)]
145pub enum LeanFile {
146    /// "Root" of the lean project, importing all the generated proofs
147    Basic,
148    /// "builtin" definitions
149    Fluxlib,
150    /// (human) opaque flux sorts, to be defined by the user in Lean
151    OpaqueSort(String),
152    /// (machine) sorts generated from flux definitions
153    Struct(String),
154    /// (human) opaque flux functions, to be defined by the user in Lean
155    OpaqueFun(String),
156    /// (machine) functions generated from flux definitions
157    Fun(String),
158    /// (machine) propositions holding the flux VCs
159    Vc(DefId),
160    /// (human) interactively written proofs of flux VCs
161    Proof(DefId),
162}
163
164impl LeanFile {
165    fn kind(&self) -> FileKind {
166        match self {
167            LeanFile::Basic
168            | LeanFile::Fluxlib
169            | LeanFile::Vc(_)
170            | LeanFile::Fun(_)
171            | LeanFile::Struct(_) => FileKind::Flux,
172            LeanFile::OpaqueSort(_) | LeanFile::OpaqueFun(_) | LeanFile::Proof(_) => FileKind::User,
173        }
174    }
175
176    fn segments(&self, genv: GlobalEnv) -> Vec<String> {
177        let project_name = snake_case_to_pascal_case(&project());
178        match self {
179            LeanFile::Basic => {
180                string_vec![project_name, "Basic"]
181            }
182            LeanFile::Fluxlib => {
183                string_vec![project_name, "Flux", "Prelude"]
184            }
185            LeanFile::OpaqueSort(name) => {
186                // let name = self.datasort_name(sort);
187                string_vec![project_name, "User", "Struct", name]
188            }
189            LeanFile::Struct(name) => {
190                // let name = self.datasort_name(sort);
191                string_vec![project_name, "Flux", "Struct", name]
192            }
193            LeanFile::OpaqueFun(name) => {
194                // let name = self.var_name(name);
195                string_vec![project_name, "User", "Fun", name]
196            }
197            LeanFile::Fun(name) => {
198                // let name = self.var_name(name);
199                string_vec![project_name, "Flux", "Fun", name]
200            }
201            LeanFile::Vc(def_id) => {
202                let name = vc_name(genv, *def_id);
203                string_vec![project_name, "Flux", "VC", name]
204            }
205            LeanFile::Proof(def_id) => {
206                let name = format!("{}Proof", vc_name(genv, *def_id));
207                string_vec![project_name, "User", "Proof", name]
208            }
209        }
210    }
211
212    /// All paths should be generated here
213    fn path(&self, genv: GlobalEnv) -> PathBuf {
214        let mut path = project_path(genv, self.kind());
215        for segment in self.segments(genv) {
216            path = path.join(segment);
217        }
218        path.set_extension("lean");
219        path
220    }
221}
222
223pub struct LeanEncoder<'genv, 'tcx> {
224    genv: GlobalEnv<'genv, 'tcx>,
225    def_id: MaybeExternId,
226    pretty_var_map: PrettyMap<fixpoint::LocalVar>,
227    sort_deps: SortDeps,
228    fun_deps: Vec<fixpoint::FunDef>,
229    kvar_solutions: KVarSolutions,
230    kvar_decls: Vec<fixpoint::KVarDecl>,
231    constraint: fixpoint::Constraint,
232    sort_files: FxHashMap<fixpoint::DataSort, LeanFile>,
233    fun_files: FxHashMap<FluxDefId, LeanFile>,
234}
235
236impl<'genv, 'tcx> LeanEncoder<'genv, 'tcx> {
237    fn lean_cx(&self) -> LeanCtxt<'_, 'genv, 'tcx> {
238        LeanCtxt {
239            genv: self.genv,
240            pretty_var_map: &self.pretty_var_map,
241            adt_map: &self.sort_deps.adt_map,
242            kvar_solutions: &self.kvar_solutions,
243            bool_mode: BoolMode::Bool,
244        }
245    }
246
247    fn datasort_name(&self, sort: &fixpoint::DataSort) -> String {
248        let name = format!("{}", WithLeanCtxt { item: sort, cx: &self.lean_cx() });
249        snake_case_to_pascal_case(&name)
250    }
251
252    fn lean_file_for_fun(&self, fun: &fixpoint::FunDef) -> LeanFile {
253        let name = self.var_name(&fun.name);
254        if fun.body.is_some() { LeanFile::Fun(name) } else { LeanFile::OpaqueFun(name) }
255    }
256
257    fn var_name(&self, var: &fixpoint::Var) -> String {
258        let name = format!("{}", WithLeanCtxt { item: var, cx: &self.lean_cx() });
259        snake_case_to_pascal_case(&name)
260    }
261
262    fn import(&self, file: &LeanFile) -> String {
263        format!("import {}", file.segments(self.genv).join("."))
264    }
265
266    fn new(
267        genv: GlobalEnv<'genv, 'tcx>,
268        def_id: MaybeExternId,
269        pretty_var_map: PrettyMap<fixpoint::LocalVar>,
270        sort_deps: SortDeps,
271        fun_deps: Vec<fixpoint::FunDef>,
272        kvar_solutions: KVarSolutions,
273        kvar_decls: Vec<fixpoint::KVarDecl>,
274        constraint: fixpoint::Constraint,
275    ) -> io::Result<Self> {
276        let mut encoder = Self {
277            genv,
278            def_id,
279            pretty_var_map,
280            sort_deps,
281            fun_deps,
282            kvar_decls,
283            kvar_solutions,
284            constraint,
285            fun_files: FxHashMap::default(),
286            sort_files: FxHashMap::default(),
287        };
288        encoder.fun_files = encoder.fun_files();
289        encoder.sort_files = encoder.sort_files();
290        Ok(encoder)
291    }
292
293    fn run(&self) -> io::Result<()> {
294        self.generate_lake_project_if_not_present()?;
295        self.generate_lib_if_absent()?;
296        self.generate_vc_file()?;
297        self.generate_proof_if_absent()?;
298        self.record_proof()?;
299        Ok(())
300    }
301
302    fn fun_files(&self) -> FxHashMap<FluxDefId, LeanFile> {
303        let mut res = FxHashMap::default();
304        for fun_def in &self.fun_deps {
305            let fixpoint::Var::Global(_, Some(did)) = fun_def.name else {
306                bug!("expected global var with id")
307            };
308            let name = self.var_name(&fun_def.name);
309            let file = LeanFile::Fun(name);
310            res.insert(did, file);
311        }
312        res
313    }
314
315    fn sort_files(&self) -> FxHashMap<fixpoint::DataSort, LeanFile> {
316        let mut res = FxHashMap::default();
317        for sort in &self.sort_deps.opaque_sorts {
318            let data_sort = sort.name.clone();
319            let name = self.datasort_name(&sort.name);
320            let file = LeanFile::OpaqueSort(name);
321            res.insert(data_sort, file);
322        }
323        for data_decl in &self.sort_deps.data_decls {
324            let data_sort = data_decl.name.clone();
325            let name = self.datasort_name(&data_decl.name);
326            let file = LeanFile::Struct(name);
327            res.insert(data_sort, file);
328        }
329        res
330    }
331
332    fn generate_lake_project_if_not_present(&self) -> io::Result<()> {
333        let path = project_path(self.genv, FileKind::User);
334        if !path.exists() {
335            Command::new("lake")
336                .current_dir(self.genv.lean_parent_dir())
337                .arg("new")
338                .arg(project())
339                .arg("lib")
340                .spawn()
341                .and_then(|mut child| child.wait())
342                .map(|_| ())?;
343        }
344        Ok(())
345    }
346
347    fn generate_opaque_sort_file_if_not_present(
348        &self,
349        sort: &fixpoint::SortDecl,
350    ) -> io::Result<()> {
351        let name = self.datasort_name(&sort.name);
352        let file = &LeanFile::OpaqueSort(name);
353
354        let path = file.path(self.genv);
355        if let Some(mut file) = create_file_with_dirs(path)? {
356            writeln!(file, "{}", self.import(&LeanFile::Fluxlib))?;
357            writeln!(file, "def {} := sorry", WithLeanCtxt { item: sort, cx: &self.lean_cx() })?;
358            file.sync_all()?;
359        }
360        Ok(())
361    }
362
363    fn data_decl_dependencies(&self, data_decl: &fixpoint::DataDecl) -> Vec<&LeanFile> {
364        let name = &data_decl.name;
365        let mut acc = vec![];
366        data_decl.deps(&mut acc);
367        acc.into_iter()
368            .map(|data_sort| {
369                self.sort_files.get(&data_sort).unwrap_or_else(|| {
370                    panic!(
371                        "Missing sort file for dependency {:?} of data decl {:?}",
372                        data_sort, name
373                    )
374                })
375            })
376            .unique()
377            .collect()
378    }
379
380    fn generate_struct_file_if_not_present(
381        &self,
382        data_decl: &fixpoint::DataDecl,
383    ) -> io::Result<()> {
384        let name = self.datasort_name(&data_decl.name);
385        let file = &LeanFile::Struct(name);
386        let path = file.path(self.genv);
387        // No need to regenerate if created in this session; but otherwise regenerate as struct may have changed
388        if let Some(mut file) = create_file_with_dirs(path)? {
389            // import prelude
390            writeln!(file, "{}", self.import(&LeanFile::Fluxlib))?;
391            // import sort dependencies
392            for dep in self.data_decl_dependencies(data_decl) {
393                writeln!(file, "{}", self.import(dep))?;
394            }
395            // write data decl
396            writeln!(file, "{}", WithLeanCtxt { item: data_decl, cx: &self.lean_cx() })?;
397            file.sync_all()?;
398        }
399        Ok(())
400    }
401
402    fn sort_file(&self, sort: &fixpoint::DataSort) -> &LeanFile {
403        self.sort_files
404            .get(sort)
405            .unwrap_or_else(|| panic!("Missing sort file for sort {:?}", sort))
406    }
407
408    fn fun_file(&self, did: &FluxDefId) -> &LeanFile {
409        self.fun_files
410            .get(did)
411            .unwrap_or_else(|| panic!("Missing fun file for fun {:?}", did))
412    }
413
414    fn fun_def_dependencies(&self, did: FluxDefId, fun_def: &fixpoint::FunDef) -> Vec<&LeanFile> {
415        let mut res = vec![];
416
417        // 1. Collect the sort dependencies
418        let mut sorts = vec![];
419        fun_def.sort.deps(&mut sorts);
420        for data_sort in sorts {
421            res.push(self.sort_file(&data_sort));
422        }
423
424        // 2. Collect the fun dependencies
425        if !self.genv.normalized_info(did).uif {
426            let body = self.genv.inlined_body(did);
427            for dep_id in local_deps(&body) {
428                res.push(self.fun_file(&dep_id.to_def_id()));
429            }
430        }
431        res
432    }
433
434    fn generate_fun_def_file_if_not_present(
435        &self,
436        did: FluxDefId,
437        fun_def: &fixpoint::FunDef,
438    ) -> io::Result<()> {
439        let path = self.lean_file_for_fun(fun_def).path(self.genv);
440        if let Some(mut file) = create_file_with_dirs(path)? {
441            // import prelude
442            writeln!(file, "{}", self.import(&LeanFile::Fluxlib))?;
443            // import sort dependencies
444            for dep in self.fun_def_dependencies(did, fun_def) {
445                writeln!(file, "{}", self.import(dep))?;
446            }
447            // write fun def
448            writeln!(file, "{}", WithLeanCtxt { item: fun_def, cx: &self.lean_cx() })?;
449            file.sync_all()?;
450        }
451        Ok(())
452    }
453
454    fn generate_lib_if_absent(&self) -> io::Result<()> {
455        let path = LeanFile::Fluxlib.path(self.genv);
456        if let Some(mut file) = create_file_with_dirs(path)? {
457            writeln!(file, "-- FLUX LIBRARY [DO NOT MODIFY] --")?;
458            // TODO: Can't we write this from a single `write!` call?
459            writeln!(
460                file,
461                "def BitVec_shiftLeft {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.shiftLeft x (s.toNat)"
462            )?;
463            writeln!(
464                file,
465                "def BitVec_ushiftRight {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.ushiftRight x (s.toNat)"
466            )?;
467            writeln!(
468                file,
469                "def BitVec_sshiftRight {{ n : Nat }} (x s : BitVec n) : BitVec n := BitVec.sshiftRight x (s.toNat)"
470            )?;
471            writeln!(
472                file,
473                "def BitVec_uge {{ n : Nat }} (x y : BitVec n) := (BitVec.ult x y).not"
474            )?;
475            writeln!(
476                file,
477                "def BitVec_sge {{ n : Nat }} (x y : BitVec n) := (BitVec.slt x y).not"
478            )?;
479            writeln!(
480                file,
481                "def BitVec_ugt {{ n : Nat }} (x y : BitVec n) := (BitVec.ule x y).not"
482            )?;
483            writeln!(
484                file,
485                "def BitVec_sgt {{ n : Nat }} (x y : BitVec n) := (BitVec.sle x y).not"
486            )?;
487            writeln!(
488                file,
489                "def SmtMap (t0 t1 : Type) [Inhabited t0] [BEq t0] [Inhabited t1] : Type := t0 -> t1"
490            )?;
491            writeln!(
492                file,
493                "def SmtMap_default {{ t0 t1: Type }} (v : t1) [Inhabited t0] [BEq t0] [Inhabited t1] : SmtMap t0 t1 := fun _ => v"
494            )?;
495            writeln!(
496                file,
497                "def 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"
498            )?;
499            writeln!(
500                file,
501                "def SmtMap_select {{ t0 t1 : Type }} [Inhabited t0] [BEq t0] [Inhabited t1] (m : SmtMap t0 t1) (k : t0) := m k"
502            )?;
503        }
504        Ok(())
505    }
506
507    fn generate_vc_prelude(&self) -> io::Result<()> {
508        // 1. Generate lake project and lib file
509        self.generate_lib_if_absent()?;
510
511        // 2. Generate Opaque Struct Files
512        for sort in &self.sort_deps.opaque_sorts {
513            self.generate_opaque_sort_file_if_not_present(sort)?;
514        }
515        // 2. Generate Struct Files
516        for data_decl in &self.sort_deps.data_decls {
517            self.generate_struct_file_if_not_present(data_decl)?;
518        }
519        // 3. Generate Func Def Files
520        for fun_def in &self.fun_deps {
521            let fixpoint::Var::Global(_, Some(did)) = fun_def.name else {
522                bug!("expected global var with id")
523            };
524            self.generate_fun_def_file_if_not_present(did, fun_def)?;
525        }
526        Ok(())
527    }
528
529    fn generate_vc_imports(&self, file: &mut fs::File) -> io::Result<()> {
530        writeln!(file, "{}", self.import(&LeanFile::Fluxlib))?;
531
532        for sort in &self.sort_deps.opaque_sorts {
533            let name = self.datasort_name(&sort.name);
534            writeln!(file, "{}", self.import(&LeanFile::OpaqueSort(name)))?;
535        }
536
537        for data_decl in &self.sort_deps.data_decls {
538            let name = self.datasort_name(&data_decl.name);
539            writeln!(file, "{}", self.import(&LeanFile::Struct(name)))?;
540        }
541
542        for fun_def in &self.fun_deps {
543            writeln!(file, "{}", self.import(&self.lean_file_for_fun(fun_def)))?;
544        }
545
546        Ok(())
547    }
548
549    fn generate_vc_file(&self) -> io::Result<()> {
550        // 1. Generate imports
551        self.generate_vc_prelude()?;
552
553        // 2. Create file and add imports
554        let def_id = self.def_id.resolved_id();
555        let path = LeanFile::Vc(def_id).path(self.genv);
556        if let Some(mut file) = create_file_with_dirs(path)? {
557            self.generate_vc_imports(&mut file)?;
558
559            let vc_name = vc_name(self.genv, def_id);
560            // 3. Write the VC
561            write!(
562                file,
563                "{}",
564                WithLeanCtxt {
565                    item: lean_format::LeanKConstraint {
566                        theorem_name: &vc_name,
567                        kvars: &self.kvar_decls,
568                        constr: &self.constraint,
569                    },
570                    cx: &self.lean_cx()
571                }
572            )?;
573            file.sync_all()?;
574        }
575
576        Ok(())
577    }
578
579    fn generate_proof_if_absent(&self) -> io::Result<()> {
580        let def_id = self.def_id.resolved_id();
581        let vc_name = vc_name(self.genv, def_id);
582        let proof_name = format!("{vc_name}_proof");
583        let path = LeanFile::Proof(def_id).path(self.genv);
584
585        if let Some(span) = self.genv.proven_externally(self.def_id.local_id()) {
586            let dst_span = SpanTrace::from_path(&path, 3, 5, proof_name.len());
587            dbg::hyperlink_json!(self.genv.tcx(), span, dst_span);
588        }
589
590        if let Some(mut file) = create_file_with_dirs(path)? {
591            writeln!(file, "{}", self.import(&LeanFile::Fluxlib))?;
592            writeln!(file, "{}", self.import(&LeanFile::Vc(def_id)))?;
593            writeln!(file, "def {proof_name} : {vc_name} := by")?;
594            writeln!(file, "  unfold {vc_name}")?;
595            writeln!(file, "  sorry")?;
596            file.sync_all()?;
597        }
598        Ok(())
599    }
600
601    fn record_proof(&self) -> io::Result<()> {
602        let path = LeanFile::Basic.path(self.genv);
603
604        let mut file = match create_file_with_dirs(&path)? {
605            Some(mut file) => {
606                // First invocation: reset VCs
607                writeln!(file, "-- Flux Basic Imports [DO NOT MODIFY] --")?;
608                file
609            }
610            None => fs::OpenOptions::new().append(true).open(path)?,
611        };
612        writeln!(file, "{}", self.import(&LeanFile::Proof(self.def_id.resolved_id())))
613    }
614
615    pub fn encode(
616        genv: GlobalEnv<'genv, 'tcx>,
617        def_id: MaybeExternId,
618        pretty_var_map: PrettyMap<fixpoint::LocalVar>,
619        sort_deps: SortDeps,
620        fun_deps: Vec<fixpoint::FunDef>,
621        kvar_decls: Vec<fixpoint::KVarDecl>,
622        constraint: fixpoint::Constraint,
623        kvar_solutions: KVarSolutions,
624    ) -> io::Result<()> {
625        let encoder = Self::new(
626            genv,
627            def_id,
628            pretty_var_map,
629            sort_deps,
630            fun_deps,
631            kvar_solutions,
632            kvar_decls,
633            constraint,
634        )?;
635        encoder.run()?;
636        Ok(())
637    }
638}
flux_infer/lean_encoding.rs

flux_infer/
lean_encoding.rs
