ED_LRR/rust/src/lib.rs

#![feature(binary_heap_retain)]
#![allow(dead_code, clippy::needless_return, clippy::too_many_arguments)]
#![warn(rust_2018_idioms, rust_2021_compatibility, clippy::disallowed_types)]
//! # Elite: Danerous Long Range Router
pub mod common;
pub mod galaxy;
pub mod journal;
pub mod mmap_csv;
#[cfg(feature = "profiling")]
pub mod profiling;
pub mod route;
pub mod search_algos;
pub mod ship;

// =========================
use stats_alloc::{Region, StatsAlloc, INSTRUMENTED_SYSTEM};
use std::alloc::System as SystemAlloc;
use std::collections::BTreeMap;
use std::io::{BufWriter, Write};
use std::path::Path;
use std::time::Instant;
#[cfg(not(feature = "profiling"))]
#[global_allocator]
static GLOBAL: &StatsAlloc<SystemAlloc> = &INSTRUMENTED_SYSTEM;
// =========================
#[cfg(not(feature = "profiling"))]
mod profiling {
    pub fn init() {}
}

use crate::common::{grid_stats, EdLrrError, SysEntry, System};
#[cfg(feature = "profiling")]
use crate::profiling::*;
use crate::route::{Router, SearchState};
use crate::ship::Ship;
use eyre::Result;
#[cfg(not(feature = "profiling"))]
use log::*;
use pyo3::create_exception;
use pyo3::exceptions::*;
use pyo3::prelude::*;
use pyo3::types::{IntoPyDict, PyDict, PyTuple};
use route::PyModeConfig;
use std::{collections::HashMap, convert::TryInto, fs::File, path::PathBuf};

#[cfg(feature = "mem_profiling")]
#[global_allocator]
static GLOBAL: ProfiledAllocator<std::alloc::System> =
    ProfiledAllocator::new(std::alloc::System, 1024);

create_exception!(_ed_lrr, RoutingError, PyException);
create_exception!(_ed_lrr, ProcessingError, PyException);
create_exception!(_ed_lrr, ResolveError, PyException);

#[derive(Debug)]
enum RangeOrShip {
    Range(f32),
    Ship(Ship),
}

impl FromPyObject<'_> for RangeOrShip {
    fn extract(ob: &PyAny) -> PyResult<Self> {
        if let Ok(n) = ob.extract() {
            return Ok(Self::Range(n));
        }
        let s: PyShip = ob.extract()?;
        return Ok(Self::Ship(s.ship));
    }
}
#[pyclass(dict)]
#[derive(Debug)]
#[pyo3(text_signature = "(callback, /)")]
struct PyRouter {
    router: Router,
    stars_path: Option<String>,
}

impl PyRouter {
    fn check_stars(&self) -> PyResult<PathBuf> {
        self.stars_path
            .as_ref()
            .ok_or_else(|| PyErr::from(EdLrrError::RuntimeError("no stars.csv loaded".to_owned())))
            .map(PathBuf::from)
    }
}

#[pymethods]
impl PyRouter {
    #[new]
    #[args(callback = "None")]
    fn new(callback: Option<PyObject>) -> Self {
        let mut router = Router::new();
        if callback.is_some() {
            router.set_callback(Box::new(move |state: &SearchState| {
                let gil_guard = Python::acquire_gil();
                let py = gil_guard.python();
                match callback.as_ref() {
                    Some(cb) => cb.call(py, (state.clone(),), None),
                    None => Ok(py.None()),
                }
            }))
        }
        PyRouter {
            router,
            stars_path: None,
        }
    }

    #[args(primary_only = "false", immediate = "false")]
    #[pyo3(text_signature = "(path, primary_only, /)")]
    fn load(&mut self, path: String, py: Python<'_>, immediate: bool) -> PyResult<PyObject> {
        self.stars_path = Some(path);
        if immediate {
            self.router
                .load(&self.check_stars()?)
                .map_err(PyErr::new::<PyValueError, _>)?;
        }
        Ok(py.None())
    }

    #[pyo3(text_signature = "(/)")]
    fn unload(&mut self, py: Python<'_>) -> PyObject {
        self.router.unload();
        py.None()
    }

