ext/
save.rs

1use std::{
2    collections::HashSet,
3    fs::File,
4    io,
5    path::{Path, PathBuf},
6    sync::{Arc, LazyLock, Mutex},
7};
8
9use algebra::Algebra;
10use anyhow::Context;
11use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
12use sseq::coordinates::Bidegree;
13
14#[derive(Debug, Clone, Eq, PartialEq)]
15pub enum SaveDirectory {
16    None,
17    Combined(AbsolutePath),
18    Split {
19        read: AbsolutePath,
20        write: AbsolutePath,
21    },
22}
23
24impl SaveDirectory {
25    pub fn read(&self) -> Option<&PathBuf> {
26        match self {
27            Self::None => None,
28            Self::Combined(x) => Some(&x.0),
29            Self::Split { read, .. } => Some(&read.0),
30        }
31    }
32
33    pub fn write(&self) -> Option<&PathBuf> {
34        match self {
35            Self::None => None,
36            Self::Combined(x) => Some(&x.0),
37            Self::Split { write, .. } => Some(&write.0),
38        }
39    }
40
41    pub fn push<P: AsRef<Path>>(&mut self, p: P) {
42        match self {
43            Self::None => {}
44            Self::Combined(d) => {
45                d.push(p);
46            }
47            Self::Split { read, write } => {
48                read.push(&p);
49                write.push(p);
50            }
51        }
52    }
53
54    pub fn is_none(&self) -> bool {
55        matches!(self, Self::None)
56    }
57
58    pub fn is_some(&self) -> bool {
59        !self.is_none()
60    }
61}
62
63impl From<Option<PathBuf>> for SaveDirectory {
64    fn from(x: Option<PathBuf>) -> Self {
65        match x {
66            None => Self::None,
67            Some(x) => Self::Combined(AbsolutePath::new(x)),
68        }
69    }
70}
71
72/// A `DashSet<PathBuf>` of files that are currently opened and being written to. When calling this
73/// function for the first time, we set the ctrlc handler to delete currently opened files then
74/// exit.
75fn open_files() -> &'static Mutex<HashSet<PathBuf>> {
76    static OPEN_FILES: LazyLock<Mutex<HashSet<PathBuf>>> = LazyLock::new(|| {
77        #[cfg(unix)]
78        ctrlc::set_handler(move || {
79            tracing::warn!("Ctrl-C detected. Deleting open files and exiting.");
80            let files = open_files().lock().unwrap();
81            for file in &*files {
82                std::fs::remove_file(file)
83                    .unwrap_or_else(|_| panic!("Error when deleting {file:?}"));
84                tracing::warn!(?file, "deleted");
85            }
86            std::process::exit(130);
87        })
88        .expect("Error setting Ctrl-C handler");
89        Default::default()
90    });
91    &OPEN_FILES
92}
93
94#[derive(Debug, Copy, Clone, Eq, PartialEq)]
95#[non_exhaustive]
96pub enum SaveKind {
97    /// The kernel of a resolution differential
98    Kernel,
99
100    /// The differential and augmentation map in a resolution
101    Differential,
102
103    /// The quasi-inverse of the resolution differential
104    ResQi,
105
106    /// The quasi-inverse of the augmentation map
107    AugmentationQi,
108
109    /// Secondary composite
110    SecondaryComposite,
111
112    /// Intermediate data used by secondary code
113    SecondaryIntermediate,
114
115    /// A secondary homotopy
116    SecondaryHomotopy,
117
118    /// A chain map
119    ChainMap,
120
121    /// A chain homotopy
122    ChainHomotopy,
123
124    /// The differential with Nassau's algorithm. This does not store the chain map data because we
125    /// always only resolve the sphere
126    NassauDifferential,
127
128    /// The quasi-inverse data in Nassau's algorithm
129    NassauQi,
130}
131
132impl SaveKind {
133    pub fn magic(self) -> u32 {
134        match self {
135            Self::Kernel => 0x0000D1FF,
136            Self::Differential => 0xD1FF0000,
137            Self::ResQi => 0x0100D1FF,
138            Self::AugmentationQi => 0x0100A000,
139            Self::SecondaryComposite => 0x00020000,
140            Self::SecondaryIntermediate => 0x00020001,
141            Self::SecondaryHomotopy => 0x00020002,
142            Self::ChainMap => 0x10100000,
143            Self::ChainHomotopy => 0x11110000,
144            Self::NassauDifferential => 0xD1FF0001,
145            Self::NassauQi => 0x0100D1FE,
146        }
147    }
148
149    pub fn name(self) -> &'static str {
150        match self {
151            Self::Kernel => "kernel",
152            Self::Differential => "differential",
153            Self::ResQi => "res_qi",
154            Self::AugmentationQi => "augmentation_qi",
155            Self::SecondaryComposite => "secondary_composite",
156            Self::SecondaryIntermediate => "secondary_intermediate",
157            Self::SecondaryHomotopy => "secondary_homotopy",
158            Self::ChainMap => "chain_map",
159            Self::ChainHomotopy => "chain_homotopy",
160            Self::NassauDifferential => "nassau_differential",
161            Self::NassauQi => "nassau_qi",
162        }
163    }
164
165    pub fn resolution_data() -> impl Iterator<Item = Self> {
166        use SaveKind::*;
167        static KINDS: [SaveKind; 4] = [Kernel, Differential, ResQi, AugmentationQi];
168        KINDS.iter().copied()
169    }
170
171    pub fn nassau_data() -> impl Iterator<Item = Self> {
172        use SaveKind::*;
173        static KINDS: [SaveKind; 2] = [NassauDifferential, NassauQi];
174        KINDS.iter().copied()
175    }
176
177    pub fn secondary_data() -> impl Iterator<Item = Self> {
178        use SaveKind::*;
179        static KINDS: [SaveKind; 3] =
180            [SecondaryComposite, SecondaryIntermediate, SecondaryHomotopy];
181        KINDS.iter().copied()
182    }
183
184    pub fn create_dir(self, p: &std::path::Path) -> anyhow::Result<()> {
185        let mut p = p.to_owned();
186
187        p.push(format!("{}s", self.name()));
188        if !p.exists() {
189            std::fs::create_dir_all(&p)
190                .with_context(|| format!("Failed to create directory {p:?}"))?;
191        } else if !p.is_dir() {
192            return Err(anyhow::anyhow!("{p:?} is not a directory"));
193        }
194        Ok(())
195    }
196}
197
198/// In addition to checking the checksum, we also keep track of which files are open, and we delete
199/// the open files if the program is terminated halfway.
200pub struct ChecksumWriter<T: io::Write> {
201    writer: T,
202    path: PathBuf,
203    adler: adler::Adler32,
204}
205
206impl<T: io::Write> ChecksumWriter<T> {
207    pub fn new(path: PathBuf, writer: T) -> Self {
208        Self {
209            path,
210            writer,
211            adler: adler::Adler32::new(),
212        }
213    }
214}
215
216/// We only implement the functions required and the ones we actually use.
217impl<T: io::Write> io::Write for ChecksumWriter<T> {
218    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
219        let bytes_written = self.writer.write(buf)?;
220        self.adler.write_slice(&buf[0..bytes_written]);
221        Ok(bytes_written)
222    }
223
224    fn flush(&mut self) -> io::Result<()> {
225        self.writer.flush()
226    }
227
228    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
229        self.writer.write_all(buf)?;
230        self.adler.write_slice(buf);
231        Ok(())
232    }
233}
234
235impl<T: io::Write> std::ops::Drop for ChecksumWriter<T> {
236    fn drop(&mut self) {
237        if !std::thread::panicking() {
238            // We may not have finished writing, so the data is wrong. It should not be given a
239            // valid checksum
240            self.writer
241                .write_u32::<LittleEndian>(self.adler.checksum())
242                .unwrap();
243            self.writer.flush().unwrap();
244            assert!(
245                open_files().lock().unwrap().remove(&self.path),
246                "File {:?} already dropped",
247                self.path
248            );
249        }
250        tracing::info!(file = ?self.path, "closing");
251    }
252}
253
254pub struct ChecksumReader<T: io::Read> {
255    reader: T,
256    adler: adler::Adler32,
257}
258
259impl<T: io::Read> ChecksumReader<T> {
260    pub fn new(reader: T) -> Self {
261        Self {
262            reader,
263            adler: adler::Adler32::new(),
264        }
265    }
266}
267
268/// We only implement the functions required and the ones we actually use.
269impl<T: io::Read> io::Read for ChecksumReader<T> {
270    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
271        let bytes_read = self.reader.read(buf)?;
272        self.adler.write_slice(&buf[0..bytes_read]);
273        Ok(bytes_read)
274    }
275
276    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
277        self.reader.read_exact(buf)?;
278        self.adler.write_slice(buf);
279        Ok(())
280    }
281}
282
283impl<T: io::Read> std::ops::Drop for ChecksumReader<T> {
284    fn drop(&mut self) {
285        if !std::thread::panicking() {
286            // If we are panicking, we may not have read everything, and panic in panic
287            // is bad.
288            assert_eq!(
289                self.adler.checksum(),
290                self.reader.read_u32::<LittleEndian>().unwrap(),
291                "Invalid file checksum"
292            );
293            let mut buf = [0];
294            // Check EOF
295            assert_eq!(self.reader.read(&mut buf).unwrap(), 0, "EOF not reached");
296        }
297    }
298}
299
300/// Open the file pointed to by `path` as a `Box<dyn Read>`. If the file does not exist, look for
301/// compressed versions.
302fn open_file(path: PathBuf) -> Option<Box<dyn io::Read>> {
303    use io::BufRead;
304
305    // We should try in decreasing order of access speed.
306    match File::open(&path) {
307        Ok(f) => {
308            let mut reader = io::BufReader::new(f);
309            if reader
310                .fill_buf()
311                .unwrap_or_else(|e| panic!("Error when reading from {path:?}: {e}"))
312                .is_empty()
313            {
314                // The file is empty. Delete the file and proceed as if it didn't exist
315                std::fs::remove_file(&path)
316                    .unwrap_or_else(|e| panic!("Error when deleting empty file {path:?}: {e}"));
317                return None;
318            }
319            return Some(Box::new(ChecksumReader::new(reader)));
320        }
321        Err(e) => {
322            if e.kind() != io::ErrorKind::NotFound {
323                panic!("Error when opening {path:?}: {e}");
324            }
325        }
326    }
327
328    #[cfg(feature = "zstd")]
329    {
330        let mut path = path;
331        path.set_extension("zst");
332        match File::open(&path) {
333            Ok(f) => {
334                return Some(Box::new(ChecksumReader::new(
335                    zstd::stream::Decoder::new(f).unwrap(),
336                )));
337            }
338            Err(e) => {
339                if e.kind() != io::ErrorKind::NotFound {
340                    panic!("Error when opening {path:?}");
341                }
342            }
343        }
344    }
345
346    None
347}
348
349pub struct SaveFile<A: Algebra> {
350    pub kind: SaveKind,
351    pub algebra: Arc<A>,
352    pub b: Bidegree,
353    pub idx: Option<usize>,
354}
355
356impl<A: Algebra> SaveFile<A> {
357    fn write_header(&self, buffer: &mut impl io::Write) -> io::Result<()> {
358        buffer.write_u32::<LittleEndian>(self.kind.magic())?;
359        buffer.write_u32::<LittleEndian>(self.algebra.magic())?;
360        buffer.write_i32::<LittleEndian>(self.b.s())?;
361        buffer.write_i32::<LittleEndian>(if let Some(i) = self.idx {
362            self.b.t() + ((i as i32) << 16)
363        } else {
364            self.b.t()
365        })
366    }
367
368    fn validate_header(&self, buffer: &mut impl io::Read) -> io::Result<()> {
369        macro_rules! check_header {
370            ($name:literal, $value:expr, $format:literal) => {
371                let data = buffer.read_u32::<LittleEndian>()?;
372                if data != $value {
373                    return Err(io::Error::new(
374                        io::ErrorKind::InvalidData,
375                        format!(
376                            "Invalid header: {} was {} but expected {}",
377                            $name,
378                            format_args!($format, data),
379                            format_args!($format, $value)
380                        ),
381                    ));
382                }
383            };
384        }
385
386        check_header!("magic", self.kind.magic(), "{:#010x}");
387        check_header!("algebra", self.algebra.magic(), "{:#06x}");
388        check_header!("s", self.b.s() as u32, "{}");
389        check_header!(
390            "t",
391            if let Some(i) = self.idx {
392                self.b.t() as u32 + ((i as u32) << 16)
393            } else {
394                self.b.t() as u32
395            },
396            "{}"
397        );
398
399        Ok(())
400    }
401
402    /// This panics if there is no save dir
403    fn get_save_path(&self, mut dir: PathBuf) -> PathBuf {
404        if let Some(idx) = self.idx {
405            dir.push(format!(
406                "{name}s/{s}_{t}_{idx}_{name}",
407                name = self.kind.name(),
408                s = self.b.s(),
409                t = self.b.t()
410            ));
411        } else {
412            dir.push(format!(
413                "{name}s/{s}_{t}_{name}",
414                name = self.kind.name(),
415                s = self.b.s(),
416                t = self.b.t()
417            ));
418        }
419        dir
420    }
421
422    pub fn open_file(&self, dir: PathBuf) -> Option<Box<dyn io::Read>> {
423        let file_path = self.get_save_path(dir);
424        let path_string = file_path.to_string_lossy().into_owned();
425        if let Some(mut f) = open_file(file_path) {
426            self.validate_header(&mut f).unwrap();
427            tracing::info!(file = path_string, "success open for reading");
428            Some(f)
429        } else {
430            tracing::info!(file = path_string, "failed open for reading");
431            None
432        }
433    }
434
435    pub fn exists(&self, dir: PathBuf) -> bool {
436        let path = self.get_save_path(dir);
437        if path.exists() {
438            return true;
439        }
440        #[cfg(not(target_arch = "wasm32"))]
441        {
442            let mut path = path;
443            path.set_extension("zst");
444            if path.exists() {
445                return true;
446            }
447        }
448        false
449    }
450
451    pub fn delete_file(&self, dir: PathBuf) -> io::Result<()> {
452        let p = self.get_save_path(dir);
453        match std::fs::remove_file(p) {
454            Ok(()) => Ok(()),
455            Err(e) if e.kind() == io::ErrorKind::NotFound => Ok(()),
456            Err(e) => Err(e),
457        }
458    }
459
460    /// # Arguments
461    ///  - `overwrite`: Whether to overwrite a file if it already exists.
462    pub fn create_file(&self, dir: PathBuf, overwrite: bool) -> impl io::Write + use<A> {
463        let p = self.get_save_path(dir);
464        tracing::info!(file = ?p, "open for writing");
465
466        // We need to do this before creating any file. The ctrlc handler does not block other threads
467        // from running, but it does lock [`open_files()`]. So this ensures we do not open new files
468        // while handling ctrlc.
469        assert!(
470            open_files().lock().unwrap().insert(p.clone()),
471            "File {p:?} is already opened"
472        );
473
474        let f = std::fs::OpenOptions::new()
475            .write(true)
476            .create_new(!overwrite)
477            .create(true)
478            .truncate(true)
479            .open(&p)
480            .with_context(|| format!("Failed to create save file {p:?}"))
481            .unwrap();
482        let mut f = ChecksumWriter::new(p, io::BufWriter::new(f));
483        self.write_header(&mut f).unwrap();
484        f
485    }
486}
487
488#[derive(Debug, Clone, Eq, PartialEq)]
489pub struct AbsolutePath(PathBuf);
490
491impl AbsolutePath {
492    pub fn new<P: AsRef<Path>>(path: P) -> Self {
493        let p = std::path::absolute(path)
494            .unwrap_or_else(|e| panic!("Error when getting absolute path: {e}"));
495        Self(p)
496    }
497
498    pub fn push<P: AsRef<Path>>(&mut self, p: P) {
499        self.0.push(p);
500    }
501}
502
503impl<P: AsRef<Path>> From<P> for AbsolutePath {
504    fn from(p: P) -> Self {
505        Self::new(p)
506    }
507}