fluidattacks-blends-domain 0.40.0

use alloc::collections::{BTreeSet, VecDeque};
use alloc::vec::Vec;

use crate::stack::resolve::{PartialPath, PathEdge, TransitionError};
use crate::stack::{StackGraph, StackNode};
use crate::NodeId;

/// The edges a search from one endpoint tries before it stops by default.
pub const DEFAULT_MAX_WORK: usize = 2_000;

/// The walks a search found, and the endpoints whose search ran out of work.
#[derive(Clone, PartialEq, Eq, Debug, Default)]
#[non_exhaustive]
pub struct MinimalPartialPaths<'a> {
    /// The walks found, endpoint by endpoint.
    pub paths: Vec<PartialPath<'a>>,
    /// The endpoints whose search stopped with edges left to try.
    pub exhausted: Vec<NodeId>,
}

/// Every walk of `graph` from one endpoint to the next, without passing another endpoint,
/// trying at most `max_work` edges from each endpoint.
///
/// Endpoints are references, definitions and the root.
#[must_use]
pub fn minimal_partial_paths(graph: &StackGraph, max_work: usize) -> MinimalPartialPaths<'_> {
    let mut found = MinimalPartialPaths::default();
    let seeds = graph
        .nodes
        .iter()
        .filter(|(_, node)| is_endpoint(node))
        .flat_map(|(id, _)| PartialPath::seed(graph, *id));
    for seed in seeds {
        let start = seed.start;
        if paths_from(graph, seed, max_work, &mut found.paths) {
            found.exhausted.push(start);
        }
    }
    found
}

fn paths_from<'a>(
    graph: &'a StackGraph,
    seed: PartialPath<'a>,
    max_work: usize,
    paths: &mut Vec<PartialPath<'a>>,
) -> bool {
    let mut visited = BTreeSet::new();
    let mut queue = VecDeque::from([seed]);
    let mut left = max_work;
    while let Some(path) = queue.pop_front() {
        for (sink, precedence) in graph.successors(path.end) {
            let Some(rest) = left.checked_sub(1) else {
                return true;
            };
            left = rest;
            let Ok(next) = append(graph, &path, sink, precedence) else {
                continue;
            };
            if !visited.insert((next.end, next.pre.clone(), next.post.clone())) {
                continue;
            }
            if graph.nodes.get(&sink).is_some_and(is_endpoint) {
                paths.push(next);
            } else {
                queue.push_back(next);
            }
        }
    }
    false
}

fn append<'a>(
    graph: &'a StackGraph,
    path: &PartialPath<'a>,
    sink: NodeId,
    precedence: i64,
) -> Result<PartialPath<'a>, TransitionError> {
    let mut next = path.clone();
    match graph.nodes.get(&sink).ok_or(TransitionError::UnknownNode)? {
        StackNode::PushSymbol(push) => next.post.symbols.push_front(&push.symbol),
        StackNode::PopSymbol(pop) => match next.post.symbols.pop_front() {
            Some(top) if top == pop.symbol => {}
            Some(_) => return Err(TransitionError::IncorrectPoppedSymbol),
            None if next.post.variable.is_some() => next.pre.symbols.push_back(&pop.symbol),
            None => return Err(TransitionError::EmptySymbolStack),
        },
        StackNode::Root | StackNode::Scope(_) => {}
    }
    next.end = sink;
    next.edges.push(PathEdge {
        source: path.end,
        precedence,
    });
    Ok(next)
}

fn is_endpoint(node: &StackNode) -> bool {
    node.is_reference() || node.is_definition() || matches!(node, StackNode::Root)
}

#[cfg(test)]
mod tests {
    use super::{minimal_partial_paths, DEFAULT_MAX_WORK};
    use crate::stack::resolve::{PartialSymbolStack, PathEdge, SymbolStackVariable};
    use crate::stack::{Definition, StackGraph, StackNode, StackScope, SymbolPop, SymbolPush};
    use crate::NodeId;
    use alloc::borrow::ToOwned;
    use alloc::vec::Vec;

    const ROOT: NodeId = NodeId(1);
    const REFERENCE: NodeId = NodeId(2);
    const SCOPE: NodeId = NodeId(3);
    const OTHER_SCOPE: NodeId = NodeId(4);
    const DEFINITION: NodeId = NodeId(5);
    const OTHER_DEFINITION: NodeId = NodeId(6);
    const RELAY: NodeId = NodeId(7);

    fn push(symbol: &str, is_reference: bool) -> StackNode {
        StackNode::PushSymbol(SymbolPush {
            symbol: symbol.to_owned(),
            precedence: 0,
            is_reference,
        })
    }

    fn pop(symbol: &str) -> StackNode {
        StackNode::PopSymbol(SymbolPop {
            symbol: symbol.to_owned(),
            precedence: 0,
            definition: Some(Definition::default()),
            import: None,
        })
    }

