swh-graph-stdlib 13.0.0

// Copyright (C) 2026  The Software Heritage developers
// See the AUTHORS file at the top-level directory of this distribution
// License: GNU General Public License version 3, or any later version
// See top-level LICENSE file for more information

//! Implements labels propagation through a topological sort

use std::borrow::Borrow;

use anyhow::{Result, bail, ensure};
use dsi_progress_logger::{ProgressLog, progress_logger};
use rapidhash::RapidHashMap;
use smallvec::SmallVec;
use swh_graph::graph::*;
use swh_graph::graph::{NodeId, SwhForwardGraph};

use super::MapReducer;
use super::labels::{
    DenseLabels, Labels, SparseLabels, StridableLabel, StriddenLabels, StriddenLabelsConfig,
};

/// Builder for [`MapReduce`]
pub struct MapReduceBuilder<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer, L: Labels> {
    graph: G,
    num_nodes: usize,
    cheap_clones: bool,
    keep_labels: bool,
    pub map_reducer: MR,
    labels_config: L::Config,
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer<Label: Sized>>
    MapReduceBuilder<G, MR, SparseLabels<MR::Label>>
{
    /// Stores labels in a HashMap. This is the best when labeling the history-hosting layer
    ///
    /// This improves memory usage at the expense of runtime and CPU use.
    pub fn new_sparse(graph: G, map_reducer: MR) -> Self {
        MapReduceBuilder {
            num_nodes: graph.actual_num_nodes().unwrap_or(graph.num_nodes()),
            graph,
            cheap_clones: false,
            keep_labels: false,
            map_reducer,
            labels_config: (),
        }
    }
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer<Label: Default + Clone + Sized>>
    MapReduceBuilder<G, MR, DenseLabels<MR::Label>>
{
    /// Stores labels in an array instead of a HashMap. This is the best when labeling the
    /// directory layer or the whole graph.
    ///
    /// This improves runtime and CPU use at the expense of memory.
    ///
    /// This should probably be used only if:
    ///
    /// * labels are small, or
    /// * computed labels are sparse (wrt. `graph.num_nodes()`)
    ///   **and** `None::<MR::Label>` is small
    pub fn new_dense(graph: G, map_reducer: MR) -> Self {
        MapReduceBuilder {
            num_nodes: graph.actual_num_nodes().unwrap_or(graph.num_nodes()),
            graph,
            cheap_clones: false,
            keep_labels: false,
            map_reducer,
            labels_config: (),
        }
    }
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer<Label: StridableLabel>>
    MapReduceBuilder<G, MR, StriddenLabels<MR::Label>>
{
    /// Specialized variant of [`new_labels`](Self::new_dense) for non-[`Sized`](Sized) labels.
    ///
    /// All labels must have the same length, but it can be computed at runtime.
    ///
    /// Like [`Self::new_dense`], this improves runtime and CPU use at the expense of memory,
    /// but less than boxing the values would.
    ///
    /// This should only be used if labels are small.
    ///
    /// `num_words` is the length of the `[Label::Word]` slice needed to store a label.
    pub fn new_stridden(graph: G, map_reducer: MR, num_words: usize) -> Self
    where
        MR::Label: StridableLabel + ToOwned,
    {
        MapReduceBuilder {
            num_nodes: graph.actual_num_nodes().unwrap_or(graph.num_nodes()),
            graph,
            cheap_clones: false,
            keep_labels: false,
            map_reducer,
            labels_config: StriddenLabelsConfig { num_words },
        }
    }
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer, L: Labels<Label = MR::Label>>
    MapReduceBuilder<G, MR, L>
{
    pub fn num_nodes(mut self, num_nodes: usize) -> Self {
        self.num_nodes = num_nodes;
        self
    }

    /// Tunes the algorithm to assume labels are cheap to clone.
    ///
    /// This is probably `true` if and only if the labels implement Copy.
    /// Setting this to `true` does not imply they are cheap to move.
    ///
    /// Defaults to `false`.
    pub fn cheap_clones(mut self, cheap_clones: bool) -> Self {
        self.cheap_clones = cheap_clones;
        self
    }

    /// Whether the algorithm should keep labels in its store in order to return them at the end
    ///
    /// This is incompatible with `cheap_clones(false)` (the default).
    ///
    /// This consumes extra memory, except when `with_labels_array()` and labels do not contain
    /// heap-allocated data.
    ///
    /// Defaults to `false`.
    pub fn keep_labels(mut self, keep_labels: bool) -> Self {
        self.keep_labels = keep_labels;
        self
    }

    pub fn build(self) -> Result<MapReduce<G, MR, L>> {
        let Self {
            graph,
            num_nodes,
            cheap_clones,
            keep_labels,
            map_reducer,
            labels_config,
        } = self;
        let pop_labels = match (cheap_clones, keep_labels) {
            (true, true) => false, // cheap to clone and we need to keep them
            (true, false) => true, // cheap to clone, but popping saves memory
            (false, true) => bail!("MapReduce cannot both keep labels and avoid expensive clones"),
            (false, false) => true, // expensive to clone, so popping saves a clone
        };
        log::info!("Allocating labels...");
        let labels = L::new(graph.num_nodes(), labels_config);
        Ok(MapReduce {
            graph,
            num_nodes,
            pop_labels,
            map_reducer,
            labels,
            pending_dependents: Default::default(),
        })
    }
}

/// Associates labels to nodes in the graph using successor nodes' labels and "bubbling up"
///
/// Use [`swh_graph::views::Subgraph`] to select the set of nodes to run this on.
/// For example, to avoid content and directory nodes (which are typically much slower to process),
/// use `Subgraph::with_node_constraint("rev,rel,snp,ori".parse().unwrap())`.
///
/// Built from [`MapReduceBuilder`]
///
/// # Example
///
/// For example, with this graph:
///
/// ```text
///      - 3
///     /
///   <-
/// 1 <--- 4 <--+
///             +--- 6
/// 2 <--- 5 <--+
/// ```
///
/// We call 1 a successor of 3, consistent with swh-graph's terminology, even though MapReducer
/// propagates labels in the other direction.
///
/// we would:
///
/// * compute label of 1
/// * compute label of 2
/// * compute label of 2 and merge it with 1's
/// * compute label of 4 and merge it with 1's
/// * compute label of 5 and merge it with 2's
/// * compute label of 6 and merge it with 4's and 5's
pub struct MapReduce<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer, L> {
    graph: G,
    num_nodes: usize,
    pop_labels: bool,
    pub map_reducer: MR,
    labels: L,
    /// For each node, counts its number of direct dependents that still need to be handled.
    ///
    /// Unused if pop_labels is false
    pending_dependents: RapidHashMap<NodeId, usize>,
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer, L: Labels<Label = MR::Label>>
    MapReduce<G, MR, L>
{
    /// Runs the configured [`MapReducer`] sequentially on all nodes in the graph.
    ///
    /// `nodes` must be an iterator of all nodes in topological order,
    /// eg. returned by [`GenerationsReader::iter_nodes`](https://docs.rs/swh_graph_topology/latest/swh_graph_topology/generations/struct.GenerationsReader.html#method.iter_nodes)
    /// with `.map(|(_depth, node)| node`).
    pub fn run_in_topological_order(
        &mut self,
        nodes: impl Iterator<Item = NodeId>,
    ) -> Result<(), MR::Error> {
        if self.pop_labels {
            self.run_in_topological_order_with_popped_labels(nodes)
        } else {
            self.run_in_topological_order_with_kept_labels(nodes)
        }
    }

    /// A single step of [`Self::run_in_topological_order`], in case the caller does not have an
    /// iterator of nodes
    pub fn push_node(&mut self, node: NodeId) -> Result<(), MR::Error> {
        if self.pop_labels {
            self.push_node_with_popped_labels(node)
        } else {
            self.push_node_with_kept_labels(node)
        }
    }

    /// Returns every node's labels, if `MapReduceBuilder::keep_labels` was set to true.
    pub fn labels(&self) -> Result<&L> {
        ensure!(
            !self.pop_labels,
            "MapReducer::labels() is not available as MapReduceBuilder::keep_labels() was not set to true"
        );
        Ok(&self.labels)
    }

    /// Returns every node's labels, if `MapReduceBuilder::keep_labels` was set to true.
    pub fn take_labels(self) -> Result<L> {
        ensure!(
            !self.pop_labels,
            "MapReducer::labels() is not available as MapReduceBuilder::keep_labels() was not set to true"
        );
        Ok(self.labels)
    }

    /// Implementation of [`run_in_topological_order`] optimized for labels that are cheap to clone
    /// or move
    fn run_in_topological_order_with_kept_labels(
        &mut self,
        nodes: impl Iterator<Item = NodeId>,
    ) -> Result<(), MR::Error> {
        let mut pl = progress_logger!(
            display_memory = true,
            item_name = "node",
            local_speed = true,
            expected_updates = Some(self.num_nodes),
        );

        pl.start("Traversing graph in topological order...");

        for node in nodes {
            pl.light_update();

            self.push_node_with_kept_labels(node)?;
        }

        pl.done();

        Ok(())
    }

    #[inline]
    fn push_node_with_kept_labels(&mut self, node: NodeId) -> Result<(), MR::Error> {
        let mut dependencies = self.graph.successors(node).into_iter();

        // get label of first dependencies
        let first_dependency_label = (&mut dependencies)
            .flat_map(|dependency| self.labels.get(dependency).map(ToOwned::to_owned))
            .next();

        // Merge other dependencies' labels with it
        let label: Option<<MR::Label as ToOwned>::Owned> = match first_dependency_label {
            Some(label) => {
                self.map_reducer.reduce(
                    label,
                    dependencies.flat_map(|dep|
                            // If 'node' is a revision, then 'dep' is its parent revision
                            self.labels.get(dep)),
                )?
            }
            None => {
                assert!(
                    dependencies.next().is_none(),
                    "first_dependency_label is None, but not all dependencies were consumed"
                );
                None
            }
        };

        // Merge this node's label with them
        let label = match label {
            Some(label) => self.map_reducer.map_reduce(node, label)?,
            None => self.map_reducer.map(node)?,
        };

        self.map_reducer
            .on_node_traversed(node, label.as_ref().map(|l| l.borrow()))?;
        if let Some(label) = label {
            let previous_label = self.labels.insert(node, label);
            assert!(previous_label.is_none(), "{node} was labeled twice");
        }

        Ok(())
    }

    /// Implementation of [`run_in_topological_order`] optimized for labels that are expensive to
    /// clone but cheap to move
    fn run_in_topological_order_with_popped_labels(
        &mut self,
        nodes: impl Iterator<Item = NodeId>,
    ) -> Result<(), MR::Error> {
        let mut pl = progress_logger!(
            display_memory = true,
            item_name = "node",
            local_speed = true,
            expected_updates = Some(self.num_nodes),
        );

        pl.start("Traversing graph in topological order...");

        for node in nodes {
            pl.light_update();

            self.push_node_with_popped_labels(node)?;
        }

        pl.done();

        debug_assert!(
            self.labels.is_empty(),
            "run_in_topological_order_with_popped_labels ended without clearing its labels store"
        );

        Ok(())
    }

    #[inline]
    fn push_node_with_popped_labels(&mut self, node: NodeId) -> Result<(), MR::Error> {
        let num_dependents = self.graph.indegree(node);

        if num_dependents > 0 {
            self.pending_dependents.insert(node, num_dependents);
        }

        let mut dependencies = self.graph.successors(node).into_iter();

        let mut merged_label: Option<<MR::Label as ToOwned>::Owned> = None;

        while let Some(first_dependency) = dependencies.next() {
            // Get label of the first dependency that has a label
            let first_dependency_label = if self.pending_dependents.get(&first_dependency)
                == Some(&1)
            {
                // Reuse the dependency's set of contributors.
                //
                // This saves a potentially expensive clone in the tight loop.
                // When working with the revision graph, this branch is almost always taken
                // because most revisions have a single parent (ie. single dependency)
                self.pending_dependents.remove(&first_dependency);
                self.labels.remove(first_dependency)
            } else {
                // Dependency is not yet ready to be popped because it has other dependents
                // to be visited.  Copy its contributor set
                let pending_dependants = self.pending_dependents.get_mut(&first_dependency).unwrap_or_else(|| panic!("Node {node} depends on node {first_dependency} but the latter's label was not computed (yet?). Check the topological order is complete."));
                *pending_dependants -= 1;
                self.labels.get(first_dependency).map(
                    |l: &MR::Label| -> <MR::Label as std::borrow::ToOwned>::Owned { l.to_owned() },
                )
            };

            // Merge it with all the others
            if let Some(first_dependency_label) = first_dependency_label {
                let mut dependencies_to_remove = SmallVec::<[_; 1]>::new();
                merged_label = self.map_reducer.reduce(
                    first_dependency_label,
                    dependencies.flat_map(|dep| {
                        *self.pending_dependents.get_mut(&dep).unwrap() -= 1;
                        if *self.pending_dependents.get(&dep).unwrap() == 0 {
                            dependencies_to_remove.push(dep);
                        }
                        // If 'node' is a revision, then 'dep' is its parent revision
                        self.labels.get(dep)
                    }),
                )?;

                // Clean up deps that have no more pending dependents
                for dep in dependencies_to_remove {
                    self.pending_dependents.remove(&dep);
                    self.labels.remove(dep);
                }
                break;
            }
        }

        let label = match merged_label {
            Some(merged_label) => self.map_reducer.map_reduce(node, merged_label)?,
            None => self.map_reducer.map(node)?,
        };
        self.map_reducer
            .on_node_traversed(node, label.as_ref().map(Borrow::borrow))?;
        if num_dependents > 0 {
            if let Some(label) = label {
                let previous_label = self.labels.insert(node, label);
                assert!(previous_label.is_none(), "{node} was labeled twice");
            }
        } else {
            assert!(
                !self.labels.contains_key(node),
                "{node} was already labeled"
            );
        }

        Ok(())
    }
}