    fn plot(&mut self, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        let mut max_v = [0f32, 0f32, 0f32];
        let mut min_v = [0f32, 0f32, 0f32];
        for node in self.router.get_tree().iter() {
            for i in 0..3 {
                if node.pos[i] > max_v[i] {
                    max_v[i] = node.pos[i];
                }
                if node.pos[i] < min_v[i] {
                    min_v[i] = node.pos[i];
                }
            }
        }
        let plot_bbox: ((f32, f32), (f32, f32)) = ((min_v[0], max_v[0]), (min_v[2], max_v[2]));
        Ok(plot_bbox.to_object(py))
    }

    fn run_bfs(&mut self, range: f32, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        self.router
            .precomp_bfs(range)
            .map_err(PyErr::new::<RoutingError, _>)
            .map(|_| py.None())
    }

    fn precompute_graph(&mut self, range: f32, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        self.router
            .precompute_graph(range)
            .map_err(PyErr::new::<RoutingError, _>)
            .map(|_| py.None())
    }

    fn nb_perf_test(&mut self, range: f32, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        let tree = self.router.get_tree();
        let total_nodes = tree.size();
        let mut total_nbs = 0;
        for (n, node) in tree.iter().enumerate() {
            total_nbs += self.router.neighbours(node, range).count();
            // nbmap.insert(node.id, nbs);
            if n % 100_000 == 0 {
                let avg = total_nbs as f64 / (n + 1) as f64;
                info!("{}/{} {} ({})", n, total_nodes, total_nbs, avg);
            }
        }
        let avg = total_nbs as f64 / total_nodes as f64;
        info!("Total: {} ({})", total_nbs, avg);
        Ok(py.None())
    }

    fn precompute_neighbors(&mut self, range: f32, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        self.router
            .precompute_all(range)
            .map_err(PyErr::new::<RoutingError, _>)
            .map(|_| py.None())
    }

    fn bfs_test(&mut self, range: f32) -> PyResult<()> {
        use rand::prelude::*;
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        let mut rng = rand::rngs::StdRng::seed_from_u64(0);
        let nodes = self.router.get_tree().size();
        loop {
            let source = *self
                .router
                .get_tree()
                .iter()
                .nth(rng.gen_range(0..nodes))
                .unwrap();
            let goal = *self
                .router
                .get_tree()
                .iter()
                .nth(rng.gen_range(0..nodes))
                .unwrap();
            self.router.bfs_loop_test(range, &source, &goal, 0);
            for w in 0..=15 {
                self.router.bfs_loop_test(range, &source, &goal, 2usize.pow(w));
            }
        }
        Ok(())
    }

