build.rs

/*!
Build script for `pointsource_pde`
*/

//use pyo3::{prepare_freethreaded_python, types::PyModule, PyErr, Python};
use conda_build::python_config;
use std::env::split_paths;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use std::process::Command;

const REQUIRED_PYTHON_MODULES: [&str; 7] = [
    "dolfinx", "ufl", "numpy", "time", "petsc4py", "mpi4py", "json",
];

const DOLFINX_ACCESS_CXX_SOURCES: [&str; 3] =
    ["nanobind_helpers.cc", "minmax_p2.cc", "function.cc"];

const PY_SRC: [&str; 7] = [
    "src/dolfinx_extras.py",
    "src/compose.py",
    "src/measure.py",
    "src/quadratic_dataterm.py",
    "src/convection_diffusion.py",
    "src/laser_sampling.py",
    "src/full_sampling.py",
];

fn main() -> Result<(), std::io::Error> {
    // Check Python module availability.
    // This doesn't necessarily work in a Conda setup.
    // prepare_freethreaded_python();
    // Python::with_gil(|py| {
    //     for m in REQUIRED_PYTHON_MODULES {
    //         PyModule::import(py, m)?;
    //     }
    //     Ok::<_, PyErr>(())
    // })?;

    let (pyc_e, py_is_conda) = python_config();
    let py_path = pyc_e.and_then(|pyc| pyc.exec_prefix_path())?;
    let py_exec = py_path.join("bin").join("python3");
    let import_check_code = REQUIRED_PYTHON_MODULES
        .map(|m| format!("import {}\n", m))
        .concat();
    let res = Command::new(&py_exec)
        .args(["-c", import_check_code.as_ref()])
        .status()?;
    if !res.success() {
        return Err(std::io::Error::other(format!(
            "Failed to load required Python modules (python execuitable {}{})",
            py_exec.display(),
            if py_is_conda { " from Conda" } else { "" }
        )));
    }

    // Set up local-to-crate paths
    let cc_mod_name = "dolfinx_access";
    let crate_name = "pointsource_pde";
    let src = PathBuf::from("src");
    let cc_mod_root = src.join(cc_mod_name);
    let bridge = src.join(format!("{cc_mod_name}.rs"));
    let cc_sources = DOLFINX_ACCESS_CXX_SOURCES.map(|f| cc_mod_root.join(f));
    let mut include_dirs = Vec::from([PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("include")]);

    // Add (an estimate of) dolfinx ldflags. This shouldn't really be our problem to deal with,
    // but everything that is not piss, is shit.
    add_dep_ldflags("DEP_DOLFINX_LDFLAGS");

    // Add external include paths
    add_dep_includes(&mut include_dirs, "DEP_DOLFINX_INCLUDE");
    add_dep_includes(&mut include_dirs, "DEP_NANOBIND_INCLUDE");

    // This may be useful to enable sometimes, when there's
    // trouble location GSL.
    // pkg_config::Config::new().probe("gsl").unwrap();

    // Build the bridge
    cxx_build::bridge(&bridge)
        .std("c++20")
        .files(&cc_sources)
        .includes(&include_dirs)
        .compile(crate_name);

    for dep in PY_SRC {
        println!("cargo::rerun-if-changed={}", dep);
    }

    let header_prefix = PathBuf::from("include").join(cc_mod_name);
    let headers = DOLFINX_ACCESS_CXX_SOURCES
        .into_iter()
        .map(|s| header_prefix.join(s.replace(".cc", ".h")));

    for dep in cc_sources
        .into_iter()
        .chain(headers)
        .chain(std::iter::once(bridge))
    {
        println!("cargo:rerun-if-changed={}", dep.display());
    }

    // Generate .clangd
    let mut dot_clangd = File::create(".clangd")?;
    write!(
        dot_clangd,
        "\
# Automatically generated by build.rs
CompileFlags:
    Add:
"
    )?;
    for i in include_dirs {
        writeln!(dot_clangd, "        - --include-directory={}", i.display())?;
    }

    Ok(())
}

/// Add list of include directions in the environment variable `dep` to
/// `include_dirs`.
fn add_dep_includes(include_dirs: &mut Vec<PathBuf>, dep: &str) {
    if let Some(include) = std::env::var_os(&dep) {
        include_dirs.extend(split_paths(&include).collect::<Vec<_>>());
    } else {
        panic!("{dep} unset.");
    }
}

fn add_dep_ldflags(dep: &str) {
    if let Some(ldflags) = std::env::var_os(&dep) {
        for lp in split_paths(&ldflags) {
            println!("cargo:rustc-link-arg=-Wl,{}", lp.display());
        }
    } else {
        panic!("{dep} unset.");
    }
}