hive-router 0.3.6

use petgraph::graph::{EdgeIndex, NodeIndex};
use std::rc::Rc;
use tracing::{instrument, trace};

use crate::query_planner::{
    ast::{arguments::ArgumentsMap, merge_path::Condition},
    graph::{error::GraphError, Graph},
    planner::walker::path::{ContextRequirement, OperationPath, PathSegment, SelectionAttributes},
};

use super::query_tree::QueryTree;

/// Represents the position of a mutation root field,
/// in the operation's selection set.
pub type MutationFieldPosition = Option<usize>;

#[derive(Debug, Clone)]
pub struct QueryTreeNode {
    /// The underlying graph node this query tree node corresponds to
    pub node_index: NodeIndex,
    /// The edge from the parent QueryTreeNode that led to this node (null for root)
    pub edge_from_parent: Option<EdgeIndex>,
    /// Nodes required to execute the move
    pub requirements: Vec<Rc<QueryTreeNode>>,
    /// Only set on nodes that deal with `@fromContext`, the rest pay a pointer for it.
    pub context: Option<Box<NodeContext>>,
    pub children: Vec<Rc<QueryTreeNode>>,
    pub selection_attributes: Option<SelectionAttributes>,
    pub condition: Option<Condition>,
    /// Distinguishes nodes originating from different top-level mutation fields.
    pub mutation_field_position: MutationFieldPosition,
}

#[derive(Debug, Clone, Default)]
pub struct NodeContext {
    /// Trees that fetch the context values for the fields below this node
    pub requirements: Vec<Rc<QueryTreeNode>>,
    /// The contextual argument a context requirement tree is for
    pub coordinate: Option<(EdgeIndex, usize)>,
    /// Context values this node's field needs
    pub resolved_requirements: Vec<Rc<ContextRequirement>>,
}

/// Implements the `PartialEq` trait for `QueryTreeNode` to allow comparison based on node index, edge from parent, and selection attributes.
/// This is also the way to "fingerprint" a node in the query tree, as it allows us to determine if two nodes are effectively
/// the same in terms of their position and attributes in the graph while merging.
impl PartialEq for QueryTreeNode {
    fn eq(&self, other: &Self) -> bool {
        self.node_index == other.node_index
            && self.edge_from_parent == other.edge_from_parent
            && self.selection_attributes == other.selection_attributes
            && self.mutation_field_position == other.mutation_field_position
            && self.condition == other.condition
            && self.context_coordinate() == other.context_coordinate()
    }
}

fn merge_query_tree_node_list(
    target_list: &mut Vec<Rc<QueryTreeNode>>,
    source_list: &[Rc<QueryTreeNode>],
) {
    if source_list.is_empty() {
        return; // nothing to merge from the source
    }

    for source_node in source_list.iter() {
        let matching_target_node = target_list
            .iter_mut()
            .find(|target_node| target_node.as_ref() == source_node.as_ref());

        match matching_target_node {
            Some(target_node) => {
                let target_node_mut = Rc::make_mut(target_node);
                target_node_mut.merge_nodes(source_node.as_ref());
            }
            None => {
                // No match found, add a clone of the source Rc to the target list.
                target_list.push(Rc::clone(source_node));
            }
        }
    }
}

impl QueryTreeNode {
    pub fn new(
        node_index: &NodeIndex,
        edge_from_parent: Option<&EdgeIndex>,
        selection_attributes: Option<&SelectionAttributes>,
        condition: Option<&Condition>,
    ) -> Self {
        QueryTreeNode {
            node_index: *node_index,
            edge_from_parent: edge_from_parent.cloned(),
            requirements: Vec::new(),
            context: None,
            children: Vec::new(),
            selection_attributes: selection_attributes.cloned(),
            mutation_field_position: None,
            condition: condition.cloned(),
        }
    }

    pub fn context_requirements(&self) -> &[Rc<QueryTreeNode>] {
        self.context
            .as_ref()
            .map_or(&[], |context| &context.requirements)
    }

    pub fn resolved_context_requirements(&self) -> &[Rc<ContextRequirement>] {
        self.context
            .as_ref()
            .map_or(&[], |context| &context.resolved_requirements)
    }

    pub fn context_coordinate(&self) -> Option<(EdgeIndex, usize)> {
        self.context.as_ref()?.coordinate
    }

    pub fn context_mut(&mut self) -> &mut NodeContext {
        self.context.get_or_insert_default()
    }

    pub fn selection_arguments(&self) -> Option<&ArgumentsMap> {
        self.selection_attributes
            .as_ref()
            .and_then(|v| v.arguments.as_ref())
    }

    pub fn selection_alias(&self) -> Option<&str> {
        self.selection_attributes
            .as_ref()
            .and_then(|v| v.alias.as_deref())
    }

    pub fn new_root(node_index: &NodeIndex) -> Self {
        QueryTreeNode::new(node_index, None, None, None)
    }

