basin 2.0.0

//! A `serde`-gated observer that snapshots the state to disk.
//!
//! Available only with the `serde` feature and off `wasm32` (it does file
//! I/O). The companion [`read_checkpoint`] reloads a state snapshot to warm
//! start a later run. Exact continuation instead uses the solver-aware
//! [`ExactCheckpointWriter`](crate::ExactCheckpointWriter).

use std::fs;
use std::io;
use std::path::{Path, PathBuf};

use serde::Serialize;
use serde::de::DeserializeOwned;

use crate::core::observer::Observe;
use crate::core::state::State;
use crate::core::termination::TerminationReport;

const MAGIC: &[u8; 8] = b"BASINST\0";
const FORMAT_VERSION: u32 = 1;

/// Write the current state to a file with [`postcard`], overwriting the
/// previous snapshot, so the file always holds the latest checkpoint.
///
/// Files start with `BASINST\0` and a little-endian `u32` format version of 1,
/// followed by the postcard payload. Basin 2.0 reads only this format. See
/// [`read_checkpoint`] for migration from legacy 1.x files.
///
/// State-only checkpointing is an observer's job in Basin. It records an
/// iterate for a later warm start without promising an identical trajectory.
/// Register it with an [`ObserverMode`](super::ObserverMode) to pick the
/// cadence ([`Every(n)`](super::ObserverMode::Every) for every `n`th iteration,
/// [`NewBest`](super::ObserverMode::NewBest) to snapshot only on improvement).
/// The writer always snapshots on `observe_final` as well.
///
/// The state type must be [`Serialize`]. Shared progress types such as
/// [`PointState`](crate::PointState), [`FirstOrderState`](crate::FirstOrderState),
/// and [`ProposalState`](crate::ProposalState) support serialization when their
/// parameter and scalar types do. Solver-owned models, neighbors, and RNGs are
/// absent from these snapshots; use [`ExactCheckpointWriter`](crate::ExactCheckpointWriter)
/// to retain them.
/// Handing a non-serializable state is a compile error.
///
/// # Resume
///
/// Read the file with [`read_checkpoint`], deserialize into the concrete state,
/// and hand it to [`Executor::new`](crate::Executor::new). The solver runs its
/// normal `init` path and begins a new run from the restored iterate.
///
/// ```no_run
/// # use basin::{FirstOrderState, CostFunction, Executor, Gradient, GradientDescent};
/// use basin::{CheckpointWriter, ObserverMode, read_checkpoint};
/// # struct Quadratic;
/// # impl basin::ProblemConstraints for Quadratic  {
///     type Constraints = basin::NoConstraints;
///     fn constraint_set(&self) -> basin::NoConstraints { basin::NoConstraints::EMPTY }
/// }
/// impl CostFunction for Quadratic {
/// #     type Param = Vec<f64>;
/// #     type Output = f64;
/// #     type Error = std::convert::Infallible;
/// #     fn cost(&self, x: &Vec<f64>) -> Result<f64, Self::Error> {
/// #         Ok(0.5 * x.iter().map(|v| v * v).sum::<f64>())
/// #     }
/// # }
/// # impl Gradient for Quadratic {
/// #     type Gradient = Vec<f64>;
/// #     fn gradient(&self, x: &Vec<f64>) -> Result<Vec<f64>, Self::Error> { Ok(x.clone()) }
/// # }
/// // First run: checkpoint every 10 iterations.
/// Executor::new(Quadratic, GradientDescent::new(0.1), FirstOrderState::new(vec![5.0, 5.0]))
///     .max_iter(50)
///     .observe_with(CheckpointWriter::new("run.ckpt"), ObserverMode::Every(10))
///     .run()
///     .unwrap();
///
/// // Later: reload and continue from where it stopped.
/// let state: FirstOrderState<Vec<f64>> = read_checkpoint("run.ckpt").unwrap();
/// Executor::new(Quadratic, GradientDescent::new(0.1), state)
///     .max_iter(50)
///     .run()
///     .unwrap();
/// ```
#[derive(Clone, Debug)]
pub struct CheckpointWriter {
    path: PathBuf,
}