    #[args(
        greedyness = "0.5",
        max_dist = "0.0",
        num_workers = "0",
        beam_width = "BeamWidth::Absolute(0)"
    )]
    #[pyo3(text_signature = "(hops, range, mode, num_workers, /)")]
    fn route(
        &mut self,
        hops: Vec<SysEntry>,
        range: RangeOrShip,
        mode: Option<PyModeConfig>,
        num_workers: usize,
    ) -> PyResult<Vec<common::System>> {
        let stars_path = self.check_stars()?;
        let route_res = self.router.load(&stars_path);
        if let Err(err_msg) = route_res {
            return Err(PyErr::new::<PyValueError, _>(err_msg));
        };
        info!("Resolving systems...");
        let ids: Vec<u32> = match resolve(&hops, &self.router.path, true) {
            Ok(sytems) => sytems.into_iter().map(|id| id.into_id()).collect(),
            Err(err_msg) => {
                return Err(err_msg.into());
            }
        };
        let mut is_default = false;
        let mut is_ship = false;
        info!("{:?}", mode);
        let mut mode = match mode {
            Some(mode) => mode,
            None => {
                let mode = PyModeConfig::default();
                is_default = true;
                mode
            }
        };
        if mode.mode.is_empty() {
            if mode.ship.is_none() {
                mode.mode = "bfs".to_string();
            } else {
                mode.mode = "ship".to_string();
                if mode.ship_mode == *"" {
                    mode.ship_mode = "jumps".to_string();
                }
            }
        }
        let range = match range {
            RangeOrShip::Range(r) => Some(r),
            RangeOrShip::Ship(ship) => {
                mode.mode = "ship".into();
                mode.ship = Some(ship);
                is_ship = true;
                None
            }
        };
        info!("{:?}", mode);
        let mode = mode.try_into()?;
        if is_default && !is_ship {
            warn!("no mode specified, defaulting to {}", mode);
        }
        #[cfg(not(feature = "profiling"))]
        let reg = Region::new(GLOBAL);
        let res = match self.router.compute_route(&ids, range, mode, num_workers) {
            Ok(route) => Ok(route),
            Err(err_msg) => Err(PyErr::new::<RoutingError, _>(err_msg)),
        };
        #[cfg(not(feature = "profiling"))]
        println!("{:?}", reg.change());
        return res;
    }

    fn perf_test(&self, callback: PyObject, py: Python<'_>) -> PyResult<PyObject> {
        use common::TreeNode;
        let node = TreeNode {
            pos: [-65.21875, 7.75, -111.03125],
            flags: 1,
            id: 0,
        };
        let goal = TreeNode {
            pos: [-9530.5, -910.28125, 19808.125],
            flags: 1,
            id: 1,
        };
        let kwargs = vec![("goal", goal), ("node", node)].into_py_dict(py);
        let mut n: usize = 0;
        let mut d: f64 = 0.0;
        let num_loops = 10_000_000;
        loop {
            let pool = unsafe { Python::new_pool(py) };
            let t_start = std::time::Instant::now();
            for _ in 0..num_loops {
                let val: f64 = callback.call(py, (), Some(kwargs))?.extract(py)?;
            }
            d += t_start.elapsed().as_secs_f64();
            drop(pool);
            n += num_loops;
            let dt = std::time::Duration::from_secs_f64(d / (n as f64));
            println!("{}: {:?}", n, dt);
        }
        Ok(py.None())
    }

    #[args(grid_size = "1.0")]
    #[pyo3(text_signature = "(grid_size)")]
    fn get_grid(&self, grid_size: f32, py: Python<'_>) -> PyResult<PyObject> {
        let stars_path = self.check_stars()?;
        grid_stats(&stars_path, grid_size)
            .map(|ret| ret.to_object(py))
            .map_err(PyErr::new::<PyRuntimeError, _>)
    }

    fn floyd_warshall_test(&mut self, range: f32) -> PyResult<Vec<common::System>> {
        let stars_path = self.check_stars()?;
        self.router
            .load(&stars_path)
            .map_err(PyErr::new::<PyValueError, _>)?;
        let res = self
            .router
            .floyd_warshall(range)
            .map_err(PyErr::new::<RoutingError, _>)?;
        Ok(res)
    }

    #[args(hops = "*")]
    #[pyo3(text_signature = "(sys_1, sys_2, ..., /)")]
    fn resolve(&self, hops: Vec<SysEntry>, py: Python<'_>) -> PyResult<PyObject> {
        info!("Resolving systems...");
        let stars_path = self.check_stars()?;
        let systems: Vec<System> = match resolve(&hops, &stars_path, false) {
            Ok(sytems) => sytems.into_iter().map(|sys| sys.into_system()).collect(),
            Err(err_msg) => {
                return Err(err_msg.into());
            }
        };
        let ret: Vec<(_, System)> = hops
            .into_iter()
            .zip(systems.iter())
            .map(|(id, sys)| (id, sys.clone()))
            .collect();
        Ok(PyDict::from_sequence(py, ret.to_object(py))?.to_object(py))
    }

    fn str_tree_test(&self) -> common::EdLrrResult<()> {
        use common::BKTree;
        const CHUNK_SIZE: usize = 4_000_000;
        let path = self.check_stars()?;
        let reader: csv::Reader<File> = csv::ReaderBuilder::new()
            .has_headers(false)
            .from_path(path)
            .map_err(EdLrrError::from)?;
        let mut data: Vec<String> = Vec::with_capacity(CHUNK_SIZE);
        let t_start = Instant::now();
        let mut base_id = 0;
        let mut wr = BufWriter::new(File::create("test.bktree")?);
        for sys in reader.into_deserialize::<System>() {
            let sys = sys?;
            data.push(sys.name);
            if data.len() > CHUNK_SIZE {
                let tree = BKTree::new(&data, base_id);
                tree.dump(&mut wr)?;
                base_id = sys.id;
            }
        }
        if !data.is_empty() {
            let tree = BKTree::new(&data, base_id);
            tree.dump(&mut wr)?;
        }
        wr.flush()?;
        println!("Took: {:?}", t_start.elapsed());
        Ok(())
    }

    fn __str__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self))
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self))
    }
}
enum ResolveResult {
    System(System),
    ID(u32),
}

impl ResolveResult {
    fn into_id(self) -> u32 {
        match self {
            Self::System(sys) => sys.id,
            Self::ID(id) => id,
        }
    }