    fn node(id: NodeId) -> StackNode {
        match id {
            ROOT => StackNode::Root,
            REFERENCE => push("x", true),
            RELAY => push("m", false),
            DEFINITION => pop("x"),
            OTHER_DEFINITION => pop("y"),
            _ => StackNode::Scope(StackScope::default()),
        }
    }

    fn graph(edges: &[(NodeId, NodeId)]) -> StackGraph {
        let mut graph = StackGraph::new();
        for (source, sink) in edges {
            graph.add_node(*source, node(*source));
            graph.add_node(*sink, node(*sink));
            graph.add_edge(*source, *sink, 0);
        }
        graph
    }

    fn stack<'a>(symbols: &[&'a str]) -> PartialSymbolStack<'a> {
        PartialSymbolStack::new(
            symbols.iter().copied().collect(),
            Some(SymbolStackVariable::INITIAL),
        )
    }

    type Walk<'a> = (
        NodeId,
        NodeId,
        PartialSymbolStack<'a>,
        PartialSymbolStack<'a>,
    );

    fn walks(graph: &StackGraph) -> Vec<Walk<'_>> {
        minimal_partial_paths(graph, DEFAULT_MAX_WORK)
            .paths
            .into_iter()
            .map(|path| (path.start, path.end, path.pre, path.post))
            .collect()
    }

    #[test]
    fn a_reference_reaches_the_definition_of_its_symbol_with_the_stack_as_it_found_it() {
        assert_eq!(
            walks(&graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)])),
            [(REFERENCE, DEFINITION, stack(&[]), stack(&[]))]
        );
    }

    #[test]
    fn a_reference_does_not_reach_the_definition_of_another_symbol() {
        assert_eq!(
            walks(&graph(&[(REFERENCE, SCOPE), (SCOPE, OTHER_DEFINITION)])),
            []
        );
    }

    #[test]
    fn a_definition_reached_from_the_root_needs_its_symbol_on_the_stack() {
        assert_eq!(
            walks(&graph(&[(ROOT, DEFINITION)])),
            [(ROOT, DEFINITION, stack(&["x"]), stack(&[]))]
        );
    }

    #[test]
    fn a_relay_push_leaves_its_symbol_without_ending_the_walk() {
        assert_eq!(
            walks(&graph(&[(REFERENCE, RELAY), (RELAY, ROOT)])),
            [(REFERENCE, ROOT, stack(&[]), stack(&["m", "x"]))]
        );
    }

    #[test]
    fn a_walk_ends_at_the_first_endpoint_it_reaches() {
        assert_eq!(
            walks(&graph(&[
                (SCOPE, ROOT),
                (ROOT, REFERENCE),
                (REFERENCE, SCOPE)
            ])),
            [
                (ROOT, REFERENCE, stack(&[]), stack(&["x"])),
                (REFERENCE, ROOT, stack(&[]), stack(&["x"])),
            ]
        );
    }

    #[test]
    fn a_cycle_without_endpoints_ends_the_search_without_walks() {
        assert_eq!(
            walks(&graph(&[
                (REFERENCE, SCOPE),
                (SCOPE, OTHER_SCOPE),
                (OTHER_SCOPE, SCOPE)
            ])),
            []
        );
    }

    #[test]
    fn two_routes_with_the_same_stacks_give_one_walk() {
        assert_eq!(
            walks(&graph(&[
                (REFERENCE, SCOPE),
                (REFERENCE, OTHER_SCOPE),
                (SCOPE, DEFINITION),
                (OTHER_SCOPE, DEFINITION),
            ])),
            [(REFERENCE, DEFINITION, stack(&[]), stack(&[]))]
        );
    }

    #[test]
    fn a_walk_follows_its_edges_from_the_start() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
        let edge = |source| PathEdge {
            source,
            precedence: 0,
        };

        assert_eq!(
            minimal_partial_paths(&graph, DEFAULT_MAX_WORK)
                .paths
                .into_iter()
                .map(|path| path.edges)
                .collect::<Vec<_>>(),
            [Vec::from([edge(REFERENCE), edge(SCOPE)])]
        );
    }

    #[test]
    fn a_cycle_that_keeps_growing_the_stack_stops_at_the_budget_and_reports_its_endpoint() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, RELAY), (RELAY, SCOPE)]);

        assert_eq!(minimal_partial_paths(&graph, 100).exhausted, [REFERENCE]);
    }

    #[test]
    fn a_search_with_exactly_the_work_it_needs_is_not_exhausted() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);

        assert_eq!(minimal_partial_paths(&graph, 2).exhausted, []);
    }

    #[test]
    fn a_search_stopped_before_its_last_edge_is_exhausted() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);

        assert_eq!(minimal_partial_paths(&graph, 1).exhausted, [REFERENCE]);
    }
}