cch_rs/cch/

mod.rs

//! Core engine and orchestration for Customizable Contraction Hierarchies.
//!
//! This module provides the primary entry point, [`CchEngine`], which manages the lifecycle
//! of the routing data. It uses atomic swaps for topology updates and per-profile locks for
//! metric updates, allowing weight changes without cloning the whole metric snapshot.
//!
//! # Key Traits
//! * [`CchGraph`]: Must be implemented by your base graph representation to provide basic
//!   topology and iteration capabilities.
//! * [`ProfileId`]: A marker trait for differentiating routing profiles (e.g., Car, Bike).

use std::cell::RefCell;
use std::hash::Hash;
use std::sync::{Arc, OnceLock};

use arc_swap::ArcSwap;
use parking_lot::RwLock;
use rustc_hash::{FxBuildHasher, FxHashMap};

use crate::cch::contract::Topology;
use crate::cch::customize::{PartialUpdateContext, Shortcuts, WeightMap, Weights};
use crate::cch::query::Query;

pub mod contract;
pub mod customize;
pub mod order;
pub mod path;
pub mod query;

/// A marker trait for differentiating routing profiles (e.g., Car, Bike, Pedestrian).
///
/// Types implementing this trait are used as keys in the engine to concurrently
/// manage and query multiple sets of metric weights over the same underlying topology.
pub trait ProfileId: Eq + Hash + Copy + Send + Sync + 'static {}
impl<T: Eq + Hash + Copy + Send + Sync + 'static> ProfileId for T {}

thread_local! {
    static QUERY_BUFFER: RefCell<Query> = RefCell::new(Query::new(0));
}

/// Represents the result of a successful shortest-path query.
#[derive(Debug, Clone)]
pub struct PathResult {
    /// The total accumulated weight (distance/time) of the shortest path.
    pub distance: f32,
    /// The sequence of original graph node IDs forming the shortest path.
    pub path: Vec<u32>,
}

/// The base graph representation required by the CCH engine.
///
/// This trait abstracts the underlying memory layout of your road network,
/// providing the core structural queries needed to build the elimination tree.
pub trait CchGraph {
    /// Returns the total number of nodes in the graph.
    fn num_nodes(&self) -> usize;
    /// Returns the total number of directed edges in the graph.
    fn num_edges(&self) -> usize;

    /// Returns a slice representing the start indices of outgoing edges for each node.
    /// Standard CSR (Compressed Sparse Row) format.
    fn first_out(&self) -> &[u32];
    /// Returns a slice of target nodes corresponding to the edges.
    fn head(&self) -> &[u32];

    /// The X-coordinate (longitude) of the given node.
    fn x(&self, node_id: u32) -> f32;
    /// The Y-coordinate (latitude) of the given node.
    fn y(&self, node_id: u32) -> f32;

    /// Returns the range of edge indices originating from node `u`.
    #[inline(always)]
    fn edge_indices(&self, u: usize) -> std::ops::Range<usize> {
        let start = self.first_out()[u] as usize;
        let end = self.first_out()[u + 1] as usize;
        start..end
    }

    /// Returns an iterator over the target nodes of all outgoing edges from node `u`.
    #[inline(always)]
    fn neighbors(&self, u: u32) -> impl Iterator<Item = u32> + '_ {
        let range = self.edge_indices(u as usize);
        self.head()[range].iter().cloned()
    }
}

/// The internal shared data state of the engine.
///
/// This is held behind an `ArcSwap` in [`CchEngine`] so topology rebuilds can be published atomically.
pub struct EngineData<P: ProfileId> {
    /// The shared, metric-independent structural topology.
    pub cch: Arc<Cch>,
    /// The customized metric data specific to each routing profile.
    pub profiles: FxHashMap<P, Arc<RwLock<ProfileData>>>,
    /// Lazily initialized auxiliary indices used to accelerate repeated partial metric updates.
    pub update_ctx: OnceLock<Arc<PartialUpdateContext>>,
}

/// The thread-safe, high-performance CCH routing engine.
///
/// This is the primary entry point for issuing queries and updating graph state.
/// By utilizing `ArcSwap`, topology updates are atomically swapped in. Metric updates mutate
/// the affected profile in place behind a read-write lock, so queries to that profile may wait briefly.
#[derive(Clone)]
pub struct CchEngine<P: ProfileId> {
    /// Lock-free pointer to the current active routing state.
    pub data: Arc<ArcSwap<EngineData<P>>>,
}

