fluidattacks-blends-domain 0.42.0

use alloc::collections::BTreeSet;
use alloc::vec::Vec;

use crate::stack::resolve::{PartialPath, PartialPathDatabase, PartialSymbolStack};
use crate::stack::StackNode;
use crate::NodeId;

/// A walk the stitcher built, and the file of the node it ends at.
#[derive(Clone, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct StitchedPath<'a> {
    /// The file of the node the walk ends at.
    pub file: &'a str,
    /// The walk.
    pub path: PartialPath<'a>,
}

type State<'a> = (NodeId, &'a str, NodeId, PartialSymbolStack<'a>);

/// Extends walks from the references of one file, phase by phase, with the stored partial paths.
#[derive(Debug)]
pub struct ForwardStitcher<'d, 'a> {
    database: &'d PartialPathDatabase<'a>,
    queue: Vec<StitchedPath<'a>>,
    seen: BTreeSet<State<'a>>,
}

impl<'d, 'a> ForwardStitcher<'d, 'a> {
    /// A stitcher whose walks start at those `references` of `file` that are references.
    pub fn from_references(
        database: &'d PartialPathDatabase<'a>,
        file: &'a str,
        references: impl IntoIterator<Item = NodeId>,
    ) -> Self {
        let mut seen = BTreeSet::new();
        let queue = database.graph(file).map_or_else(Vec::new, |graph| {
            references
                .into_iter()
                .filter(|node| graph.nodes.get(node).is_some_and(StackNode::is_reference))
                .flat_map(|node| PartialPath::seed(graph, node))
                .map(|path| StitchedPath { file, path })
                .filter(|seed| seen.insert(state(seed)))
                .collect()
        });
        Self {
            database,
            queue,
            seen,
        }
    }

    /// Extends every pending walk by one stored path, and returns the walks that now reach a
    /// definition with no symbol left to find. A walk that reaches an import also goes on to
    /// what the import names. A walk that needs symbols on the stack before the reference is
    /// dropped, since a reference starts from an empty stack.
    pub fn process_next_phase(&mut self) -> Vec<StitchedPath<'a>> {
        let pending = core::mem::take(&mut self.queue);
        let extensions: Vec<_> = pending.iter().flat_map(|walk| self.extend(walk)).collect();
        let mut complete = Vec::new();
        for walk in extensions {
            if !self.seen.insert(state(&walk)) {
                continue;
            }
            let graph = self.database.graph(walk.file);
            let end = graph.and_then(|graph| graph.nodes.get(&walk.path.end));
            let answers =
                walk.path.post.symbols.is_empty() && end.is_some_and(StackNode::is_definition);
            if answers {
                complete.push(walk.clone());
            }
            if !answers || end.is_some_and(|node| node.import_binding().is_some()) {
                self.queue.push(walk);
            }
        }
        complete
    }

    /// Whether no walk is left to extend.
    #[must_use]
    pub fn is_complete(&self) -> bool {
        self.queue.is_empty()
    }

    fn extend(&self, walk: &StitchedPath<'a>) -> Vec<StitchedPath<'a>> {
        let Some(graph) = self.database.graph(walk.file) else {
            return Vec::new();
        };
        let candidates: Vec<_> = if matches!(graph.nodes.get(&walk.path.end), Some(StackNode::Root))
        {
            let roots = self.database.starting_at_root(&walk.path.post);
            roots
                .into_iter()
                .map(|(id, path)| (id.file, path))
                .collect()
        } else {
            let here = self.database.starting_at(walk.file, walk.path.end);
            here.map(|(_, path)| (walk.file, path)).collect()
        };
        let mut extended = Vec::new();
        for (file, candidate) in candidates {
            let at_this_end = PartialPath {
                start: walk.path.end,
                ..candidate.clone()
            };
            match walk.path.concatenate(graph, &at_this_end) {
                Ok(path) if path.pre.symbols.is_empty() => {
                    extended.push(StitchedPath { file, path });
                }
                _ => {}
            }
        }
        extended
    }
}

fn state<'a>(walk: &StitchedPath<'a>) -> State<'a> {
    (
        walk.path.start,
        walk.file,
        walk.path.end,
        walk.path.post.clone(),
    )
}