    fn into_system(self) -> System {
        if let Self::System(sys) = self {
            return sys;
        }
        panic!("Tried to unwrap ID into System");
    }
}

fn resolve(
    entries: &[SysEntry],
    path: &Path,
    id_only: bool,
) -> Result<Vec<ResolveResult>, EdLrrError> {
    let mut names: Vec<String> = Vec::new();
    let mut ret: Vec<u32> = Vec::new();
    let mut needs_rtree = false;
    for ent in entries {
        match ent {
            SysEntry::Name(name) => names.push(name.to_owned()),
            SysEntry::Pos(_) => {
                needs_rtree |= true;
            }
            _ => (),
        }
    }
    if !path.exists() {
        return Err(EdLrrError::ResolveError(format!(
            "Source file {:?} does not exist!",
            path.display()
        )));
    }
    let name_ids = if !names.is_empty() {
        mmap_csv::mmap_csv(path, names)?
    } else {
        HashMap::new()
    };
    let tmp_r = needs_rtree
        .then(|| {
            let mut r = Router::new();
            r.load(path).map(|_| r)
        })
        .transpose()?;
    for ent in entries {
        match ent {
            SysEntry::Name(name) => {
                let ent_res = name_ids.get(name).ok_or_else(|| {
                    EdLrrError::ResolveError(format!("System {} not found", name))
                })?;
                let sys = ent_res.as_ref().ok_or_else(|| {
                    EdLrrError::ResolveError(format!("System {} not found", name))
                })?;
                ret.push(*sys);
            }
            SysEntry::ID(id) => ret.push(*id),
            SysEntry::Pos((x, y, z)) => ret.push(
                tmp_r
                    .as_ref()
                    .unwrap()
                    .closest(&[*x, *y, *z])
                    .ok_or_else(|| EdLrrError::ResolveError("No systems loaded!".to_string()))?
                    .id,
            ),
        }
    }
    if id_only {
        return Ok(ret.iter().map(|id| ResolveResult::ID(*id)).collect());
    } else {
        let mut lc = route::LineCache::create(path)?;
        let mut systems = vec![];
        for id in ret {
            let sys = ResolveResult::System(lc.get(id)?.unwrap());
            systems.push(sys)
        }
        return Ok(systems);
    }
}

#[pyclass(dict)]
#[derive(Debug, Clone)]
struct PyShip {
    ship: Ship,
}
#[pymethods]
impl PyShip {
    #[staticmethod]
    fn from_loadout(py: Python<'_>, loadout: &str) -> PyResult<PyObject> {
        match Ship::new_from_json(loadout) {
            Ok(ship) => Ok((PyShip { ship: ship.1 }).into_py(py)),
            Err(err_msg) => Err(PyErr::new::<PyValueError, _>(err_msg)),
        }
    }
    #[staticmethod]
    fn from_journal(py: Python<'_>) -> PyResult<PyObject> {
        let mut ship = match Ship::new_from_journal() {
            Ok(ship) => ship,
            Err(err_msg) => {
                return Err(PyErr::new::<PyValueError, _>(err_msg));
            }
        };
        let ships: Vec<(PyObject, PyObject)> = ship
            .drain()
            .map(|(ship_name, ship)| {
                let ship_name_py = ship_name.to_object(py);
                let ship_py = (PyShip { ship }).into_py(py);
                (ship_name_py, ship_py)
            })
            .collect();
        Ok(PyDict::from_sequence(py, ships.to_object(py))?.to_object(py))
    }

    fn to_dict(&self, py: Python<'_>) -> PyResult<PyObject> {
        self.ship.to_object(py)
    }

    #[pyo3(text_signature = "(dist, /)")]
    fn fuel_cost(&self, _py: Python<'_>, dist: f32) -> f32 {
        self.ship.fuel_cost(dist)
    }

    #[getter]
    fn range(&self, _py: Python<'_>) -> f32 {
        self.ship.range()
    }

    #[getter]
    fn max_range(&self, _py: Python<'_>) -> f32 {
        self.ship.max_range()
    }

    #[pyo3(text_signature = "(dist, /)")]
    fn make_jump(&mut self, dist: f32, _py: Python<'_>) -> Option<f32> {
        self.ship.make_jump(dist)
    }

    #[pyo3(text_signature = "(dist, /)")]
    fn can_jump(&self, dist: f32, _py: Python<'_>) -> bool {
        self.ship.can_jump(dist)
    }

