bosk 0.1.0

//! # bosk
//!
//! *A bosk is a small thicket of trees — here, an ensemble of decision trees.*
//!
//! Gradient-boosting model **inference** behind one small [`Model`] trait: load a
//! trained LightGBM, ONNX, or CatBoost model and call [`Model::predict`].
//!
//! The headline is the **LightGBM backend: it parses the `.lgb` text format
//! directly**, with no C dependency and no FFI to `lib_lightgbm`. Built with
//! `--no-default-features` the LightGBM predictor has **zero** dependencies —
//! trivial to cross-compile, embed, or drop into a `musl` target.
//!
//! ```no_run
//! use bosk::{load_model, Model};
//!
//! # fn main() -> bosk::Result<()> {
//! let model = load_model(std::path::Path::new("model.lgb"))?;
//! let features = [0.5, -0.2, 0.7, 1.1, -0.9];
//! let p = model.predict(&features)?; // probability for a binary classifier
//! # let _ = p;
//! # Ok(())
//! # }
//! ```
//!
//! `load_model` auto-detects the format by extension; [`supported_extensions`]
//! reports what the current build can load.

#![forbid(unsafe_code)]

mod error;

pub use error::{Error, Result};

mod lgb;
pub use lgb::LgbModel;

#[cfg(feature = "onnx")]
mod onnx;
#[cfg(feature = "onnx")]
pub use onnx::OnnxModel;

#[cfg(feature = "catboost")]
mod catboost;
#[cfg(feature = "catboost")]
pub use catboost::{CatBoostModel, Output};

use std::path::Path;

/// Unified inference interface across all backends.
///
/// [`predict`](Model::predict) returns a `Result` because a wrong-length
/// feature vector is refused with [`Error::FeatureCount`] (rather than
/// silently evaluated with missing values) and because the ONNX and CatBoost
/// backends can fail at inference time. The one case where the pre-check is
/// impossible is an ONNX graph whose input feature dimension is dynamic —
/// there the shape mismatch surfaces from ONNX Runtime as
/// [`Error::Backend`]. On the pure-Rust LightGBM path the feature-count
/// refusal is the only possible error: evaluation itself cannot fail once
/// the model is loaded and validated — hold a concrete [`LgbModel`] and use
/// [`LgbModel::predict_unchecked`] where an infallible call matters more
/// than the length check.
pub trait Model: Send + Sync {
    /// Predict a single sample. For a binary classifier this is the positive-class
    /// probability.
    fn predict(&self, features: &[f64]) -> Result<f64>;

    /// Predict a batch of samples, one output per row.
    ///
    /// `flat` holds the samples row-major: `n_features` values per row, with
    /// no copying or per-row allocation required of the caller. Returns
    /// [`Error::BatchShape`] when `flat` is not a whole number of rows (or
    /// `n_features` is zero). Backends may override this with a more
    /// efficient batched implementation.
    fn predict_batch(
        &self,
        flat: &[f64],
        n_features: usize,
    ) -> Result<Vec<f64>> {
        check_batch_shape(flat, n_features)?;
        flat.chunks_exact(n_features)
            .map(|s| self.predict(s))
            .collect()
    }
}

/// Shared [`Model::predict_batch`] input validation: `flat` must split into
/// whole rows of `n_features`.
fn check_batch_shape(flat: &[f64], n_features: usize) -> Result<()> {
    if n_features == 0 || flat.len() % n_features != 0 {
        return Err(Error::BatchShape {
            len: flat.len(),
            n_features,
        });
    }
    Ok(())
}

/// Load a model from a file, auto-detecting the format by extension
/// (case-insensitively).
///
/// Available formats depend on compiled features: `lgb`/`txt` always (both
/// map to the LightGBM text format — `.txt` is what LightGBM's own
/// `save_model` conventionally writes), `onnx` and `cbm` only when the
/// corresponding feature is enabled in this build. Loading an extension this
/// build cannot handle returns [`Error::UnsupportedFormat`].
///
/// A `.cbm` file is assumed to be a **binary classifier** (the CatBoost C API
/// does not expose the trained loss function, so this cannot be verified).
/// For a CatBoost regression or ranking model, construct the backend
/// explicitly with `CatBoostModel::load(path, Output::Raw)`.
pub fn load_model(path: &Path) -> Result<Box<dyn Model>> {
    let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
    match ext.to_ascii_lowercase().as_str() {
        "lgb" | "txt" => Ok(Box::new(LgbModel::load(path)?)),
        #[cfg(feature = "onnx")]
        "onnx" => Ok(Box::new(OnnxModel::load(path)?)),
        #[cfg(feature = "catboost")]
        "cbm" => Ok(Box::new(CatBoostModel::load(
            path,
            Output::Probability,
        )?)),
        _ => Err(Error::UnsupportedFormat {
            ext: ext.to_string(),
            supported: supported_extensions(),
        }),
    }
}

/// Extensions [`load_model`] accepts in this build, in preferred order.
///
/// The slice always matches [`load_model`] exactly: an extension is listed
/// here if and only if the current feature set can load it. Use it to look
/// a model up by basename (`model.onnx`, then `model.lgb`, …) without
/// hard-coding which formats the build has. If you instead glob a whole
/// directory against this list, note that `.txt` — LightGBM's conventional
/// extension — also matches unrelated text files, so treat a per-file
/// [`Error::Parse`] as "not a model", not as a fatal condition.
#[must_use]
pub fn supported_extensions() -> &'static [&'static str] {
    &[
        #[cfg(feature = "onnx")]
        "onnx",
        "lgb",
        "txt",
        #[cfg(feature = "catboost")]
        "cbm",
    ]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_load_model_lgb_extensions() {
        // Uppercase and the conventional LightGBM `.txt` must both load.
        for name in ["bosk_case_test.LGB", "bosk_case_test.txt"] {
            let path = std::env::temp_dir().join(name);
            std::fs::copy("tests/fixtures/tiny_binary.lgb", &path)
                .expect("copy fixture");
            let model = load_model(&path).expect(name);
            let _ = std::fs::remove_file(&path);
            assert!(model
                .predict(&[0.5, -0.2, 0.7, 1.1, -0.9])
                .is_ok());
        }
    }

    /// Pins the documented contract: an extension is listed in
    /// [`supported_extensions`] iff [`load_model`] accepts it in this build.
    /// For a listed extension a nonexistent path must fail with anything
    /// but [`Error::UnsupportedFormat`].
    #[test]
    fn test_supported_extensions_match_load_model() {
        for ext in supported_extensions() {
            let path = format!("no_such_model.{ext}");
            let Err(err) = load_model(Path::new(&path)) else {
                panic!(
                    "{ext}: expected an error for a nonexistent file"
                );
            };
            assert!(
                !matches!(err, Error::UnsupportedFormat { .. }),
                "{ext} is listed but load_model refused it: {err:?}"
            );
        }
    }

    #[test]
    fn test_load_model_unknown_extension() {
        // `Box<dyn Model>` has no Debug impl, so destructure by hand.
        let Err(err) = load_model(Path::new("model.xgboost")) else {
            panic!("expected UnsupportedFormat, got Ok");
        };
        assert!(
            matches!(err, Error::UnsupportedFormat { .. }),
            "got {err:?}"
        );
    }
}