hive-router 0.3.2

use std::collections::VecDeque;

use petgraph::{
    graph::NodeIndex,
    visit::{EdgeRef, NodeRef},
    Direction,
};
use tracing::{instrument, trace};

use crate::query_planner::{
    ast::selection_set::selection_items_are_subset_of,
    planner::fetch::{
        error::FetchGraphError,
        fetch_graph::FetchGraph,
        fetch_step_data::{FetchStepData, FetchStepFlags},
        optimize::utils::MergedSteps,
        state::MultiTypeFetchStep,
    },
};

impl FetchGraph<MultiTypeFetchStep> {
    /// When a child has the input identical as the output,
    /// it gets squashed into its parent.
    /// Its children becomes children of the parent.
    #[instrument(level = "trace", skip_all)]
    pub(crate) fn merge_passthrough_child(&mut self) -> Result<(), FetchGraphError> {
        let root_index = self
            .root_index
            .ok_or(FetchGraphError::NonSingleRootStep(0))?;
        // Breadth-First Search (BFS) starting from the root node.
        let mut queue = VecDeque::from([root_index]);
        let mut merged_steps = MergedSteps::default();

        while let Some(parent_index) = queue.pop_front() {
            // Store pairs of sibling nodes that can be merged.
            let mut merges_to_perform: Vec<(NodeIndex, NodeIndex)> = Vec::new();
            let parent_index = merged_steps.resolve(parent_index);

            let children: Vec<_> = self
                .graph
                .neighbors_directed(parent_index, Direction::Outgoing)
                .collect();

            let parent = self.get_step_data(parent_index)?;

            for child_index in children.iter() {
                queue.push_back(*child_index);
                // Add the current child to the queue for further processing (BFS).
                let child = self.get_step_data(*child_index)?;

                if parent.can_merge_passthrough_child(parent_index, *child_index, child, self) {
                    trace!(
                        "passthrough optimization found: merge [{}] <-- [{}]",
                        parent_index.index(),
                        child_index.index()
                    );
                    merges_to_perform.push((parent_index, *child_index));
                }
            }

            for (parent_index, child_index) in merges_to_perform {
                let parent_index = merged_steps.resolve(parent_index);
                let child_index = merged_steps.resolve(child_index);
                // An earlier merge may have joined them already, or changed one of them.
                if parent_index == child_index {
                    continue;
                }
                let parent = self.get_step_data(parent_index)?;
                let child = self.get_step_data(child_index)?;
                if !parent.can_merge_passthrough_child(parent_index, child_index, child, self) {
                    continue;
                }

                perform_passthrough_child_merge(parent_index, child_index, self)?;
                merged_steps.record(child_index, parent_index);
            }
        }

        Ok(())
    }
}

impl FetchStepData<MultiTypeFetchStep> {
    pub(crate) fn can_merge_passthrough_child(
        &self,
        self_index: NodeIndex,
        other_index: NodeIndex,
        other: &Self,
        fetch_graph: &FetchGraph<MultiTypeFetchStep>,
    ) -> bool {
        if self_index == other_index {
            return false;
        }

        if other
            .flags
            .contains(FetchStepFlags::USED_FOR_TYPE_CONDITION)
        {
            return false;
        }

        // if the `other` FetchStep has a single parent and it's `this` FetchStep
        if fetch_graph.parents_of(other_index).count() != 1 {
            return false;
        }

        if fetch_graph.parents_of(other_index).next().unwrap().source() != self_index {
            return false;
        }

        // Every type it fetches has to be in its input already, not just one of them.
        other
            .output
            .iter_selections()
            .all(|(output_def_name, output_selections)| {
                other
                    .input
                    .selections_for_definition(output_def_name)
                    .is_some_and(|input_selections| {
                        selection_items_are_subset_of(
                            &input_selections.items,
                            &output_selections.items,
                        )
                    })
            })
    }
}

#[instrument(level = "trace", skip_all)]
fn perform_passthrough_child_merge(
    self_index: NodeIndex,
    other_index: NodeIndex,
    fetch_graph: &mut FetchGraph<MultiTypeFetchStep>,
) -> Result<(), FetchGraphError> {
    let (me, other) = fetch_graph.get_pair_of_steps_mut(self_index, other_index)?;
    let path = other.response_path.slice_from(me.response_path.len());

    trace!(
        "merging fetch steps [{}] + [{}] at path {}",
        self_index.index(),
        other_index.index(),
        path
    );

    me.output.migrate_from_another(&other.output, &path)?;

    let mut children_indexes: Vec<NodeIndex> = vec![];
    let mut parents_indexes: Vec<NodeIndex> = vec![];
    for edge_ref in fetch_graph.children_of(other_index) {
        children_indexes.push(edge_ref.target().id());
    }

    for edge_ref in fetch_graph.parents_of(other_index) {
        // We ignore self_index
        if edge_ref.source().id() != self_index {
            parents_indexes.push(edge_ref.source().id());
        }
    }

    // Replace parents:
    // 1. Add self -> child
    for child_index in children_indexes {
        trace!(
            "migrating parent [{}] to child [{}]",
            self_index.index(),
            child_index.index()
        );

        fetch_graph.connect(self_index, child_index);
    }

    // 2. Add parent -> self
    for parent_index in parents_indexes {
        trace!(
            "linking parent [{}] to self [{}]",
            parent_index.index(),
            self_index.index()
        );

        fetch_graph.connect(parent_index, self_index);
    }

    // 3. Drop other -> child and parent -> other
    trace!("removing other [{}] from graph", other_index.index());
    fetch_graph.remove_step(other_index);

    Ok(())
}