dyntri-core 0.10.2

//! Breadth-first search algorithms
//!
//! If you are only interested in the volume of spherical shells [`bfs_sphere_volumes()`]
//! is significantly faster than the other implementations.
//! If you want to determine the full distance matrix and all spherical shells,
//! use [`bfs()`].
//!
//! The others serve specific purposes for other implementations, but are likely not necessary in
//! general applications.

use std::collections::VecDeque;

use crate::graph::Graph;

/// Perform Breadth First Search up to and including `max_distance` returning the sphere volume.
// TODO: This can probably be made even faster by updating the spherevol only every level
pub fn bfs_sphere_volumes(graph: &Graph, origin: usize, max_distance: u16) -> Vec<usize> {
    let mut visited = vec![false; graph.len()];
    visited[origin] = true;
    // TODO: select a good default for the capacity of the queue
    let mut queue = VecDeque::with_capacity(max_distance as usize * 4);
    queue.push_back((origin, 0));

    let mut spherevol = vec![0; max_distance as usize + 1];
    spherevol[0] = 1;
    loop {
        let Some((node, r)) = queue.pop_front() else {
            break;
        };
        if r >= max_distance {
            // Exit early when max_distance has been reached
            break;
        }
        for &nbr in graph.get_neighbours(node) {
            // Only consider unvisited nodes
            let visited_nbr = &mut visited[nbr];
            if *visited_nbr {
                continue;
            }
            *visited_nbr = true;
            queue.push_back((nbr, r + 1));
            spherevol[r as usize + 1] += 1;
        }
    }
    spherevol
}

/// Results of a breadth-first search
///
/// These results contain the distance list, volume of spherical shells,
/// and the node labels of the spheres.
pub struct BFSResult {
    /// Distance list of the BFS, e.g. the entry at index i
    /// gives the distance of node i to the origin
    /// Note: all undetermined distances are set to [`u16::MAX`]
    pub distance_list: Vec<u16>,
    /// Flattened list of node labels making up spheres
    pub spheres: Vec<usize>,
    /// Indices of the spheres, s.t. `spheres[sphere_idx[radius]..sphere_idx[radius+1]]`
    /// yields the node labels of the sphere at `radius`
    pub sphere_idx: Vec<usize>,
    /// The length of the spheres, must be equal to "diff(sphere_idx)"
    pub sphere_volumes: Vec<usize>,
}

impl BFSResult {
    #[inline]
    pub fn get_sphere(&self, radius: u16) -> &[usize] {
        let r = radius as usize;
        if r + 1 >= self.sphere_idx.len() {
            return &[];
        }
        &self.spheres[self.sphere_idx[r]..self.sphere_idx[r + 1]]
    }

    #[inline]
    pub fn get_sphere_mut(&mut self, radius: u16) -> &mut [usize] {
        let r = radius as usize;
        if r + 1 >= self.sphere_idx.len() {
            return &mut [];
        }
        &mut self.spheres[self.sphere_idx[r]..self.sphere_idx[r + 1]]
    }

    #[inline]
    pub fn get_ball_volumes(&self) -> &[usize] {
        &self.sphere_idx[1..]
    }

    /// Return the number of nodes in the BFS, i.e. the size of the `ball_volume` at `max_radius`
    #[inline]
    pub fn len(&self) -> usize {
        self.spheres.len()
    }

    #[must_use]
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.spheres.len() == 0
    }

    #[inline]
    pub fn origin(&self) -> usize {
        self.spheres[0]
    }

    #[inline]
    pub fn max_radius(&self) -> u16 {
        self.sphere_volumes.len() as u16
    }
}

