fluidattacks-blends-domain 0.37.0

use core::fmt;

use crate::alloc::string::ToString;
use crate::ast::{AstGraph, AstNode};
use crate::NodeId;
use alloc::format;
use alloc::string::String;
use alloc::vec::Vec;

#[derive(Clone, PartialEq, Eq, Debug)]
pub enum NodeResolutionError {
    Error(String),
    NodeNotFound(String),
}
impl fmt::Display for NodeResolutionError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Error(msg) => f.write_str(msg),
            Self::NodeNotFound(node) => {
                f.write_str(&format!("The node {node} is not in the graph"))
            }
        }
    }
}

pub fn ensure_node(graph: &AstGraph, node: NodeId) -> Result<&AstNode, NodeResolutionError> {
    graph.nodes.get(&node).map_or_else(
        || Err(NodeResolutionError::NodeNotFound(node.0.to_string())),
        Ok,
    )
}

/// Resolves the single child of `node`, failing unless it has exactly one.
pub fn ensure_only_one_child<'a>(
    graph: &'a AstGraph,
    node: &AstNode,
) -> Result<&'a AstNode, NodeResolutionError> {
    let [child_id] = node.children.as_slice() else {
        return Err(NodeResolutionError::Error(format!(
            "A {} must have exactly one child, found {}",
            node.kind,
            node.children.len()
        )));
    };
    ensure_node(graph, *child_id)
}

pub fn children<'a>(
    graph: &'a AstGraph,
    node: &AstNode,
) -> Result<Vec<&'a AstNode>, NodeResolutionError> {
    node.children
        .iter()
        .map(|child_id| ensure_node(graph, *child_id))
        .collect()
}

pub fn ensure_only_two_children<'a>(
    graph: &'a AstGraph,
    node: &AstNode,
) -> Result<(&'a AstNode, &'a AstNode), NodeResolutionError> {
    let [child_1, child_2] = node.children.as_slice() else {
        return Err(NodeResolutionError::Error(format!(
            "A {} must have exactly two children, found {}",
            node.kind,
            node.children.len()
        )));
    };
    ensure_node(graph, *child_1).and_then(|c1| ensure_node(graph, *child_2).map(|c2| (c1, c2)))
}