swh-graph-stdlib 11.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

//! Associate information to a subset of nodes

use std::collections::HashMap;

use dsi_progress_logger::{ProgressLog, progress_logger};
use rapidhash::RapidHashMap;
use swh_graph::graph::*;
use swh_graph::graph::{NodeId, SwhForwardGraph};

/// Callbacks for [`MapReduce`]
pub trait MapReducer {
    type Label: Clone;
    type Error;

    /// Returns the label to assign to the given node, independently of its successors, if any
    fn map(&mut self, node: NodeId) -> Result<Option<Self::Label>, Self::Error>;

    /// Given the labels of the children of a node, merge them into a single label.
    ///
    /// Not guaranteed to be called for every node.
    fn reduce<'a, I: Iterator<Item = &'a Self::Label>>(
        &mut self,
        first_label: Self::Label,
        other_labels: I,
    ) -> Result<Option<Self::Label>, Self::Error>
    where
        Self::Label: 'a;

    /// Special-case of [`Self::reduce`] for merging a node's label with its successors' label
    ///
    /// Defaults to calling [`Self::reduce`].
    ///
    /// Not guaranteed to be called for every node.
    fn map_reduce(
        &mut self,
        node: NodeId,
        successors_label: Self::Label,
    ) -> Result<Option<Self::Label>, Self::Error> {
        match self.map(node)? {
            Some(own_label) => self.reduce(own_label, [&successors_label].into_iter()),
            None => Ok(Some(successors_label)),
        }
    }

    /// Called once a node's initial label was computed and merged with its successors'
    ///
    /// Defaults to no-op.
    fn on_node_traversed(
        &mut self,
        _node: NodeId,
        _label: Option<&Self::Label>,
    ) -> Result<(), Self::Error> {
        Ok(())
    }
}

/// 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())`.
///
/// # 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> {
    graph: G,
    num_nodes: usize,
    pub map_reducer: MR,
}

impl<G: SwhForwardGraph + SwhBackwardGraph, MR: MapReducer> MapReduce<G, MR> {
    pub fn new(graph: G, map_reducer: MR) -> Self {
        MapReduce {
            num_nodes: graph.actual_num_nodes().unwrap_or(graph.num_nodes()),
            graph,
            map_reducer,
        }
    }

    pub fn with_num_nodes(mut self, num_nodes: usize) -> Self {
        self.num_nodes = num_nodes;
        self
    }

    pub fn run_in_topological_order(
        &mut self,
        nodes: impl Iterator<Item = NodeId>,
    ) -> Result<(), MR::Error> {
        let mut labels: RapidHashMap<NodeId, MR::Label> = RapidHashMap::default();

        // For each node it, counts its number of direct dependents that still need to be handled
        let mut pending_dependents = HashMap::<NodeId, usize>::new();

        let mut pl = progress_logger!(
            display_memory = true,
            item_name = "node",
            local_speed = true,
            expected_updates = Some(self.num_nodes),
        );

        pl.start("Traversing graph");

        for node in nodes {
            pl.light_update();
            let num_dependents = self.graph.indegree(node);

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

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

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

            while let Some(first_dependency) = dependencies.next() {
                // Get label of the first dependency that has a label
                let first_dependency_label =
                    if 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)
                        pending_dependents.remove(&first_dependency);
                        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
                        *pending_dependents.get_mut(&first_dependency).unwrap() -= 1;
                        labels.get(&first_dependency).cloned()
                    };

                // Merge it with all the others
                if let Some(first_dependency_label) = first_dependency_label {
                    merged_label = self.map_reducer.reduce(
                        first_dependency_label,
                        dependencies.flat_map(|dep|
                            // If 'node' is a revision, then 'dep' is its parent revision
                            labels.get(&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())?;
            if num_dependents > 0 {
                if let Some(label) = label {
                    let previous_label = labels.insert(node, label);
                    assert!(previous_label.is_none(), "{node} was labeled twice");
                }
            } else {
                assert!(!labels.contains_key(&node), "{node} was already labeled");
            }
        }

        pl.done();

        Ok(())
    }
}