#[cfg(test)]
mod tests {
    use super::ForwardStitcher;
    use crate::stack::resolve::{PartialPathDatabase, DEFAULT_MAX_WORK};
    use crate::stack::{
        Definition, ImportBinding, 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 DEFINITION: NodeId = NodeId(4);
    const RELAY: NodeId = NodeId(5);
    const MODULE: NodeId = NodeId(6);
    const MODULE_DEFINITION: NodeId = NodeId(7);
    const IMPORT: NodeId = NodeId(8);
    const ORIGIN: NodeId = NodeId(9);
    const MODULE_REFERENCE: NodeId = NodeId(10);
    const ALL_PHASES: usize = 16;
    const LIBRARY: (&str, &[(NodeId, NodeId)]) = ("m.py", &[(ROOT, MODULE), (MODULE, DEFINITION)]);

    type Edges = &'static [(NodeId, NodeId)];

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

    fn pop(symbol: &str, definition: bool, import: bool) -> StackNode {
        StackNode::PopSymbol(SymbolPop {
            symbol: symbol.to_owned(),
            precedence: 0,
            definition: definition.then(Definition::default),
            import: import.then(|| ImportBinding {
                module_stem: "m".to_owned(),
                source_symbol: None,
                relative_level: 0,
                binds_module: false,
            }),
        })
    }

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

    fn graph(edges: Edges) -> 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 stitch(
        files: &[(&'static str, Edges)],
        reference: NodeId,
        phases: usize,
    ) -> (Vec<(&'static str, NodeId)>, bool) {
        let graphs: Vec<StackGraph> = files.iter().map(|(_, edges)| graph(edges)).collect();
        let mut database = PartialPathDatabase::default();
        for ((file, _), graph) in files.iter().zip(&graphs) {
            database.add_file(file, graph, DEFAULT_MAX_WORK);
        }
        let mut stitcher = ForwardStitcher::from_references(&database, "a.py", [reference]);
        let answers = (0..phases)
            .flat_map(|_| stitcher.process_next_phase())
            .filter_map(|walk| {
                Some((
                    files.iter().find(|(file, _)| *file == walk.file)?.0,
                    walk.path.end,
                ))
            })
            .collect();
        (answers, stitcher.is_complete())
    }

    #[test]
    fn a_reference_reaches_a_definition_of_its_own_file() {
        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
        assert_eq!(stitch(&[own], REFERENCE, 1).0, [("a.py", DEFINITION)]);
    }

    #[test]
    fn a_node_that_is_not_a_reference_starts_no_walk() {
        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
        assert_eq!(stitch(&[own], SCOPE, 0), (Vec::new(), true));
    }

    #[test]
    fn a_reference_reaches_a_definition_of_another_file_through_the_root() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let answers = stitch(&[importer, LIBRARY], REFERENCE, ALL_PHASES).0;
        assert_eq!(answers, [("m.py", DEFINITION)]);
    }

    #[test]
    fn a_reference_that_reaches_an_import_also_reaches_what_it_imports() {
        let edges = &[
            (REFERENCE, IMPORT),
            (IMPORT, ORIGIN),
            (ORIGIN, RELAY),
            (RELAY, ROOT),
        ];
        let answers = stitch(&[("a.py", edges), LIBRARY], REFERENCE, ALL_PHASES).0;
        assert_eq!(answers, [("a.py", IMPORT), ("m.py", DEFINITION)]);
    }

    #[test]
    fn a_definition_reached_with_symbols_left_to_find_is_walked_past() {
        let edges = &[
            (REFERENCE, RELAY),
            (RELAY, MODULE_DEFINITION),
            (MODULE_DEFINITION, DEFINITION),
        ];
        let answers = stitch(&[("a.py", edges)], REFERENCE, ALL_PHASES).0;
        assert_eq!(answers, [("a.py", DEFINITION)]);
    }

    #[test]
    fn a_walk_that_needs_symbols_the_reference_did_not_push_is_dropped() {
        let importer = ("a.py", &[(MODULE_REFERENCE, ROOT)][..]);
        assert_eq!(
            stitch(&[importer, LIBRARY], MODULE_REFERENCE, 2),
            (Vec::new(), true)
        );
    }

    #[test]
    fn a_walk_back_to_a_state_already_seen_is_not_extended_again() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let looping = (
            "m.py",
            &[(ROOT, MODULE), (MODULE, RELAY), (RELAY, ROOT)][..],
        );
        assert!(stitch(&[importer, looping], REFERENCE, ALL_PHASES).1);
    }
}