/// Perform breadth-first search up to and including `max_distance`,
/// returning the full [`BFSResult`]
///
/// Note that the result is only valid for radii up to and including `max_distance`
pub fn bfs(graph: &Graph, origin: usize, max_distance: u16) -> BFSResult {
    let mut distance = vec![u16::MAX; graph.len()];

    distance[origin] = 0;

    let mut spheres = Vec::with_capacity(graph.len());
    spheres.push(origin);
    let mut sphere_sizes = Vec::with_capacity(max_distance as usize + 1);
    sphere_sizes.push(1);
    let mut sphere_idx = Vec::with_capacity(max_distance as usize + 1);
    sphere_idx.push(0);
    sphere_idx.push(1);

    // TODO: select a good default for the capacity of the queue
    let mut queue = VecDeque::with_capacity(graph.len() / 4);
    queue.push_back(origin);

    let mut current_distance = 0;
    loop {
        let Some(node) = queue.pop_front() else {
            sphere_sizes.push(graph.len() - sphere_idx.last().expect("Not empty by constuction"));
            break;
        };

        let r = distance[node];
        if r > current_distance {
            let current_sphere_idx = spheres.len();
            let last_sphere_idx = sphere_idx.last().expect("Not empty by construction");
            sphere_sizes.push(current_sphere_idx - last_sphere_idx);
            sphere_idx.push(current_sphere_idx);
            if r >= max_distance {
                // Exit early when max_distance has been reached
                break;
            }
            current_distance = r;
        }

        for &nbr in graph.get_neighbours(node) {
            // Only consider unvisited nodes
            if distance[nbr] != u16::MAX {
                continue;
            }
            spheres.push(nbr);
            queue.push_back(nbr);
            distance[nbr] = r + 1;
        }
    }

    BFSResult {
        distance_list: distance,
        spheres,
        sphere_idx,
        sphere_volumes: sphere_sizes,
    }
}

/// Results of a breadth-first search containing only a distance list
/// and the nodes in sphere at a requested radius
///
/// See [`BFSResult`] for a more complete overview
pub struct BFSSphere {
    pub distance_list: Vec<u16>,
    pub sphere: Vec<usize>,
}

/// Perform BFS returning [`BFSSphere`] for radius `distance`
///
/// Note this method requires a hint for the size of the sphere. It is designed to
/// be used in double sphere average sphere distance calculations, where another
/// sphere at the same radius is already known and provides a good guess for the
/// size of the other sphere.
/// If this sizehint is a poor guess the implementation will still work fine,
/// but reallocation may be necessary slowing it down.
pub fn bfs_sphere_sizehint(
    graph: &Graph,
    origin: usize,
    distance: u16,
    sizehint: usize,
) -> BFSSphere {
    let mut distlist = vec![u16::MAX; graph.len()];
    distlist[origin] = 0;

    let mut sphere = Vec::with_capacity(sizehint);
    let mut queue = VecDeque::with_capacity(sizehint);
    queue.push_back(origin);

    // Explore until `distance - 1` layer is reached
    loop {
        let Some(node) = queue.pop_front() else {
            break;
        };
        let r = distlist[node];
        if r >= distance - 1 {
            queue.push_back(node);
            break;
        }
        for &nbr in graph.get_neighbours(node) {
            // Only consider unvisited nodes
            let dist_nbr = &mut distlist[nbr];

            if *dist_nbr != u16::MAX {
                continue;
            }
            queue.push_back(nbr);
            *dist_nbr = r + 1;
        }
    }
    // Explore last layer and add found nodes to sphere
    loop {
        let Some(node) = queue.pop_front() else {
            break;
        };
        let r = distlist[node];
        if r >= distance {
            break;
        }
        for &nbr in graph.get_neighbours(node) {
            // Only consider unvisited nodes
            let dist_nbr = &mut distlist[nbr];
            if *dist_nbr != u16::MAX {
                continue;
            }
            *dist_nbr = r + 1;
            sphere.push(nbr);
        }
    }

    BFSSphere {
        distance_list: distlist,
        sphere,
    }
}

/// Perform BFS and return the distance list
///
/// Uses [`bfs_with_dist()`] internally, allocating space for distance list itself
pub fn bfs_dist(graph: &Graph, origin: usize, max_distance: u16) -> Vec<u16> {
    let mut distance = vec![u16::MAX; graph.len()];
    bfs_with_dist(graph, origin, &mut distance, max_distance);

    distance
}

