gbrt_rs/utils/

serialization.rs

#![allow(dead_code)]

//! Model serialization and persistence for gradient boosting models.
//!
//! This module provides robust utilities for saving and loading trained models
//! with metadata, versioning, and compression support. It supports multiple
//! serialization formats and ensures model compatibility across different
//! versions of the library.
//!
//! # Features
//!
//! - **Multiple formats**: JSON (human-readable) and Bincode (compact binary)
//! - **Compression**: Optional gzip compression to reduce file size
//! - **Versioning**: Automatic model version tracking and validation
//! - **Metadata**: Rich model information (timestamps, hyperparameters, etc.)
//! - **Validation**: Verify model files without full deserialization

use serde::{Serialize, de::DeserializeOwned};
use std::fs::{File, create_dir_all};
use std::io::{BufReader, BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use thiserror::Error;
use chrono::{DateTime, Utc};

/// Errors that can occur during model serialization/deserialization.
///
/// Covers IO errors, format-specific errors, version mismatches, and
/// data integrity issues.
#[derive(Error, Debug)]
pub enum SerializationError {
    /// Underlying filesystem or IO operation failed.
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),
    
    /// JSON serialization/deserialization failed.
    #[error("JSON serialization error: {0}")]
    JsonError(#[from] serde_json::Error),
    
    /// Bincode encoding failed.
    #[error("Bincode encoding error: {0}")]
    BincodeEncodeError(#[from] bincode::error::EncodeError),
    
    /// Bincode decoding failed.
    #[error("Bincode decoding error: {0}")]
    BincodeDecodeError(#[from] bincode::error::DecodeError),
    
    /// Model file is corrupted or has invalid structure.
    #[error("Invalid model file: {0}")]
    InvalidModelFile(String),
    
    /// Model version doesn't match current library version.
    #[error("Version mismatch: expected {expected}, got {actual}")]
    VersionMismatch { expected: String, actual: String },
   
    /// General serialization failure.
    #[error("Serialization failed: {0}")]
    SerializationFailed(String),
}

/// Result type for serialization operations.
pub type SerializationResult<T> = std::result::Result<T, SerializationError>;

/// Metadata describing a serialized model.
///
/// Tracks versioning, creation info, and model characteristics to ensure
/// reproducibility and compatibility.
#[derive(Debug, Clone, Serialize, serde::Deserialize)]
pub struct ModelMetadata {
    /// Semantic version of the library that created the model.
    pub version: String,
    /// Timestamp when the model was first saved.
    pub created_at: DateTime<Utc>,
    /// Timestamp when the model was last updated.
    pub updated_at: DateTime<Utc>,
    /// Type of model (e.g., "gradient_booster", "decision_tree").
    pub model_type: String,
    /// Number of features the model expects.
    pub num_features: usize,
    /// Number of trees in the ensemble (if applicable).
    pub num_trees: usize,
    /// Model hyperparameters and configuration.
    pub parameters: std::collections::HashMap<String, String>,
}

impl ModelMetadata {
    /// Creates new metadata for a model.
    ///
    /// Automatically sets version to current crate version and timestamps to now.
    ///
    /// # Arguments
    ///
    /// * `model_type` - Identifier for the model type
    /// * `num_features` - Number of input features
    /// * `num_trees` - Number of trees (0 for single models)
    pub fn new(model_type: &str, num_features: usize, num_trees: usize) -> Self {
        let now = Utc::now();
        Self {
            version: env!("CARGO_PKG_VERSION").to_string(),
            created_at: now,
            updated_at: now,
            model_type: model_type.to_string(),
            num_features,
            num_trees,
            parameters: std::collections::HashMap::new(),
        }
    }
    
    /// Adds a hyperparameter to the metadata (builder pattern).
    ///
    /// # Arguments
    ///
    /// * `key` - Parameter name
    /// * `value` - Parameter value (converted to string)
    ///
    /// # Returns
    ///
    /// Self with updated parameters
    pub fn with_parameter(mut self, key: &str, value: &str) -> Self {
        self.parameters.insert(key.to_string(), value.to_string());
        self
    }
    
    /// Updates the `updated_at` timestamp to current time.
    pub fn update(&mut self) {
        self.updated_at = Utc::now();
    }
   
    /// Validates that the model version matches the current library version.
    ///
    /// # Returns
    ///
    /// `Ok(())` if versions match, `Err(SerializationError::VersionMismatch)` otherwise
    pub fn validate_version(&self) -> SerializationResult<()> {
        let current_version = env!("CARGO_PKG_VERSION");
        if self.version != current_version {
            return Err(SerializationError::VersionMismatch {
                expected: current_version.to_string(),
                actual: self.version.clone(),
            });
        }
        Ok(())
    }
}

/// Main serializer with configurable format and compression.
///
/// Provides a unified interface for saving and loading models with
/// consistent error handling and metadata management.
pub struct ModelSerializer {
    /// Whether to compress the output with gzip.
    compression: bool,
    /// Serialization format to use.
    format: SerializationFormat,
}

/// Supported serialization formats.
///
/// Each format has different tradeoffs between size, speed, and readability.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum SerializationFormat {
    /// JSON format (human-readable, good interoperability).
    Json,
    /// Bincode format (compact binary, fastest serialization).
    Bincode,
}

impl Default for ModelSerializer {
    fn default() -> Self {
        Self {
            compression: true,
            format: SerializationFormat::Bincode,
        }
    }
}

impl ModelSerializer {
    /// Creates a new serializer with default settings (Bincode + compression).
    pub fn new() -> Self {
        Self::default()
    }
    
    /// Enables or disables gzip compression.
    ///
    /// # Arguments
    ///
    /// * `compression` - `true` to compress (default), `false` for uncompressed
    ///
    /// # Returns
    ///
    /// Self with updated compression setting (builder pattern)
    pub fn with_compression(mut self, compression: bool) -> Self {
        self.compression = compression;
        self
    }
   
    /// Sets the serialization format.
    ///
    /// # Arguments
    ///
    /// * `format` - [`SerializationFormat`] to use
    ///
    /// # Returns
    ///
    /// Self with updated format (builder pattern)
    pub fn with_format(mut self, format: SerializationFormat) -> Self {
        self.format = format;
        self
    }
    
    /// Saves a model to disk with metadata.
    ///
    /// Automatically creates parent directories if they don't exist.
    ///
    /// # Arguments
    ///
    /// * `model` - Model to serialize (must implement `Serialize`)
    /// * `metadata` - Model metadata
    /// * `path` - Destination file path
    ///
    /// # Returns
    ///
    /// `Ok(())` on success, [`SerializationError`] on failure 
    pub fn save_model<T: Serialize>(
        &self,
        model: &T,
        metadata: &ModelMetadata,
        path: &Path,
    ) -> SerializationResult<()> {
        save_model(model, metadata, path, self.format, self.compression)
    }
    
    /// Loads a model from disk with metadata.
    ///
    /// Automatically detects format and compression from file content.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to model file
    ///
    /// # Returns
    ///
    /// Tuple of `(deserialized_model, metadata)`
    ///
    /// # Errors
    ///
    /// Returns error if file doesn't exist, format is invalid, or version mismatch 
    pub fn load_model<T: DeserializeOwned>(&self, path: &Path) -> SerializationResult<(T, ModelMetadata)> {
        load_model(path)
    }
    
    /// Validates a model file without loading the full model.
    ///
    /// Useful for checking version compatibility or file integrity.
    ///
    /// # Arguments
    ///
    /// * `path` - Path to model file
    ///
    /// # Returns
    ///
    /// Metadata if validation succeeds
    pub fn validate_model_file(&self, path: &Path) -> SerializationResult<ModelMetadata> {
        validate_model_file(path)
    }
}


// Standalone serialization functions

/// Saves model to file with explicit format and compression settings.
///
/// This is a lower-level function; prefer [`ModelSerializer::save_model`] for most cases.
///
/// # Arguments
///
/// * `model` - Model to serialize
/// * `metadata` - Model metadata
/// * `path` - Destination file path
/// * `format` - Serialization format
/// * `compression` - Whether to compress
///
/// # Returns
///
/// `Ok(())` on success
pub fn save_model<T: Serialize>(
    model: &T,
    metadata: &ModelMetadata,
    path: &Path,
    format: SerializationFormat,
    compression: bool,
) -> SerializationResult<()> {
    // Create directory if it doesn't exist
    if let Some(parent) = path.parent() {
        create_dir_all(parent)?;
    }
    
    match format {
        SerializationFormat::Json => {
            let file = File::create(path)?;
            let writer = BufWriter::new(file);
            
            if compression {
                let mut gz_writer = flate2::write::GzEncoder::new(writer, flate2::Compression::default());
                serde_json::to_writer(&mut gz_writer, &(model, metadata))?;
                gz_writer.finish()?;
            } else {
                serde_json::to_writer(writer, &(model, metadata))?;
            }
        }
        SerializationFormat::Bincode => {
            let file = File::create(path)?;
            let mut writer = BufWriter::new(file);
            
            // Use bincode's serde integration with default configuration
            let config = bincode::config::standard();
            
            if compression {
                let mut gz_writer = flate2::write::GzEncoder::new(writer, flate2::Compression::default());
                bincode::serde::encode_into_std_write(&(model, metadata), &mut gz_writer, config)?;
                gz_writer.finish()?;
            } else {
                bincode::serde::encode_into_std_write(&(model, metadata), &mut writer, config)?;
            }
        }
    }
    
    Ok(())
}

/// Loads model from file (auto-detects format and compression).
///
/// Automatically infers format from file extension and compression from magic bytes.
///
/// # Arguments
///
/// * `path` - Path to model file
///
/// # Returns
///
/// Tuple of `(model, metadata)`
///
/// # Errors
///
/// Returns error if file is corrupted, format is unknown, or version mismatch
pub fn load_model<T: DeserializeOwned>(path: &Path) -> SerializationResult<(T, ModelMetadata)> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);
    
    // Try to detect format and compression
    let (model, metadata) = if path.extension().map_or(false, |ext| ext == "json") {
        // JSON format
        if is_gzip_compressed(path)? {
            let gz_reader = flate2::read::GzDecoder::new(reader);
            serde_json::from_reader(gz_reader)?
        } else {
            serde_json::from_reader(reader)?
        }
    } else if path.extension().map_or(false, |ext| ext == "bin") {
        // Bincode format
        let config = bincode::config::standard();
        
        if is_gzip_compressed(path)? {
            let mut gz_reader = flate2::read::GzDecoder::new(reader);
            bincode::serde::decode_from_std_read(&mut gz_reader, config)?
        } else {
            bincode::serde::decode_from_std_read(&mut reader, config)?
        }
    } else {
        // Try to auto-detect
        if let Ok(result) = load_as_bincode(&mut reader, path) {
            result
        } else if let Ok(result) = load_as_json(&mut reader, path) {
            result
        } else {
            return Err(SerializationError::InvalidModelFile(
                "Could not determine file format".to_string()
            ));
        }
    };
    
    // Validate metadata
    metadata.validate_version()?;
    
    Ok((model, metadata))
}