impl CheckpointWriter {
    /// Snapshot to `path`, overwriting it on each fire.
    pub fn new(path: impl Into<PathBuf>) -> Self {
        Self { path: path.into() }
    }

    /// Serialize `state` and write it to the checkpoint path. Writes to a
    /// sibling `*.tmp` first and renames into place so a crash mid-write can't
    /// truncate the previous good checkpoint.
    fn write<S: Serialize>(&self, state: &S) -> io::Result<()> {
        let payload = postcard::to_allocvec(state).map_err(io::Error::other)?;
        let mut bytes =
            Vec::with_capacity(MAGIC.len() + size_of::<u32>() + payload.len());
        bytes.extend_from_slice(MAGIC);
        bytes.extend_from_slice(&FORMAT_VERSION.to_le_bytes());
        bytes.extend_from_slice(&payload);
        let tmp = self.path.with_extension("tmp");
        fs::write(&tmp, &bytes)?;
        fs::rename(&tmp, &self.path)
    }

    /// Attempt the write; on failure, log to stderr and carry on. The
    /// [`Observe`] contract is infallible—a failing checkpoint must not kill
    /// the optimization run.
    fn try_write<S: Serialize>(&self, state: &S) {
        if let Err(err) = self.write(state) {
            eprintln!(
                "CheckpointWriter: failed to write {}: {err}",
                self.path.display()
            );
        }
    }
}

impl<S> Observe<S> for CheckpointWriter
where
    S: State + Serialize,
{
    fn observe_iter(&mut self, state: &S) {
        self.try_write(state);
    }

    fn observe_final(
        &mut self,
        state: &S,
        _reason: &TerminationReport<S::Float>,
    ) {
        self.try_write(state);
    }
}

/// Load a checkpoint previously written by [`CheckpointWriter`] into a concrete
/// state, ready for [`Executor::new`](crate::Executor::new) as a warm start.
///
/// Reads postcard state format version 1 without modifying the file. Missing
/// format markers, unsupported versions, malformed data, and trailing bytes
/// return [`io::ErrorKind::InvalidData`].
///
/// # Migrating from Basin 1.x
///
/// Basin 2.0 no longer reads unprefixed bincode files. Before upgrading, load
/// them with a compatible 1.x release and rewrite them using
/// [`CheckpointWriter`] from Basin 1.15 or later in the 1.x series, which writes
/// postcard. The concrete state's serialized layout must still match. If it
/// changes, export parameters from the 1.x application and construct a fresh
/// state in 2.0. State snapshots do not promise exact continuation.
pub fn read_checkpoint<S: DeserializeOwned>(
    path: impl AsRef<Path>,
) -> io::Result<S> {
    let bytes = fs::read(path)?;
    let invalid_data =
        |message: &str| io::Error::new(io::ErrorKind::InvalidData, message);
    let body = bytes
        .strip_prefix(MAGIC)
        .ok_or_else(|| invalid_data("invalid state checkpoint magic"))?;
    if body.len() < size_of::<u32>() {
        return Err(invalid_data("truncated state checkpoint prefix"));
    }
    let version = u32::from_le_bytes(
        body[..size_of::<u32>()]
            .try_into()
            .expect("version slice has fixed length"),
    );
    if version != FORMAT_VERSION {
        return Err(invalid_data(&format!(
            "unsupported state checkpoint format version {version}; expected {FORMAT_VERSION}"
        )));
    }
    let (state, remaining) =
        postcard::take_from_bytes(&body[size_of::<u32>()..])
            .map_err(|error| invalid_data(&error.to_string()))?;
    if !remaining.is_empty() {
        return Err(invalid_data("trailing data in state checkpoint"));
    }
    Ok(state)
}