sail_scripts/kernel/src/lib.rs

//! SAIL Kernel — Memory-Safe Rust Extension
//! =========================================
//! This is the top-level PyO3 module that exposes all SAIL kernel
//! subsystems to Python. Every function here is provably memory-safe
//! (no raw pointers, no unsafe blocks except in FFI boundary).
//!
//! Modules:
//!   - text_cleaner   : UTF-8 normalization, HTML stripping, whitespace
//!   - tokenizer      : Memory-safe BPE tokenizer with bounded buffer
//!   - dataset_writer : Arrow IPC writer for tokenized shards
//!   - compressor     : zstd compression/decompression of shards
//!   - token_budget   : Dynamic token budget calculator
//!   - reward_scorer  : GRPO reward signal computation

use pyo3::prelude::*;

mod text_cleaner;
mod tokenizer;
mod dataset_writer;
mod compressor;
mod token_budget;
mod reward_scorer;
mod errors;

/// Register all submodules into the Python `sail_kernel` module.
#[pymodule]
fn sail_kernel(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // ── Text Cleaner ─────────────────────────────────────────────────────────
    m.add_function(wrap_pyfunction!(text_cleaner::clean_text, m)?)?;
    m.add_function(wrap_pyfunction!(text_cleaner::clean_batch, m)?)?;

    // ── Tokenizer ────────────────────────────────────────────────────────────
    m.add_class::<tokenizer::RustTokenizer>()?;

    // ── Dataset Writer ───────────────────────────────────────────────────────
    m.add_function(wrap_pyfunction!(dataset_writer::write_arrow_shard, m)?)?;
    m.add_function(wrap_pyfunction!(dataset_writer::merge_shards, m)?)?;

    // ── Compressor ───────────────────────────────────────────────────────────
    m.add_function(wrap_pyfunction!(compressor::compress_shard, m)?)?;
    m.add_function(wrap_pyfunction!(compressor::decompress_shard, m)?)?;

    // ── Token Budget ─────────────────────────────────────────────────────────
    m.add_function(wrap_pyfunction!(token_budget::compute_budget, m)?)?;

    // ── Reward Scorer ────────────────────────────────────────────────────────
    m.add_function(wrap_pyfunction!(reward_scorer::score_output, m)?)?;
    m.add_function(wrap_pyfunction!(reward_scorer::score_batch, m)?)?;

    // ── Version info ─────────────────────────────────────────────────────────
    m.add("__version__", "0.1.0")?;
    m.add("__doc__", "SAIL Memory-Safe Kernel — Rust native extension")?;

    Ok(())
}