cp2k-rs 0.2.3

//! Rust bindings for CP2K
//!
//! This crate provides Rust bindings for the CP2K quantum chemistry package,
//! with additional Python bindings through PyO3.

mod ffi;
#[cfg(feature = "extended")]
mod ffi_extended;
/// Raw FFI bindings for extended interface (for direct shm writes in worker).
#[cfg(feature = "extended")]
pub mod ffi_extended_raw {
    pub use crate::ffi_extended::*;
}
mod force_env;

/// POSIX shared memory helpers for zero-copy IPC of large arrays.
pub mod shm;

/// IPC protocol shared between the Python frontend and the MPI worker binary.
pub mod worker_protocol;

/// Pure-Rust worker management (spawning, IPC, lifecycle). GIL-free.
pub mod worker;

#[cfg(feature = "mpi")]
use mpi::{ffi::MPI_Comm_c2f, raw::AsRaw};

use std::sync::atomic::{AtomicBool, Ordering};

#[cfg(feature = "extended")]
pub use force_env::GridInfo;
pub use force_env::{CP2KError, CP2KResult, ForceEnv};

#[cfg(feature = "python")]
pub mod python;

/// Tracks whether CP2K was initialized with or without MPI support.
/// Used to ensure finalize() calls the matching teardown function.
static INIT_WITH_MPI: AtomicBool = AtomicBool::new(true);

/// Initialize CP2K with MPI support
///
/// This function automatically detects if MPI is already initialized (e.g., by mpi4py
/// in Python or mpi crate in Rust) and uses the existing MPI environment if available.
/// Otherwise, it initializes MPI as part of CP2K initialization.
///
/// This is the recommended initialization function when using MPI, as it avoids
/// double MPI initialization conflicts.
pub fn init() -> CP2KResult<()> {
    unsafe {
        let result = ffi::cp2k_init_with_mpi_check();
        if result != 0 {
            return Err(CP2KError::InitializationError(format!(
                "CP2K initialization failed with code {}",
                result
            )));
        }
    }
    INIT_WITH_MPI.store(true, Ordering::SeqCst);
    Ok(())
}

/// Initialize CP2K without MPI support
pub fn init_without_mpi() -> CP2KResult<()> {
    unsafe {
        ffi::cp2k_init_without_mpi();
    }
    INIT_WITH_MPI.store(false, Ordering::SeqCst);
    Ok(())
}

/// Finalize CP2K with MPI support
///
/// This function finalizes CP2K without finalizing MPI if MPI was already
/// initialized by the host application (e.g., mpi4py or mpi crate).
/// The host application is responsible for finalizing MPI.
///
/// # Important
/// All [`ForceEnv`](crate::ForceEnv) instances must be dropped before calling
/// this function.  CP2K's `f77_interface` will abort with "invalid env_id" if
/// a live force environment is still registered.  Explicitly drop your
/// environments before finalization:
///
/// ```ignore
/// drop(force_env);
/// cp2k_rs::finalize()?;
/// ```
pub fn finalize() -> CP2KResult<()> {
    unsafe {
        let result = if INIT_WITH_MPI.load(Ordering::SeqCst) {
            ffi::cp2k_finalize_with_mpi_check()
        } else {
            ffi::cp2k_finalize_without_mpi();
            0
        };
        if result != 0 {
            return Err(CP2KError::FinalizationError(format!(
                "CP2K finalization failed with code {}",
                result
            )));
        }
    }
    Ok(())
}

/// Finalize CP2K without MPI support
pub fn finalize_without_mpi() -> CP2KResult<()> {
    unsafe {
        ffi::cp2k_finalize_without_mpi();
    }
    Ok(())
}

/// Get the CP2K version string
pub fn get_version() -> CP2KResult<String> {
    const MAX_VERSION_LEN: usize = 256;
    let mut version_buf = vec![0u8; MAX_VERSION_LEN];

    unsafe {
        ffi::cp2k_get_version(version_buf.as_mut_ptr() as *mut i8, MAX_VERSION_LEN as i32);
    }

    // Find the null terminator
    let end = version_buf
        .iter()
        .position(|&c| c == 0)
        .unwrap_or(version_buf.len());
    version_buf.truncate(end);

    // Convert to String
    Ok(String::from_utf8(version_buf)?)
}

/// Run a CP2K input file
pub fn run_input(input_file: &str, output_file: &str) -> CP2KResult<()> {
    let input_c = std::ffi::CString::new(input_file)?;
    let output_c = std::ffi::CString::new(output_file)?;

    unsafe {
        ffi::cp2k_run_input(input_c.as_ptr(), output_c.as_ptr());
    }

    Ok(())
}

/// Run a CP2K input file with custom MPI communicator
#[cfg(feature = "mpi")]
pub fn run_input_comm(
    input_file: &str,
    output_file: &str,
    comm: &mpi::topology::SimpleCommunicator,
) -> CP2KResult<()> {
    let input_c = std::ffi::CString::new(input_file)?;
    let output_c = std::ffi::CString::new(output_file)?;
    let raw_comm = comm.as_raw();
    let fortran_comm = unsafe { MPI_Comm_c2f(raw_comm) };

    unsafe {
        ffi::cp2k_run_input_comm(input_c.as_ptr(), output_c.as_ptr(), fortran_comm);
    }

    Ok(())
}