/// Perform BFS to update a distance list
pub fn bfs_with_dist(graph: &Graph, origin: usize, distance: &mut [u16], max_distance: u16) {
    distance[origin] = 0;

    // TODO: select a good default for the capacity of the queue
    let mut queue = VecDeque::with_capacity(max_distance as usize * 4);
    queue.push_back(origin);

    loop {
        let Some(node) = queue.pop_front() else {
            break;
        };
        let r = distance[node];
        if r >= max_distance {
            break;
        }
        for &nbr in graph.get_neighbours(node) {
            // Only consider unvisited nodes
            let dist_nbr = &mut distance[nbr];

            if *dist_nbr != u16::MAX {
                continue;
            }
            queue.push_back(nbr);
            *dist_nbr = r + 1;
        }
    }
}

/// Perform BFS for multiple origins returning the flattened partial distance matrix
///
/// This function is here to be used in other algorithms that can benefit from knowing multiple
/// distance lists at the same time. This would be ideally be replaced by a parallelized function,
/// which the algorithm is set-up for to do easily.
pub fn multi_bfs_dist(graph: &Graph, origins: &[usize], max_distance: u16) -> Vec<u16> {
    let mut distance = vec![u16::MAX; origins.len() * graph.len()];

    multi_bfs_with_dist(graph, origins, max_distance, &mut distance);

    distance
}

/// Perform BFS for multiple origins computing the preallocated flattened partial distance matrix
pub fn multi_bfs_with_dist(
    graph: &Graph,
    origins: &[usize],
    max_distance: u16,
    distance: &mut [u16],
) {
    let nsize = graph.len();
    for (i, &origin) in origins.iter().enumerate() {
        let dist_slice = &mut distance[i * nsize..(i + 1) * nsize];
        bfs_with_dist(graph, origin, dist_slice, max_distance);
    }
}

#[cfg(test)]
mod tests {
    use itertools::assert_equal;

    use super::*;

    #[test]
    fn test_bfs_results() {
        let distance: u16 = 15;
        let size = 400;
        let graph = Graph::cubic2d(size);

        let origin = 0;

        let bfs_result = bfs(&graph, origin, distance);

        assert_eq!(bfs_result.len(), *bfs_result.sphere_idx.last().unwrap());
        assert_eq!(
            bfs_result.sphere_volumes.len() + 1,
            bfs_result.sphere_idx.len()
        );
        assert_eq!(bfs_result.origin(), origin);
    }

    #[test]
    fn test_bfs_sphere() {
        let distance: u16 = 15;
        let size = 400;
        let graph = Graph::cubic2d(size);

        let origin = 0;

        let bfs0 = bfs_sphere_sizehint(&graph, origin, distance, distance as usize * 4);
        let mut sphere0 = bfs0.sphere;
        sphere0.sort();

        let mut bfs1 = bfs(&graph, origin, distance);
        let sphere1 = bfs1.get_sphere_mut(distance);
        sphere1.sort();

        assert_equal(&sphere0, sphere1);

        let svol2 = bfs_sphere_volumes(&graph, origin, distance);
        assert_eq!(sphere0.len(), svol2[distance as usize]);
        assert_equal(svol2, bfs1.sphere_volumes);
    }

    #[test]
    fn test_bfs_spherevol() {
        let origin = 0;
        let max_distance: u16 = 15;
        let size = (max_distance as usize) * 2 + 1;

        let graph = Graph::cubic2d(size);
        let rprofile = bfs_sphere_volumes(&graph, origin, max_distance);
        assert_equal(
            rprofile.into_iter().skip(1),
            (1..=max_distance as usize).map(|r| 4 * r),
        );

        let graph = Graph::cubic3d(size);
        let rprofile = bfs_sphere_volumes(&graph, origin, max_distance);
        assert_equal(
            rprofile.into_iter().skip(1),
            (1..=max_distance as usize).map(|r| 4 * r * r + 2),
        );

        let graph = Graph::cubic4d(size);
        let rprofile = bfs_sphere_volumes(&graph, origin, max_distance);
        assert_equal(
            rprofile.into_iter().skip(1),
            // Note that r(r^2 + 2) is always divisible by 3, so this int division is exact
            (1..=max_distance as usize).map(|r| (8 * r * (r * r + 2)) / 3),
        );

        let graph = Graph::hexagonal(size);
        let rprofile = bfs_sphere_volumes(&graph, origin, max_distance);
        assert_equal(
            rprofile.into_iter().skip(1),
            (1..=max_distance as usize).map(|r| 6 * r),
        );
    }
}