/// Validates a model file without loading the complete model.
///
/// For JSON files, this reads only the metadata portion. For bincode files,
/// full deserialization may be required.
///
/// # Arguments
///
/// * `path` - Path to model file
///
/// # Returns
///
/// Metadata if file is valid
pub fn validate_model_file(path: &Path) -> SerializationResult<ModelMetadata> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);
    
    // Try to read just the metadata
    if path.extension().map_or(false, |ext| ext == "json") {
        if is_gzip_compressed(path)? {
            let gz_reader = flate2::read::GzDecoder::new(reader);
            let value: serde_json::Value = serde_json::from_reader(gz_reader)?;
            extract_metadata(&value)
        } else {
            let value: serde_json::Value = serde_json::from_reader(reader)?;
            extract_metadata(&value)
        }
    } else {
        // For bincode, we need a different approach
        load_model::<serde::de::IgnoredAny>(path)?;
        Err(SerializationError::SerializationFailed(
            "Metadata validation for bincode requires full deserialization".to_string()
        ))
    }
}

/// Serializes a value to a pretty-printed JSON string.
///
/// Useful for debugging or human-readable output.
///
/// # Arguments
///
/// * `value` - Value to serialize
///
/// # Returns
///
/// JSON string
pub fn serialize_to_json<T: Serialize>(value: &T) -> SerializationResult<String> {
    Ok(serde_json::to_string_pretty(value)?)
}