    pub fn merge_nodes(&mut self, other: &Self) {
        merge_query_tree_node_list(&mut self.children, &other.children);
        merge_query_tree_node_list(&mut self.requirements, &other.requirements);
        let Some(other_context) = &other.context else {
            return;
        };
        let context = self.context_mut();
        merge_query_tree_node_list(&mut context.requirements, &other_context.requirements);
        for requirement in &other_context.resolved_requirements {
            if !context.resolved_requirements.iter().any(|existing| {
                existing.argument == requirement.argument
                    && existing.selection == requirement.selection
                    && existing.levels == requirement.levels
            }) {
                context.resolved_requirements.push(Rc::clone(requirement));
            }
        }
    }

    pub fn from_paths(
        graph: &Graph,
        paths: &[OperationPath],
        mutation_field_position: MutationFieldPosition,
    ) -> Result<Option<Rc<Self>>, GraphError> {
        if paths.is_empty() {
            return Ok(None);
        }

        let mut trees = paths
            .iter()
            .map(|path| {
                // TODO: cover with ?
                QueryTree::from_path(graph, path, mutation_field_position)
                    .expect("expected tree to be built but it failed")
            })
            .collect::<Vec<_>>();

        if trees.len() == 1 {
            return Ok(Some(trees.remove(0).root));
        }

        Ok(Some(QueryTree::merge_trees(trees).root))
    }

    #[instrument(level = "trace",skip(graph, segments), fields(
      total_segments = segments.len()
    ))]
    fn from_path_segment_sequences(
        graph: &Graph,
        segments: &[Rc<PathSegment>],
        current_index: usize,
        mutation_field_position: MutationFieldPosition,
    ) -> Result<Option<Rc<Self>>, GraphError> {
        if current_index >= segments.len() {
            return Ok(None);
        }

        let segment_at_index: &Rc<PathSegment> = &segments[current_index];
        let edge_at_index = &segment_at_index.edge_index;
        let requirements_tree_at_index = &segment_at_index.requirement_tree;
        let selection_attributes_at_index = &segment_at_index.selection_attributes;

        trace!(
            "Processing edge: {}",
            graph.pretty_print_edge(*edge_at_index, false)
        );

        // Creates the QueryTreeNode representing the state after traversing this edge
        let tail_node_index = graph.get_edge_tail(edge_at_index)?;
        let mut tree_node = QueryTreeNode::new(
            &tail_node_index,
            Some(edge_at_index),
            (*selection_attributes_at_index).as_ref(),
            segment_at_index.condition.as_ref(),
        );

        // Only apply the position to the first node created from the segments
        if current_index == 0 {
            tree_node.mutation_field_position = mutation_field_position;
        }

        if let Some(requirements_tree_arc) = requirements_tree_at_index {
            tree_node
                .requirements
                .push(Rc::clone(requirements_tree_arc));
        }

        if !segment_at_index.context_requirements().is_empty() {
            tree_node.context_mut().resolved_requirements =
                segment_at_index.context_requirements().to_vec();
        }
        for (index, segment) in segments.iter().enumerate().skip(current_index + 1) {
            for requirement in segment.context_requirements() {
                if index.checked_sub(requirement.segments_back) == Some(current_index) {
                    merge_query_tree_node_list(
                        &mut tree_node.context_mut().requirements,
                        &[Rc::clone(&requirement.tree)],
                    );
                }
            }
        }

        let subsequent_query_tree_node =
            Self::from_path_segment_sequences(graph, segments, current_index + 1, None)?;

        match subsequent_query_tree_node {
            Some(subsequent_query_tree_node) => {
                trace!("Adding subsequent step as child");
                tree_node.children.push(subsequent_query_tree_node);
            }
            None => {
                trace!("No subsequent steps (leaf or end of path)");
            }
        }
        Ok(Some(Rc::new(tree_node)))
    }

    #[instrument(level = "trace",skip_all, fields(
      root_node = graph.pretty_print_node(root_node_index),
      segments_count = segments.len()
    ))]
    pub fn create_root_for_path_sequences(
        graph: &Graph,
        root_node_index: &NodeIndex,
        segments: &Vec<Rc<PathSegment>>,
        mutation_field_position: MutationFieldPosition,
    ) -> Result<QueryTreeNode, GraphError> {
        trace!(
            "Building root query tree node: {}",
            graph.pretty_print_node(root_node_index)
        );

        let mut root_tree_node = Self::new_root(root_node_index);

        if segments.is_empty() {
            trace!("Path has no edges beyond the root.");
        } else {
            let subsequent_node = QueryTreeNode::from_path_segment_sequences(
                graph,
                segments.as_slice(),
                0,
                mutation_field_position,
            )?;

            if let Some(subsequent_node) = subsequent_node {
                root_tree_node.children.push(subsequent_node);
            }
        }

        Ok(root_tree_node)
    }
}

#[cfg(test)]
#[cfg(target_pointer_width = "64")]
mod size_tests {
    use super::*;
    use std::mem::size_of;

    #[test]
    fn planner_structs_stay_small() {
        // pathfinding makes a lot of these, so growing them makes planning slower
        for (name, actual, expected) in [
            ("PathSegment", size_of::<PathSegment>(), 152),
            ("QueryTreeNode", size_of::<QueryTreeNode>(), 192),
        ] {
            assert_eq!(
                actual, expected,
                "{name} is {actual} B, expected {expected} B"
            );
        }
    }
}