use serde::{de::DeserializeOwned, Serialize};
#[cfg(feature = "checkpoint")]
use bincode::Options;
#[cfg(feature = "checkpoint")]
use std::fs::{self, File};
#[cfg(feature = "checkpoint")]
use std::io::{BufReader, BufWriter, Read, Write};
#[cfg(feature = "checkpoint")]
use std::path::{Path, PathBuf};
use super::state::{Checkpoint, CHECKPOINT_VERSION};
use crate::error::CheckpointError;
pub const DEFAULT_MAX_CHECKPOINT_BYTES: u64 = 256 * 1024 * 1024;
pub const MIN_SUPPORTED_CHECKPOINT_VERSION: u32 = 1;
#[cfg(feature = "checkpoint")]
fn check_version(version: u32) -> Result<(), CheckpointError> {
if version > CHECKPOINT_VERSION {
return Err(CheckpointError::VersionMismatch {
expected: CHECKPOINT_VERSION,
found: version,
});
}
if version < MIN_SUPPORTED_CHECKPOINT_VERSION {
return Err(CheckpointError::VersionTooOld(version));
}
Ok(())
}
#[cfg(feature = "checkpoint")]
fn bincode_read_options(limit: u64) -> impl Options {
bincode::DefaultOptions::new()
.with_fixint_encoding()
.allow_trailing_bytes()
.with_limit(limit)
}
#[cfg(feature = "checkpoint")]
fn temp_path_for(path: &Path) -> PathBuf {
let mut os = path.as_os_str().to_owned();
os.push(".tmp");
PathBuf::from(os)
}
#[cfg(feature = "checkpoint")]
fn parse_checkpoint_index(file_name: &str, base_name: &str) -> Option<usize> {
let rest = file_name.strip_prefix(base_name)?.strip_prefix('_')?;
let digits: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect();
if digits.is_empty() {
return None;
}
let tail = &rest[digits.len()..];
if !tail.is_empty() && !tail.starts_with('.') {
return None;
}
digits.parse::<usize>().ok()
}
#[cfg(feature = "checkpoint")]
fn scan_max_index(directory: &Path, base_name: &str) -> Option<usize> {
std::fs::read_dir(directory)
.ok()?
.filter_map(|e| e.ok())
.filter_map(|e| {
let name = e.file_name().to_string_lossy().into_owned();
parse_checkpoint_index(&name, base_name)
})
.max()
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CheckpointFormat {
Json,
Binary,
CompressedBinary,
}
impl Default for CheckpointFormat {
fn default() -> Self {
Self::Binary
}
}
#[cfg(feature = "checkpoint")]
pub fn save_checkpoint<G>(
checkpoint: &Checkpoint<G>,
path: impl AsRef<Path>,
format: CheckpointFormat,
) -> Result<(), CheckpointError>
where
G: Clone + Serialize + crate::genome::traits::EvolutionaryGenome,
{
let path = path.as_ref();
let tmp_path = temp_path_for(path);
let write_result = (|| -> Result<(), CheckpointError> {
let file = File::create(&tmp_path)?;
let mut writer = BufWriter::new(file);
match format {
CheckpointFormat::Json => {
serde_json::to_writer_pretty(&mut writer, checkpoint)
.map_err(|e| CheckpointError::SerializeError(Box::new(e)))?;
}
CheckpointFormat::Binary => {
writer.write_all(&CHECKPOINT_VERSION.to_le_bytes())?;
writer.write_all(b"FEVO")?;
bincode::serialize_into(&mut writer, checkpoint)
.map_err(|e| CheckpointError::SerializeError(Box::new(e)))?;
}
CheckpointFormat::CompressedBinary => {
writer.write_all(&CHECKPOINT_VERSION.to_le_bytes())?;
writer.write_all(b"FEVC")?; let bytes = bincode::serialize(checkpoint)
.map_err(|e| CheckpointError::SerializeError(Box::new(e)))?;
let compressed = compress_data(&bytes);
writer.write_all(&(compressed.len() as u64).to_le_bytes())?;
writer.write_all(&compressed)?;
}
}
writer.flush()?;
let file = writer.into_inner().map_err(|e| e.into_error())?;
file.sync_all()?;
Ok(())
})();
if let Err(e) = write_result {
let _ = fs::remove_file(&tmp_path);
return Err(e);
}
fs::rename(&tmp_path, path)?;
Ok(())
}
#[cfg(feature = "checkpoint")]
pub fn load_checkpoint<G>(path: impl AsRef<Path>) -> Result<Checkpoint<G>, CheckpointError>
where
G: Clone + Serialize + DeserializeOwned + crate::genome::traits::EvolutionaryGenome,
{
load_checkpoint_with_limit(path, DEFAULT_MAX_CHECKPOINT_BYTES)
}
#[cfg(feature = "checkpoint")]
pub fn load_checkpoint_with_limit<G>(
path: impl AsRef<Path>,
max_bytes: u64,
) -> Result<Checkpoint<G>, CheckpointError>
where
G: Clone + Serialize + DeserializeOwned + crate::genome::traits::EvolutionaryGenome,
{
let path = path.as_ref();
if !path.exists() {
return Err(CheckpointError::NotFound(path.display().to_string()));
}
let file_len = fs::metadata(path)?.len();
if file_len > max_bytes {
return Err(CheckpointError::TooLarge {
size: file_len,
limit: max_bytes,
});
}
let file = File::open(path)?;
let mut reader = BufReader::new(file);
let mut header = [0u8; 8];
reader.read_exact(&mut header)?;
if &header[4..8] == b"FEVO" {
let version = u32::from_le_bytes([header[0], header[1], header[2], header[3]]);
check_version(version)?;
let checkpoint: Checkpoint<G> = bincode_read_options(max_bytes)
.deserialize_from(&mut reader)
.map_err(|e| CheckpointError::DeserializeError(e))?;
check_version(checkpoint.version)?;
Ok(checkpoint)
} else if &header[4..8] == b"FEVC" {
let version = u32::from_le_bytes([header[0], header[1], header[2], header[3]]);
check_version(version)?;
let mut len_bytes = [0u8; 8];
reader.read_exact(&mut len_bytes)?;
let compressed_len = u64::from_le_bytes(len_bytes);
if compressed_len > max_bytes {
return Err(CheckpointError::TooLarge {
size: compressed_len,
limit: max_bytes,
});
}
let mut compressed = vec![0u8; compressed_len as usize];
reader.read_exact(&mut compressed)?;
let decompressed = decompress_data(&compressed).map_err(CheckpointError::Corrupted)?;
let checkpoint: Checkpoint<G> = bincode_read_options(max_bytes)
.deserialize(&decompressed)
.map_err(|e| CheckpointError::DeserializeError(e))?;
check_version(checkpoint.version)?;
Ok(checkpoint)
} else {
drop(reader);
let file = File::open(path)?;
let reader = BufReader::new(file);
let checkpoint: Checkpoint<G> = serde_json::from_reader(reader)
.map_err(|e| CheckpointError::DeserializeError(Box::new(e)))?;
check_version(checkpoint.version)?;
Ok(checkpoint)
}
}
#[cfg(feature = "checkpoint")]
fn compress_data(data: &[u8]) -> Vec<u8> {
if data.is_empty() {
return Vec::new();
}
let mut compressed = Vec::with_capacity(data.len());
let mut i = 0;
while i < data.len() {
let byte = data[i];
let mut count = 1u8;
while i + (count as usize) < data.len() && data[i + (count as usize)] == byte && count < 255
{
count += 1;
}
if count >= 4 || byte == 0xFF {
compressed.push(0xFF);
compressed.push(count);
compressed.push(byte);
} else {
for _ in 0..count {
if byte == 0xFF {
compressed.push(0xFF);
compressed.push(1);
compressed.push(0xFF);
} else {
compressed.push(byte);
}
}
}
i += count as usize;
}
compressed
}
#[cfg(feature = "checkpoint")]
fn decompress_data(data: &[u8]) -> Result<Vec<u8>, String> {
let mut decompressed = Vec::new();
let mut i = 0;
while i < data.len() {
if data[i] == 0xFF {
if i + 2 >= data.len() {
return Err("Truncated RLE sequence".to_string());
}
let count = data[i + 1] as usize;
let byte = data[i + 2];
for _ in 0..count {
decompressed.push(byte);
}
i += 3;
} else {
decompressed.push(data[i]);
i += 1;
}
}
Ok(decompressed)
}
#[cfg(feature = "checkpoint")]
pub struct CheckpointManager {
pub directory: std::path::PathBuf,
pub base_name: String,
pub format: CheckpointFormat,
pub keep_n: usize,
pub interval: usize,
pub max_bytes: u64,
current_index: usize,
}
#[cfg(feature = "checkpoint")]
impl CheckpointManager {
pub fn new(directory: impl Into<std::path::PathBuf>, base_name: impl Into<String>) -> Self {
let directory = directory.into();
let base_name = base_name.into();
let current_index = scan_max_index(&directory, &base_name)
.map(|max| max + 1)
.unwrap_or(0);
Self {
directory,
base_name,
format: CheckpointFormat::Binary,
keep_n: 3,
interval: 100,
max_bytes: DEFAULT_MAX_CHECKPOINT_BYTES,
current_index,
}
}
pub fn with_format(mut self, format: CheckpointFormat) -> Self {
self.format = format;
self
}
pub fn with_max_bytes(mut self, max_bytes: u64) -> Self {
self.max_bytes = max_bytes;
self
}
pub fn current_index(&self) -> usize {
self.current_index
}
pub fn keep(mut self, n: usize) -> Self {
self.keep_n = n;
self
}
pub fn every(mut self, generations: usize) -> Self {
self.interval = generations;
self
}
pub fn should_save(&self, generation: usize) -> bool {
generation > 0 && generation.is_multiple_of(self.interval)
}
pub fn current_path(&self) -> std::path::PathBuf {
let extension = match self.format {
CheckpointFormat::Json => "json",
CheckpointFormat::Binary | CheckpointFormat::CompressedBinary => "ckpt",
};
self.directory.join(format!(
"{}_{:08}.{}",
self.base_name, self.current_index, extension
))
}
pub fn save<G>(&mut self, checkpoint: &Checkpoint<G>) -> Result<(), CheckpointError>
where
G: Clone + Serialize + crate::genome::traits::EvolutionaryGenome,
{
std::fs::create_dir_all(&self.directory)?;
let path = self.current_path();
save_checkpoint(checkpoint, &path, self.format)?;
self.current_index += 1;
if self.current_index > self.keep_n {
let old_index = self.current_index - self.keep_n - 1;
let extension = match self.format {
CheckpointFormat::Json => "json",
CheckpointFormat::Binary | CheckpointFormat::CompressedBinary => "ckpt",
};
let old_path = self
.directory
.join(format!("{}_{:08}.{}", self.base_name, old_index, extension));
let _ = std::fs::remove_file(old_path); }
Ok(())
}
pub fn load_latest<G>(&self) -> Result<Option<Checkpoint<G>>, CheckpointError>
where
G: Clone + Serialize + DeserializeOwned + crate::genome::traits::EvolutionaryGenome,
{
let extension = match self.format {
CheckpointFormat::Json => "json",
CheckpointFormat::Binary | CheckpointFormat::CompressedBinary => "ckpt",
};
let _pattern = format!("{}_*.{}", self.base_name, extension);
let mut checkpoints: Vec<_> = std::fs::read_dir(&self.directory)?
.filter_map(|e| e.ok())
.filter(|e| e.file_name().to_string_lossy().starts_with(&self.base_name))
.collect();
if checkpoints.is_empty() {
return Ok(None);
}
checkpoints.sort_by(|a, b| {
let ia = parse_checkpoint_index(&a.file_name().to_string_lossy(), &self.base_name);
let ib = parse_checkpoint_index(&b.file_name().to_string_lossy(), &self.base_name);
ib.cmp(&ia)
});
for entry in checkpoints {
match load_checkpoint_with_limit(entry.path(), self.max_bytes) {
Ok(checkpoint) => return Ok(Some(checkpoint)),
Err(_) => continue, }
}
Ok(None)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::genome::real_vector::RealVector;
use crate::population::individual::Individual;
use tempfile::tempdir;
#[test]
fn test_save_load_json() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.json");
let population: Vec<Individual<RealVector>> = vec![
Individual::new(RealVector::new(vec![1.0, 2.0])),
Individual::new(RealVector::new(vec![3.0, 4.0])),
];
let checkpoint = Checkpoint::new(10, population);
save_checkpoint(&checkpoint, &path, CheckpointFormat::Json).unwrap();
let loaded: Checkpoint<RealVector> = load_checkpoint(&path).unwrap();
assert_eq!(loaded.generation, 10);
assert_eq!(loaded.population.len(), 2);
}
#[test]
fn test_save_load_binary() {
let dir = tempdir().unwrap();
let path = dir.path().join("test.ckpt");
let population: Vec<Individual<RealVector>> =
vec![Individual::new(RealVector::new(vec![1.0, 2.0, 3.0]))];
let checkpoint = Checkpoint::new(5, population)
.with_evaluations(500)
.with_metadata("test", "value");
save_checkpoint(&checkpoint, &path, CheckpointFormat::Binary).unwrap();
let loaded: Checkpoint<RealVector> = load_checkpoint(&path).unwrap();
assert_eq!(loaded.generation, 5);
assert_eq!(loaded.evaluations, 500);
assert_eq!(loaded.metadata.get("test"), Some(&"value".to_string()));
}
#[test]
fn test_save_load_compressed() {
let dir = tempdir().unwrap();
let path = dir.path().join("test_compressed.ckpt");
let population: Vec<Individual<RealVector>> = (0..100)
.map(|i| Individual::new(RealVector::new(vec![i as f64; 10])))
.collect();
let checkpoint = Checkpoint::new(100, population);
save_checkpoint(&checkpoint, &path, CheckpointFormat::CompressedBinary).unwrap();
let loaded: Checkpoint<RealVector> = load_checkpoint(&path).unwrap();
assert_eq!(loaded.generation, 100);
assert_eq!(loaded.population.len(), 100);
}
#[test]
fn test_compression_decompression() {
let original = vec![0u8, 0, 0, 0, 0, 1, 2, 3, 3, 3, 3, 3, 3, 4, 5];
let compressed = compress_data(&original);
let decompressed = decompress_data(&compressed).unwrap();
assert_eq!(original, decompressed);
}
#[test]
fn test_checkpoint_manager() {
let dir = tempdir().unwrap();
let mut manager = CheckpointManager::new(dir.path(), "evolution")
.with_format(CheckpointFormat::Binary)
.keep(2)
.every(10);
for gen in [10, 20, 30, 40] {
let population: Vec<Individual<RealVector>> =
vec![Individual::new(RealVector::new(vec![gen as f64]))];
let checkpoint = Checkpoint::new(gen, population);
manager.save(&checkpoint).unwrap();
}
let loaded: Option<Checkpoint<RealVector>> = manager.load_latest().unwrap();
assert!(loaded.is_some());
assert_eq!(loaded.unwrap().generation, 40);
}
#[test]
fn test_version_check() {
let population: Vec<Individual<RealVector>> = vec![];
let checkpoint = Checkpoint::new(0, population);
assert!(checkpoint.is_compatible());
}
#[test]
fn test_json_version_gate_too_new() {
let dir = tempdir().unwrap();
let path = dir.path().join("future.json");
let population: Vec<Individual<RealVector>> = vec![];
let mut checkpoint = Checkpoint::new(3, population);
checkpoint.version = CHECKPOINT_VERSION + 1;
save_checkpoint(&checkpoint, &path, CheckpointFormat::Json).unwrap();
let err = load_checkpoint::<RealVector>(&path).unwrap_err();
assert!(
matches!(err, CheckpointError::VersionMismatch { .. }),
"expected VersionMismatch, got {err:?}"
);
}
#[test]
fn test_json_version_gate_too_old() {
let dir = tempdir().unwrap();
let path = dir.path().join("ancient.json");
let population: Vec<Individual<RealVector>> = vec![];
let mut checkpoint = Checkpoint::new(3, population);
checkpoint.version = MIN_SUPPORTED_CHECKPOINT_VERSION - 1;
save_checkpoint(&checkpoint, &path, CheckpointFormat::Json).unwrap();
let err = load_checkpoint::<RealVector>(&path).unwrap_err();
assert!(
matches!(err, CheckpointError::VersionTooOld(v) if v == MIN_SUPPORTED_CHECKPOINT_VERSION - 1),
"expected VersionTooOld, got {err:?}"
);
}
#[test]
fn test_manager_is_restart_safe() {
let dir = tempdir().unwrap();
{
let mut manager = CheckpointManager::new(dir.path(), "evolution")
.with_format(CheckpointFormat::Binary)
.keep(10);
for gen in [10usize, 20, 30] {
let population: Vec<Individual<RealVector>> =
vec![Individual::new(RealVector::new(vec![gen as f64]))];
manager.save(&Checkpoint::new(gen, population)).unwrap();
}
assert_eq!(manager.current_index(), 3);
}
let manager2 = CheckpointManager::new(dir.path(), "evolution")
.with_format(CheckpointFormat::Binary)
.keep(10);
assert_eq!(
manager2.current_index(),
3,
"restarted manager must continue after the highest existing index"
);
let loaded: Option<Checkpoint<RealVector>> = manager2.load_latest().unwrap();
assert_eq!(loaded.unwrap().generation, 30);
}
#[test]
fn test_atomic_save_preserves_destination_on_failure() {
let dir = tempdir().unwrap();
let path = dir.path().join("evolution.ckpt");
let good: Vec<Individual<RealVector>> = vec![Individual::new(RealVector::new(vec![1.0]))];
save_checkpoint(&Checkpoint::new(111, good), &path, CheckpointFormat::Binary).unwrap();
let mut tmp = path.clone().into_os_string();
tmp.push(".tmp");
std::fs::create_dir(&tmp).unwrap();
let newer: Vec<Individual<RealVector>> = vec![Individual::new(RealVector::new(vec![2.0]))];
let result = save_checkpoint(
&Checkpoint::new(222, newer),
&path,
CheckpointFormat::Binary,
);
assert!(
result.is_err(),
"save should fail when temp file is blocked"
);
let loaded: Checkpoint<RealVector> = load_checkpoint(&path).unwrap();
assert_eq!(loaded.generation, 111);
std::fs::remove_dir(&tmp).unwrap();
}
#[test]
fn test_atomic_save_leaves_no_temp_file() {
let dir = tempdir().unwrap();
let path = dir.path().join("evolution.ckpt");
let population: Vec<Individual<RealVector>> =
vec![Individual::new(RealVector::new(vec![1.0]))];
save_checkpoint(
&Checkpoint::new(5, population),
&path,
CheckpointFormat::Binary,
)
.unwrap();
let mut tmp = path.clone().into_os_string();
tmp.push(".tmp");
assert!(!PathBuf::from(tmp).exists(), "temp file must not remain");
}
#[test]
fn test_load_rejects_oversized_file() {
let dir = tempdir().unwrap();
let path = dir.path().join("big.ckpt");
let population: Vec<Individual<RealVector>> = (0..50)
.map(|i| Individual::new(RealVector::new(vec![i as f64; 8])))
.collect();
save_checkpoint(
&Checkpoint::new(1, population),
&path,
CheckpointFormat::Binary,
)
.unwrap();
let err = load_checkpoint_with_limit::<RealVector>(&path, 16).unwrap_err();
assert!(
matches!(err, CheckpointError::TooLarge { limit: 16, .. }),
"expected TooLarge, got {err:?}"
);
}
#[test]
fn test_load_rejects_corrupt_length_prefix() {
let dir = tempdir().unwrap();
let path = dir.path().join("corrupt.ckpt");
let mut bytes = Vec::new();
bytes.extend_from_slice(&CHECKPOINT_VERSION.to_le_bytes());
bytes.extend_from_slice(b"FEVC");
bytes.extend_from_slice(&u64::MAX.to_le_bytes()); std::fs::write(&path, &bytes).unwrap();
let err = load_checkpoint::<RealVector>(&path).unwrap_err();
assert!(
matches!(err, CheckpointError::TooLarge { .. }),
"expected TooLarge, got {err:?}"
);
}
#[test]
fn test_load_latest_orders_across_digit_boundary() {
let dir = tempdir().unwrap();
for (idx, gen) in [(99999usize, 99999usize), (100000, 100000)] {
let path = dir.path().join(format!("evolution_{idx}.ckpt"));
let population: Vec<Individual<RealVector>> =
vec![Individual::new(RealVector::new(vec![gen as f64]))];
save_checkpoint(
&Checkpoint::new(gen, population),
&path,
CheckpointFormat::Binary,
)
.unwrap();
}
let manager = CheckpointManager::new(dir.path(), "evolution");
let loaded: Option<Checkpoint<RealVector>> = manager.load_latest().unwrap();
assert_eq!(
loaded.unwrap().generation,
100000,
"index 100000 is newer than 99999 despite lexicographic order"
);
}
#[test]
fn test_current_path_is_zero_padded_to_8() {
let dir = tempdir().unwrap();
let manager = CheckpointManager::new(dir.path(), "evolution");
let name = manager.current_path();
assert!(
name.file_name().unwrap().to_string_lossy() == "evolution_00000000.ckpt",
"got {name:?}"
);
}
}