/// Deserializes from a JSON string.
///
/// # Arguments
///
/// * `json` - JSON string
///
/// # Returns
///
/// Deserialized value
pub fn deserialize_from_json<T: DeserializeOwned>(json: &str) -> SerializationResult<T> {
    Ok(serde_json::from_str(json)?)
}

/// Serializes a value to bincode bytes.
///
/// # Arguments
///
/// * `value` - Value to serialize
///
/// # Returns
///
/// Binary representation
pub fn serialize_to_bincode<T: Serialize>(value: &T) -> SerializationResult<Vec<u8>> {
    let config = bincode::config::standard();
    let mut bytes = Vec::new();
    bincode::serde::encode_into_std_write(value, &mut bytes, config)?;
    Ok(bytes)
}

/// Deserializes from bincode bytes.
///
/// # Arguments
///
/// * `bytes` - Binary data
///
/// # Returns
///
/// Deserialized value
pub fn deserialize_from_bincode<T: DeserializeOwned>(bytes: &[u8]) -> SerializationResult<T> {
    let config = bincode::config::standard();
    let mut cursor = std::io::Cursor::new(bytes);
    Ok(bincode::serde::decode_from_std_read(&mut cursor, config)?)
}

// Helper functions

/// Checks if a file is gzip compressed by reading magic bytes.
///
/// # Arguments
///
/// * `path` - File path
///
/// # Returns
///
/// `true` if file starts with gzip magic number
fn is_gzip_compressed(path: &Path) -> SerializationResult<bool> {
    let file = File::open(path)?;
    let mut reader = BufReader::new(file);
    let mut buffer = [0; 2];
    
    reader.read_exact(&mut buffer)?;
    
    // Gzip magic number: 0x1f 0x8b
    Ok(buffer[0] == 0x1f && buffer[1] == 0x8b)
}