    #[args(fuel_amount = "None")]
    #[pyo3(text_signature = "(fuel_amount, /)")]
    fn refuel(&mut self, fuel_amount: Option<f32>, _py: Python<'_>) {
        if let Some(fuel) = fuel_amount {
            self.ship.fuel_mass = (self.ship.fuel_mass + fuel).min(self.ship.fuel_capacity)
        } else {
            self.ship.fuel_mass = self.ship.fuel_capacity;
        }
    }

    #[pyo3(text_signature = "(factor, /)")]
    fn boost(&mut self, factor: f32, _py: Python<'_>) {
        self.ship.boost(factor);
    }

    fn __str__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self.ship))
    }

    fn __repr__(&self) -> PyResult<String> {
        Ok(format!("{:?}", &self.ship))
    }
}

impl PyShip {
    fn get_ship(&self) -> Ship {
        self.ship.clone()
    }
}

#[pyfunction]
fn preprocess_edsm(
    _bodies_path: &str,
    _systems_path: &str,
    _out_path: &str,
    _py: Python<'_>,
) -> PyResult<()> {
    Err(pyo3::exceptions::PyNotImplementedError::new_err(
        "please use Spansh's Galaxy dump and preprocess_galaxy()",
    ))
}

fn to_py_value(value: eval::Value, py: Python<'_>) -> PyResult<PyObject> {
    type Value = eval::Value;
    match value {
        Value::String(s) => Ok(s.to_object(py)),
        Value::Number(n) => {
            if let Some(n) = n.as_u64() {
                return Ok(n.to_object(py));
            }
            if let Some(n) = n.as_i64() {
                return Ok(n.to_object(py));
            }
            return Ok(n.as_f64().unwrap().to_object(py));
        }
        Value::Bool(b) => Ok(b.to_object(py)),
        Value::Array(mut t) => {
            let mut res: Vec<PyObject> = vec![];
            for v in t.drain(..) {
                res.push(to_py_value(v, py)?);
            }
            Ok(PyTuple::new(py, &res).to_object(py))
        }
        Value::Object(o) => {
            let res = PyDict::new(py);
            for (k, v) in o.iter() {
                res.set_item(k, to_py_value(v.clone(), py)?)?;
            }
            Ok(res.to_object(py))
        }
        Value::Null => Ok(py.None()),
    }
}

fn to_py(res: Result<eval::Value, eval::Error>, py: Python<'_>) -> PyResult<PyObject> {
    res.map_err(|e| PyErr::from(EdLrrError::EvalError(e)))
        .and_then(|r| to_py_value(r, py))
}

#[pyfunction]
#[pyo3(text_signature = "(expr)")]
fn expr_test(expr: &str, py: Python<'_>) -> PyResult<PyObject> {
    use eval::{to_value, Expr, Value};
    let mut res = Expr::new(expr)
        .compile()
        .map_err(|e| PyErr::from(EdLrrError::EvalError(e)))?;
    let mut hm: HashMap<&str, Value> = HashMap::new();
    hm.insert("foo", to_value(42));
    hm.insert("bar", to_value((-1, -2, -3)));
    res = res.value("x", vec!["Hello", "world", "!"]);
    res = res.value("y", 42);
    res = res.value("p", (2.17, 5.14, 1.62));
    res = res.value("hw", "Hello World!");
    res = res.value("hm", hm);
    to_py(res.exec(), py)
}

#[pyfunction]
#[pyo3(text_signature = "(path, out_path, /)")]
fn preprocess_galaxy(path: &str, out_path: &str) -> PyResult<()> {
    use common::build_index;
    use galaxy::process_galaxy_dump;
    let path = PathBuf::from(path);
    let out_path = PathBuf::from(out_path);
    process_galaxy_dump(&path, &out_path).unwrap();
    build_index(&out_path)?;
    Ok(())
}

#[pymodule]
pub fn _ed_lrr(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    better_panic::install();
    pyo3_log::init();
    profiling::init();
    m.add_class::<PyRouter>()?;
    m.add_class::<PyShip>()?;
    m.add_wrapped(pyo3::wrap_pyfunction!(preprocess_galaxy))?;
    m.add_wrapped(pyo3::wrap_pyfunction!(preprocess_edsm))?;
    m.add_wrapped(pyo3::wrap_pyfunction!(expr_test))?;
    Ok(())
}