ED_LRR/rust/src/lib.rs

mod common;
mod preprocess;
mod route;
use pyo3::prelude::*;
use pyo3::types::{PyDict,PyList};
use pyo3::exceptions::*;
use std::path::PathBuf;

#[pymodule]
pub fn _ed_lrr(_py: Python, m: &PyModule) -> PyResult<()> {
    
    /// preprocess(infile_systems, infile_bodies, outfile, callback)
    /// --
    ///
    /// Preprocess bodies.json and systemsWithCoordinates.json into stars.csv
    #[pyfn(m, "preprocess")]
    fn ed_lrr_preprocess(
        py: Python<'static>,
        infile_systems: String,
        infile_bodies: String,
        outfile: String,
        callback: PyObject,
    ) -> PyResult<PyObject> {
        use preprocess::*;
        let state = PyDict::new(py);
        let state_dict = PyDict::new(py);
        callback.call(py,(state_dict,),None).unwrap();
        let callback_wrapped = move |state: &PreprocessState| {
            // println!("SEND: {:?}",state);
            state_dict.set_item("file",state.file.clone())?;
            state_dict.set_item("total",state.total)?;
            state_dict.set_item("done",state.done)?;
            callback.call(py,(state_dict,),None)
        };
        preprocess_files(&PathBuf::from(infile_bodies), &PathBuf::from(infile_systems), &PathBuf::from(outfile), Box::new(callback_wrapped)).unwrap();
        return Ok(state.to_object(py));
    }

    /// route(infile, hops, range, mode,primary, greedyness, precomp, callback)
    /// --
    ///
    /// Compute a Route using the suplied parameters
    #[pyfn(m, "route")]
    fn route(
        py: Python<'static>,
        hops: Vec<String>,
        range: f32,
        mode: String,
        primary: bool,
        permute: bool,
        greedyness: Option<f32>,
        precomp: Option<String>,
        path: String,
        callback: PyObject,
    ) -> PyResult<PyObject> {
        use route::*;
        // TODO: Verify Parameters
        let mode = match Mode::parse(&mode) {
            Ok(val) => val,
            Err(e) => {
                return Err(PyErr::new::<ValueError,_>(e));
            }
        };
        let state_dict = PyDict::new(py);
        callback.call(py,(state_dict,),None).unwrap();
        let callback_wrapped = move |state: &SearchState| {
            println!("SEND: {:?}",state);
            state_dict.set_item("mode",state.mode.clone())?;
            state_dict.set_item("system",state.system.clone())?;
            state_dict.set_item("body",state.body.clone())?;
            state_dict.set_item("depth",state.depth)?;
            state_dict.set_item("queue_size",state.queue_size)?;
            state_dict.set_item("d_rem",state.d_rem)?;
            state_dict.set_item("d_total",state.d_total)?;
            state_dict.set_item("prc_done",state.prc_done)?;
            state_dict.set_item("n_seen",state.n_seen)?;
            state_dict.set_item("prc_seen",state.prc_seen)?;
            callback.call(py,(state_dict,),None)
        };
        let opts=RouteOpts{
            systems:hops,
            range:Some(range),
            file_path: PathBuf::from(path),
            precomp_file: precomp.map(PathBuf::from),
            callback: Box::new(callback_wrapped),
            mode,
            factor: greedyness,
            precompute: false,
            permute: permute,
            keep_first: true,
            keep_last: true,
            primary
        };
        let none=().to_object(py);
        match route(opts) {
            Ok(Some(route)) => {
                let hops=route.iter().map(|hop| {
                    let pos=PyList::new(py,hop.pos.iter());
                    let elem=PyDict::new(py);
                    elem.set_item("star_type",hop.star_type.clone()).unwrap();
                    elem.set_item("system",hop.system.clone()).unwrap();
                    elem.set_item("body",hop.body.clone()).unwrap();
                    elem.set_item("distance",hop.distance).unwrap();
                    elem.set_item("pos",pos).unwrap();
                    return elem;
                });
                let lst=PyList::new(py,hops);
                return Ok(lst.to_object(py));
            }
            Ok(None) => {
                return Ok(none);
            },
            Err(e) => {
                return Err(PyErr::new::<ValueError,_>(e));
            }
        };
        /*
        let state = PyDict::new(py);
        state.set_item("infile", infile)?;
        state.set_item("source", source)?;
        state.set_item("dest", dest)?;
        state.set_item("range", range)?;
        state.set_item("mode", mode)?;
        state.set_item("greedyness", greedyness)?;
        state.set_item("precomp", precomp)?;
        state.set_item("callback", callback)?;
        return callback.call(py,(state.to_object(py),),None);
        */
    }
    Ok(())
}