cch_rs/cch/

customize.rs

//! Parallelized metric customization phase.
//!
//! Customization is the "C" in CCH. Once the structural topology is built, this module
//! takes a specific set of edge weights (e.g., travel times for a specific vehicle profile)
//! and propagates them through the shortcut arcs.
//!
//! This phase is heavily optimized and parallelized using `rayon`. It utilizes a lock-free
//! approach with atomic spin-locks (`AtomicBool`) and thread-local workspaces to evaluate
//! the elimination tree levels concurrently.

use std::cell::RefCell;
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::sync::atomic::{AtomicBool, Ordering};

use rayon::prelude::*;

use crate::cch::contract::Topology;
use crate::cch::{Cch, CchGraph};

thread_local! {
    static WORKSPACE: RefCell<Vec<WeightPair>> = const { RefCell::new(Vec::new()) };
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct WeightPair {
    up: f32,
    down: f32,
}
#[derive(Clone)]
pub struct Weights {
    pub up: Vec<f32>,
    pub down: Vec<f32>,
}
#[repr(C)]
#[derive(Clone, Copy)]
pub struct WitnessPair {
    pub up: u32,
    pub down: u32,
}
#[derive(Clone)]
pub struct Shortcuts {
    pub middles: Vec<WitnessPair>,
}

#[derive(Debug, Clone)]
pub struct WeightMap {
    pub up_map: Vec<u32>,
    pub down_map: Vec<u32>,
}
impl WeightMap {
    pub fn build<G: CchGraph>(graph: &G, ranks: &[u32], topology: &Topology) -> Self {
        let (mut up_map, mut down_map) = (vec![u32::MAX; graph.num_edges()], vec![u32::MAX; graph.num_edges()]);
        (0..graph.num_nodes() as u32).for_each(|u| {
            let u_rank = ranks[u as usize];
            graph.edge_indices(u as usize).for_each(|edge_idx| {
                let v_rank = ranks[graph.head()[edge_idx] as usize];
                let (from_r, to_r) = if u_rank < v_rank { (u_rank, v_rank) } else { (v_rank, u_rank) };
                if let Some(arc_id) = unsafe { topology.find_edge_id(from_r, to_r) } {
                    let target = if u_rank < v_rank { &mut up_map } else { &mut down_map };
                    target[edge_idx] = arc_id;
                }
            });
        });
        Self { up_map, down_map }
    }
}
impl WeightPair {
    const INFINITY: Self = Self {
        up: f32::INFINITY,
        down: f32::INFINITY,
    };
    #[inline(always)]
    fn is_infinite(&self) -> bool { self.up == f32::INFINITY && self.down == f32::INFINITY }
    #[inline(always)]
    fn reset(&mut self) {
        self.up = f32::INFINITY;
        self.down = f32::INFINITY;
    }
}

#[derive(Clone, Copy)]
struct RawSlice<T>(*mut T);
unsafe impl<T> Send for RawSlice<T> {}
unsafe impl<T> Sync for RawSlice<T> {}
impl<T> RawSlice<T> {
    #[inline(always)]
    unsafe fn get(self, i: usize) -> *mut T { unsafe { self.0.add(i) } }
}

struct MinArcQueue {
    heap: BinaryHeap<Reverse<u32>>,
    queued: Vec<bool>,
}
impl MinArcQueue {
    fn new(arc_count: usize) -> Self {
        Self {
            heap: BinaryHeap::new(),
            queued: vec![false; arc_count],
        }
    }

    fn push(&mut self, arc: u32) {
        let slot = &mut self.queued[arc as usize];
        if !*slot {
            *slot = true;
            self.heap.push(Reverse(arc));
        }
    }

    fn pop(&mut self) -> Option<u32> {
        self.heap.pop().map(|Reverse(arc)| {
            self.queued[arc as usize] = false;
            arc
        })
    }
}

#[derive(Clone)]
pub struct PartialUpdateContext {
    tail: Vec<u32>,
    first_in: Vec<u32>,
    in_arc: Vec<u32>,
    up_first: Vec<u32>,
    up_src: Vec<u32>,
    down_first: Vec<u32>,
    down_src: Vec<u32>,
}
impl PartialUpdateContext {
    pub fn build(topology: &Topology, mapper: &WeightMap) -> Self {
        let arc_count = topology.head.len();
        let mut tail = vec![0; arc_count];
        topology.first_out.windows(2).enumerate().for_each(|(u, w)| tail[w[0] as usize..w[1] as usize].fill(u as u32));

        let mut first_in = vec![0; topology.first_out.len()];
        topology.head.iter().for_each(|&v| first_in[v as usize + 1] += 1);
        (1..first_in.len()).for_each(|i| first_in[i] += first_in[i - 1]);

        let mut in_arc = vec![0; arc_count];
        let mut next_in = first_in[..first_in.len() - 1].to_vec();
        topology.head.iter().enumerate().for_each(|(arc, &v)| {
            in_arc[next_in[v as usize] as usize] = arc as u32;
            next_in[v as usize] += 1;
        });

        let (up_first, up_src) = Self::build_source_index(&mapper.up_map, arc_count);
        let (down_first, down_src) = Self::build_source_index(&mapper.down_map, arc_count);

        Self {
            tail,
            first_in,
            in_arc,
            up_first,
            up_src,
            down_first,
            down_src,
        }
    }

    fn build_source_index(map: &[u32], arc_count: usize) -> (Vec<u32>, Vec<u32>) {
        let mut f = vec![0; arc_count + 1];
        map.iter().filter(|&&a| a != u32::MAX).for_each(|&a| f[a as usize + 1] += 1);
        (1..f.len()).for_each(|i| f[i] += f[i - 1]);
        let (mut s, mut n) = (vec![0; f[arc_count] as usize], f.clone());
        map.iter().enumerate().filter(|&(_, &a)| a != u32::MAX).for_each(|(i, &a)| {
            s[n[a as usize] as usize] = i as u32;
            n[a as usize] += 1;
        });
        (f, s)
    }

    #[inline(always)]
    fn intersect_by<K: Ord, F: FnMut(usize, usize)>(n1: usize, n2: usize, mut k1: impl FnMut(usize) -> K, mut k2: impl FnMut(usize) -> K, mut f: F) {
        let (mut i, mut j) = (0, 0);
        while i < n1 && j < n2 {
            match k1(i).cmp(&k2(j)) {
                std::cmp::Ordering::Less => i += 1,
                std::cmp::Ordering::Greater => j += 1,
                std::cmp::Ordering::Equal => {
                    f(i, j);
                    i += 1;
                    j += 1;
                },
            }
        }
    }

    pub fn for_each_upper<F: FnMut(u32, u32, u32)>(&self, topo: &Topology, arc: u32, mut f: F) {
        let (x, y) = (self.tail[arc as usize] as usize, topo.head[arc as usize] as usize);
        let (xo, yo) = (arc as usize + 1, topo.first_out[y] as usize);
        Self::intersect_by(
            (topo.first_out[x + 1] as usize).saturating_sub(xo),
            (topo.first_out[y + 1] as usize).saturating_sub(yo),
            |i| topo.head[xo + i],
            |j| topo.head[yo + j],
            |i, j| f(arc, (xo + i) as u32, (yo + j) as u32),
        );
    }

    pub fn for_each_intermediate<F: FnMut(u32, u32, u32)>(&self, topo: &Topology, arc: u32, mut f: F) {
        let (x, y) = (self.tail[arc as usize] as usize, topo.head[arc as usize] as usize);
        let (xo, yo) = (topo.first_out[x] as usize, self.first_in[y] as usize);
        Self::intersect_by(
            (arc as usize).saturating_sub(xo),
            (self.first_in[y + 1] as usize).saturating_sub(yo),
            |i| topo.head[xo + i],
            |j| self.tail[self.in_arc[yo + j] as usize],
            |i, j| f((xo + i) as u32, arc, self.in_arc[yo + j]),
        );
    }

    pub fn for_each_lower<F: FnMut(u32, u32, u32)>(&self, topo: &Topology, arc: u32, mut f: F) {
        let (x, y) = (self.tail[arc as usize] as usize, topo.head[arc as usize] as usize);
        let (xo, yo) = (self.first_in[x] as usize, self.first_in[y] as usize);
        Self::intersect_by(
            (self.first_in[x + 1] as usize).saturating_sub(xo),
            (self.first_in[y + 1] as usize).saturating_sub(yo),
            |i| self.tail[self.in_arc[xo + i] as usize],
            |j| self.tail[self.in_arc[yo + j] as usize],
            |i, j| f(self.in_arc[xo + i], self.in_arc[yo + j], arc),
        );
    }
}

impl Cch {
    pub fn build_partial_update_context(&self) -> PartialUpdateContext { PartialUpdateContext::build(&self.topology, &self.weight_map) }

    #[inline]
    fn ensure_workspace_size(&self, ws: &mut Vec<WeightPair>, required_size: usize) {
        if ws.len() < required_size {
            ws.resize(required_size, WeightPair::INFINITY);
        }
    }

    #[inline(always)]
    fn initial_arc_weight(&self, ctx: &PartialUpdateContext, input_weights: &[f32], arc: u32) -> WeightPair {
        let mut pair = WeightPair::INFINITY;
        ctx.up_src[ctx.up_first[arc as usize] as usize..ctx.up_first[arc as usize + 1] as usize]
            .iter()
            .for_each(|&e| pair.up = pair.up.min(input_weights[e as usize]));
        ctx.down_src[ctx.down_first[arc as usize] as usize..ctx.down_first[arc as usize + 1] as usize]
            .iter()
            .for_each(|&e| pair.down = pair.down.min(input_weights[e as usize]));
        pair
    }

    fn recompute_arc(&self, ctx: &PartialUpdateContext, input_weights: &[f32], weights: &mut Weights, shortcuts: &mut Shortcuts, arc: u32) -> WeightPair {
        let mut pair = self.initial_arc_weight(ctx, input_weights, arc);
        let mut middle = WitnessPair { up: u32::MAX, down: u32::MAX };

        ctx.for_each_lower(&self.topology, arc, |bottom_arc, mid_arc, _| {
            let bottom = bottom_arc as usize;
            let mid = mid_arc as usize;
            let witness = ctx.tail[bottom];

            let next_up = weights.down[bottom] + weights.up[mid];
            if next_up < pair.up {
                pair.up = next_up;
                middle.up = witness;
            }

            let next_down = weights.up[bottom] + weights.down[mid];
            if next_down < pair.down {
                pair.down = next_down;
                middle.down = witness;
            }
        });

        weights.up[arc as usize] = pair.up;
        weights.down[arc as usize] = pair.down;
        shortcuts.middles[arc as usize] = middle;
        pair
    }

    pub fn customize(&self, mapper: &WeightMap, scheduler: &[Vec<u32>], original_weights: &[f32]) -> (Weights, Shortcuts) {
        let num_arcs = self.topology.head.len();
        let mut weights = vec![WeightPair::INFINITY; num_arcs];
        let mut middles = vec![WitnessPair { up: u32::MAX, down: u32::MAX }; num_arcs];

        let weights_ptr = RawSlice(weights.as_mut_ptr());
        let witness_ptr = RawSlice(middles.as_mut_ptr());

        original_weights.iter().enumerate().for_each(|(i, &weight)| unsafe {
            let set_weight = |idx: u32, is_up: bool| {
                if idx != u32::MAX {
                    let p = &mut *weights_ptr.get(idx as usize);
                    if is_up {
                        p.up = p.up.min(weight);
                    } else {
                        p.down = p.down.min(weight);
                    }
                }
            };
            set_weight(mapper.up_map[i], true);
            set_weight(mapper.down_map[i], false);
        });

        let num_nodes = self.ranks.len();
        let locks: Vec<AtomicBool> = (0..num_nodes).map(|_| AtomicBool::new(false)).collect();

        scheduler.iter().for_each(|l| l.par_iter().for_each(|&u| self.relax_node(u, weights_ptr, witness_ptr, &locks)));

        (
            Weights {
                up: weights.iter().map(|w| w.up).collect(),
                down: weights.iter().map(|w| w.down).collect(),
            },
            Shortcuts { middles },
        )
    }

    pub fn customize_partial(&self, mapper: &WeightMap, ctx: &PartialUpdateContext, input_weights: &mut [f32], weights: &mut Weights, shortcuts: &mut Shortcuts, updates: &[(usize, f32)]) {
        if updates.is_empty() {
            return;
        }

        let mut queue = MinArcQueue::new(self.topology.head.len());
        updates.iter().for_each(|&(edge_idx, new_weight)| {
            assert!(edge_idx < input_weights.len(), "Input edge index {edge_idx} out of bounds");
            input_weights[edge_idx] = new_weight;

            let up_arc = mapper.up_map[edge_idx];
            if up_arc != u32::MAX {
                queue.push(up_arc);
            }

            let down_arc = mapper.down_map[edge_idx];
            if down_arc != u32::MAX {
                queue.push(down_arc);
            }
        });

        while let Some(arc) = queue.pop() {
            let old_up = weights.up[arc as usize];
            let old_down = weights.down[arc as usize];
            let new_pair = self.recompute_arc(ctx, input_weights, weights, shortcuts, arc);

            if old_up == new_pair.up && old_down == new_pair.down {
                continue;
            }

            ctx.for_each_intermediate(&self.topology, arc, |bottom_arc, _, top_arc| {
                let bottom = bottom_arc as usize;
                let top = top_arc as usize;
                if weights.down[bottom] + old_up == weights.up[top]
                    || weights.up[bottom] + old_down == weights.down[top]
                    || weights.down[bottom] + new_pair.up < weights.up[top]
                    || weights.up[bottom] + new_pair.down < weights.down[top]
                {
                    queue.push(top_arc);
                }
            });

            ctx.for_each_upper(&self.topology, arc, |_, mid_arc, top_arc| {
                let mid = mid_arc as usize;
                let top = top_arc as usize;
                if weights.up[mid] + old_down == weights.up[top]
                    || weights.down[mid] + old_up == weights.down[top]
                    || weights.up[mid] + new_pair.down < weights.up[top]
                    || weights.down[mid] + new_pair.up < weights.down[top]
                {
                    queue.push(top_arc);
                }
            });
        }

        #[cfg(debug_assertions)]
        {
            (0..self.topology.head.len() as u32).for_each(|arc| {
                ctx.for_each_lower(&self.topology, arc, |bottom_arc, mid_arc, top_arc| {
                    let bottom = bottom_arc as usize;
                    let mid = mid_arc as usize;
                    let top = top_arc as usize;
                    assert!(weights.up[top] <= weights.down[bottom] + weights.up[mid]);
                    assert!(weights.down[top] <= weights.up[bottom] + weights.down[mid]);
                });
            });
        }
    }

    #[inline(always)]
    fn relax_node(&self, u: u32, weights_ptr: RawSlice<WeightPair>, witnesses_ptr: RawSlice<WitnessPair>, locks: &[AtomicBool]) {
        WORKSPACE.with(|cell| {
            let mut ws = cell.borrow_mut();
            self.ensure_workspace_size(&mut ws, self.ranks.len());

            unsafe {
                let ws_ptr = ws.as_mut_ptr();
                let u_rng = self.topology.get_range(u as usize);
                u_rng.clone().for_each(|arc| *ws_ptr.add(self.topology.get_head(arc)) = *weights_ptr.get(arc));

                for uv_arc in u_rng.clone() {
                    let v = self.topology.get_head(uv_arc);
                    let uv_pair = *ws_ptr.add(v);
                    if uv_pair.is_infinite() {
                        continue;
                    }

                    let mut backoff = 1;
                    while locks[v].compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
                        for _ in 0..backoff {
                            std::hint::spin_loop();
                        }
                        if backoff < 64 {
                            backoff <<= 1;
                        }
                    }

                    for vw_arc in self.topology.get_range(v) {
                        let w = self.topology.get_head(vw_arc);
                        let uw_pair = *ws_ptr.add(w);

                        let (target_weight, target_witness) = (weights_ptr.get(vw_arc), witnesses_ptr.get(vw_arc));
                        let (n_up, n_down) = (uv_pair.down + uw_pair.up, uv_pair.up + uw_pair.down);

                        if n_up < (*target_weight).up {
                            (*target_weight).up = n_up;
                            (*target_witness).up = u;
                        }
                        if n_down < (*target_weight).down {
                            (*target_weight).down = n_down;
                            (*target_witness).down = u;
                        }
                    }
                    locks[v].store(false, Ordering::Release);
                }
                u_rng.into_iter().for_each(|arc| (*ws_ptr.add(self.topology.get_head(arc))).reset());
            }
        });
    }
}