impl<P: ProfileId> CchEngine<P> {
    /// Bootstraps a new routing engine from a base graph and a set of initial profiles.
    ///
    /// This performs the initial Nested Dissection ordering, topology contraction,
    /// and heavily parallelized metric customization for all provided profiles.
    pub fn new<'a, G: CchGraph + Sync>(graph: &G, profile_weights: impl IntoIterator<Item = (&'a P, &'a Vec<f32>)>) -> Self {
        let cch = Arc::new(Cch::new(graph));
        let mut profiles = FxHashMap::default();
        for (pid, weights) in profile_weights {
            let (w, s) = cch.customize(&cch.weight_map, &cch.scheduler, weights);
            profiles.insert(
                *pid,
                Arc::new(RwLock::new(ProfileData {
                    input_weights: weights.clone(),
                    weights: w,
                    shortcuts: s,
                })),
            );
        }
        Self {
            data: Arc::new(ArcSwap::from_pointee(EngineData {
                cch,
                profiles,
                update_ctx: OnceLock::new(),
            })),
        }
    }

    /// Performs a zero-downtime update of the edge weights for a specific profile (e.g., live traffic).
    ///
    /// This triggers the parallel Customization phase for the new weights and updates the
    /// profile in place. Queries to the same profile will wait on the profile lock.
    pub fn update_weights(&self, profile_id: P, new_weights: &[f32]) {
        let current = self.data.load();
        let Some(profile) = current.profiles.get(&profile_id) else {
            return;
        };
        let mut profile = profile.write();
        current.cch.recustomize_profile(&mut profile, new_weights);
    }

    /// Performs an incremental metric update for an existing profile.
    ///
    /// `updates` contains `(original_edge_index, new_weight)` pairs in the base graph's CSR edge order.
    /// The engine only recomputes the affected hierarchy arcs and their dependents.
    pub fn update_weights_partial(&self, profile_id: P, updates: &[(usize, f32)]) {
        if updates.is_empty() {
            return;
        }

        let current = self.data.load();
        let Some(profile) = current.profiles.get(&profile_id) else {
            return;
        };
        let update_ctx = current.update_ctx.get_or_init(|| Arc::new(current.cch.build_partial_update_context())).clone();

        let mut profile = profile.write();
        let ProfileData { input_weights, weights, shortcuts } = &mut *profile;
        current.cch.customize_partial(&current.cch.weight_map, update_ctx.as_ref(), input_weights, weights, shortcuts, updates);
    }

    /// Performs a partial structural update to the elimination tree (e.g., minor road closures).
    ///
    /// Instead of rebuilding the entire graph from scratch, this identifies the affected
    /// branches of the elimination tree and only recontracts the modified segments.
    /// It then re-customizes the metrics and hot-swaps the active state.
    pub fn update_topology<G: CchGraph + Sync>(&self, graph: &G, modified_old_nodes: &[u32], new_nodes: &[u32], profile_weights: &FxHashMap<P, Vec<f32>>) {
        let current = self.data.load();
        let mut new_cch_inner = (*current.cch).clone();

        let start_rank = new_cch_inner.update_order(graph, modified_old_nodes, new_nodes);
        let new_topology = new_cch_inner.recontract(graph, start_rank, &current.cch.topology);

        new_cch_inner.weight_map = WeightMap::build(graph, &new_cch_inner.ranks, &new_topology);
        new_cch_inner.scheduler = new_topology.build_scheduler();
        new_cch_inner.topology = new_topology;

        let cch = Arc::new(new_cch_inner);
        let mut profiles = FxHashMap::with_capacity_and_hasher(profile_weights.len(), FxBuildHasher);

        for (pid, weights) in profile_weights {
            let (w, s) = cch.customize(&cch.weight_map, &cch.scheduler, weights);
            profiles.insert(
                *pid,
                Arc::new(RwLock::new(ProfileData {
                    input_weights: weights.clone(),
                    weights: w,
                    shortcuts: s,
                })),
            );
        }

        let update_ctx = OnceLock::new();
        if current.update_ctx.get().is_some() {
            let _ = update_ctx.set(Arc::new(cch.build_partial_update_context()));
        }

        self.data.store(Arc::new(EngineData { cch, profiles, update_ctx }));
    }

    /// Triggers a complete rebuild of the underlying topology and metric data.
    ///
    /// This behaves exactly like `new()`, but atomically swaps the newly built engine state
    /// into the current instance.
    pub fn rebuild_topology<G: CchGraph + Sync>(&self, graph: &G, profile_weights: &FxHashMap<P, Vec<f32>>) {
        let current = self.data.load();
        let new_engine = Self::new(graph, profile_weights);
        let next = new_engine.data.load_full();
        if current.update_ctx.get().is_some() {
            let _ = next.update_ctx.set(Arc::new(next.cch.build_partial_update_context()));
        }
        self.data.store(next);
    }

    /// Executes a shortest-path query between a source and target node for the given profile.
    ///
    /// This method leverages thread-local storage to ensure zero dynamic memory allocation
    /// during the query. It evaluates a bidirectional Dijkstra search over the shortcut hierarchy.
    ///
    /// Returns `None` if the profile doesn't exist or if no path can be found.
    pub fn query_path(&self, profile: P, from: u32, to: u32) -> Option<PathResult> {
        let current = self.data.load();
        let profile_data = current.profiles.get(&profile)?;
        QUERY_BUFFER.with(|q| {
            let mut server = q.borrow_mut();
            let num_nodes = current.cch.ranks.len();
            if server.fw.len() < num_nodes {
                *server = Query::new(num_nodes);
            }
            let profile_data = profile_data.read();
            let result = current.cch.query(&mut server, &profile_data.weights, from, to)?;
            let path = current.cch.query_path(&server, &profile_data.shortcuts, &result);
            Some(PathResult { distance: result.distance, path })
        })
    }
}