/// Loads model assuming bincode format.
fn load_as_bincode<T: DeserializeOwned>(
    reader: &mut BufReader<File>,
    path: &Path,
) -> SerializationResult<(T, ModelMetadata)> {
    let config = bincode::config::standard();
    
    if is_gzip_compressed(path)? {
        let mut gz_reader = flate2::read::GzDecoder::new(reader);
        Ok(bincode::serde::decode_from_std_read(&mut gz_reader, config)?)
    } else {
        Ok(bincode::serde::decode_from_std_read(reader, config)?)
    }
}

/// Loads model assuming JSON format.
fn load_as_json<T: DeserializeOwned>(
    reader: &mut BufReader<File>,
    path: &Path,
) -> SerializationResult<(T, ModelMetadata)> {
    if is_gzip_compressed(path)? {
        let gz_reader = flate2::read::GzDecoder::new(reader);
        Ok(serde_json::from_reader(gz_reader)?)
    } else {
        Ok(serde_json::from_reader(reader)?)
    }
}

/// Extracts metadata from a JSON value.
fn extract_metadata(value: &serde_json::Value) -> SerializationResult<ModelMetadata> {
    if let Some(metadata_value) = value.get("1") {
        Ok(serde_json::from_value(metadata_value.clone())?)
    } else {
        Err(SerializationError::InvalidModelFile(
            "No metadata found in JSON file".to_string()
        ))
    }
}

/// Utility functions for model file management.
///
/// Provides helpers for working with model files on disk, including
/// path construction, file listing, and format detection.
pub struct ModelFileUtils;

impl ModelFileUtils {
    /// Gets the appropriate file extension for a serialization format.
    ///
    /// # Arguments
    ///
    /// * `format` - Serialization format
    /// * `compressed` - Whether compression is enabled
    ///
    /// # Returns
    ///
    /// File extension string (e.g., "json.gz", "bin") 
    pub fn get_extension(format: SerializationFormat, compressed: bool) -> String {
        let base_ext = match format {
            SerializationFormat::Json => "json",
            SerializationFormat::Bincode => "bin",
        };
        
        if compressed {
            format!("{}.gz", base_ext)
        } else {
            base_ext.to_string()
        }
    }
    
    /// Constructs a full model file path with appropriate extension.
    ///
    /// # Arguments
    ///
    /// * `base_path` - Directory for model files
    /// * `model_name` - Base name of the model
    /// * `format` - Serialization format
    /// * `compressed` - Whether to compress
    ///
    /// # Returns
    ///
    /// Complete file path 
    pub fn create_model_path(
        base_path: &Path,
        model_name: &str,
        format: SerializationFormat,
        compressed: bool,
    ) -> PathBuf {
        let ext = Self::get_extension(format, compressed);
        base_path.join(format!("{}.{}", model_name, ext))
    }
    
    /// Lists all model files in a directory.
    ///
    /// Scans the directory for files with recognized extensions (.json, .bin, .gz).
    ///
    /// # Arguments
    ///
    /// * `dir` - Directory to scan
    ///
    /// # Returns
    ///
    /// Vector of model file paths 
    pub fn list_model_files(dir: &Path) -> SerializationResult<Vec<PathBuf>> {
        let mut model_files = Vec::new();
        
        for entry in std::fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            
            if path.is_file() {
                if let Some(ext) = path.extension() {
                    if ext == "json" || ext == "bin" || ext == "gz" {
                        model_files.push(path);
                    }
                }
            }
        }
        
        Ok(model_files)
    }
    
    /// Checks if a file is a valid model file based on extension.
    ///
    /// # Arguments
    ///
    /// * `path` - File path to check
    ///
    /// # Returns
    ///
    /// `true` if file has a recognized model extension
    pub fn is_model_file(path: &Path) -> bool {
        path.is_file() && path.extension().map_or(false, |ext| {
            ext == "json" || ext == "bin" || ext == "gz"
        })
    }
}