/// Profile-specific routing metrics resulting from the Customization phase.
#[derive(Clone)]
pub struct ProfileData {
    /// The input weight vector in the base graph's CSR edge order.
    pub input_weights: Vec<f32>,
    /// The computed weights for the hierarchy arcs.
    pub weights: Weights,
    /// Witness nodes required to unpack shortcut arcs back into original graph node sequences.
    pub shortcuts: Shortcuts,
}

/// The profile-independent, structural core of the Contraction Hierarchy.
///
/// This maintains the node ordering, the elimination tree, and the mapping
/// from the original graph's edges to the hierarchy's arcs.
#[derive(Clone)]
pub struct Cch {
    /// The structural contraction hierarchy containing forward and backward shortcut arcs.
    pub topology: Topology,
    /// Maps original edge indices to their corresponding arc in the `Topology`.
    pub weight_map: WeightMap,
    /// Grouped nodes mapped by tree depth, allowing safe parallelization during customization.
    pub scheduler: Vec<Vec<u32>>,

    /// Maps original node IDs to their rank in the contraction order.
    pub ranks: Vec<u32>,
    /// Maps ranks back to their original node IDs.
    pub order: Vec<u32>,
}

impl Cch {
    /// Constructs a new structural hierarchy from a given graph.
    ///
    /// Computes the node contraction order via METIS, contracts the graph to build
    /// the topology, and maps out the parallel scheduler.
    pub fn new<G: CchGraph + Sync>(graph: &G) -> Self {
        let ranks = Self::get_metis_order(graph);

        let mut order = vec![0; graph.num_nodes()];
        for (node_id, &rank) in ranks.iter().enumerate() {
            order[rank as usize] = node_id as u32;
        }

        let topology = Self::contract(&ranks, graph);
        let weight_map = WeightMap::build(graph, &ranks, &topology);
        let scheduler = topology.build_scheduler();

        Self {
            topology,
            weight_map,
            scheduler,
            ranks,
            order,
        }
    }

    /// Builds a reusable mutable profile for a fixed CCH topology.
    ///
    /// Unlike [`CchEngine`], this is intended for single-owner workflows such as
    /// benchmarking or offline customization, where in-place updates are acceptable.
    pub fn build_profile(&self, original_weights: &[f32]) -> ProfileData {
        let (weights, shortcuts) = self.customize(&self.weight_map, &self.scheduler, original_weights);
        ProfileData {
            input_weights: original_weights.to_vec(),
            weights,
            shortcuts,
        }
    }

    /// Re-customizes an owned profile in place, reusing its allocations.
    pub fn recustomize_profile(&self, profile: &mut ProfileData, new_weights: &[f32]) {
        profile.input_weights.clear();
        profile.input_weights.extend_from_slice(new_weights);
        let (weights, shortcuts) = self.customize(&self.weight_map, &self.scheduler, &profile.input_weights);
        profile.weights = weights;
        profile.shortcuts = shortcuts;
    }

    /// Applies a partial metric update to an owned profile in place.
    pub fn customize_profile_partial(&self, profile: &mut ProfileData, updates: &[(usize, f32)]) {
        let update_ctx = self.build_partial_update_context();
        self.customize_profile_partial_with_context(profile, &update_ctx, updates);
    }

    /// Applies a partial metric update to an owned profile in place using a reusable context.
    pub fn customize_profile_partial_with_context(&self, profile: &mut ProfileData, update_ctx: &PartialUpdateContext, updates: &[(usize, f32)]) {
        self.customize_partial(&self.weight_map, update_ctx, &mut profile.input_weights, &mut profile.weights, &mut profile.shortcuts, updates);
    }

    /// Incrementally updates the node order mapping when new nodes are added or modified.
    ///
    /// Returns the lowest rank affected by the changes, which serves as the starting
    /// point for a partial topology `recontract`.
    pub fn update_order<G: CchGraph + Sync>(&mut self, graph: &G, modified_old_nodes: &[u32], new_nodes: &[u32]) -> u32 {
        let old_num = self.ranks.len();
        let new_num = graph.num_nodes();

        self.ranks.resize(new_num, 0);
        self.order.resize(new_num, 0);

        for (i, &u) in new_nodes.iter().enumerate() {
            let rank = (old_num + i) as u32;
            self.ranks[u as usize] = rank;
            self.order[rank as usize] = u;
        }

        let mut min_dirty_rank = old_num as u32;
        for &u in modified_old_nodes {
            let r = self.ranks[u as usize];
            if r < min_dirty_rank {
                min_dirty_rank = r;
            }
        }

        min_dirty_rank
    }

    /// Returns the mapping from rank to original node ID.
    pub fn get_order(&self) -> &[u32] { &self.order }

    /// Returns the mapping from original node ID to rank.
    pub fn get_ranks(&self) -> &[u32] { &self.ranks }
}