fluidattacks-blends-domain 0.43.0

use alloc::collections::{BTreeMap, 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>,
}

/// The stored paths a stitcher tries from one reference when no other limit is given.
pub const DEFAULT_MAX_STITCH_WORK: usize = 1_000;

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>>,
    max_work: usize,
    work: BTreeMap<NodeId, usize>,
    exhausted: BTreeSet<NodeId>,
}

impl<'d, 'a> ForwardStitcher<'d, 'a> {
    /// A stitcher whose walks start at those `references` of `file` that are references, and
    /// that tries at most `max_work` stored paths from each of them.
    pub fn from_references(
        database: &'d PartialPathDatabase<'a>,
        file: &'a str,
        references: impl IntoIterator<Item = NodeId>,
        max_work: usize,
    ) -> 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,
            max_work,
            work: BTreeMap::new(),
            exhausted: BTreeSet::new(),
        }
    }

    /// 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: at a root, no stored path that needs
    /// symbols below the reference is tried. A reference that runs out of work with stored paths
    /// left to try stops there and is reported by [`Self::exhausted`].
    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()
    }

    /// The references that stopped with stored paths left to try, in id order. A reference
    /// listed here may still have answers that were never found.
    pub fn exhausted(&self) -> impl Iterator<Item = NodeId> + '_ {
        self.exhausted.iter().copied()
    }

    /// The stored paths tried so far, over every reference.
    #[must_use]
    pub fn work(&self) -> usize {
        self.work.values().sum()
    }

    fn extend(&mut 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 below_reference = walk.path.post.variable == walk.path.pre.variable;
            let left = PartialSymbolStack {
                symbols: walk.path.post.symbols.clone(),
                variable: walk.path.post.variable.filter(|_| !below_reference),
            };
            let roots = self.database.starting_at_root(&left);
            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();
        let spent = self.work.entry(walk.path.start).or_default();
        for (file, candidate) in candidates {
            if *spent >= self.max_work {
                self.exhausted.insert(walk.path.start);
                break;
            }
            *spent = spent.saturating_add(1);
            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, DEFAULT_MAX_STITCH_WORK};
    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 LOCAL_REFERENCE: NodeId = NodeId(11);
    const INNER: NodeId = NodeId(12);
    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 | LOCAL_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),
            INNER => pop("x", false, 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
    }

    #[derive(PartialEq, Debug)]
    struct Run {
        answers: Vec<(&'static str, NodeId)>,
        complete: bool,
        exhausted: Vec<NodeId>,
        work: usize,
    }

    fn stitch(
        files: &[(&'static str, Edges)],
        reference: NodeId,
        phases: usize,
    ) -> (Vec<(&'static str, NodeId)>, bool) {
        let run = stitch_within(files, &[reference], phases, DEFAULT_MAX_STITCH_WORK);
        (run.answers, run.complete)
    }

    fn stitch_within(
        files: &[(&'static str, Edges)],
        references: &[NodeId],
        phases: usize,
        max_work: usize,
    ) -> Run {
        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",
            references.iter().copied(),
            max_work,
        );
        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();
        Run {
            answers,
            complete: stitcher.is_complete(),
            exhausted: stitcher.exhausted().collect(),
            work: stitcher.work(),
        }
    }

    #[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);
    }

    #[test]
    fn a_reference_with_no_work_left_is_reported_as_exhausted() {
        let own = ("a.py", &[(REFERENCE, SCOPE), (SCOPE, DEFINITION)][..]);
        let run = stitch_within(&[own], &[REFERENCE], ALL_PHASES, 0);
        assert_eq!(
            (run.answers, run.exhausted),
            (Vec::new(), Vec::from([REFERENCE]))
        );
    }

    #[test]
    fn a_reference_with_work_for_every_stored_path_is_not_exhausted() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let run = stitch_within(&[importer, LIBRARY], &[REFERENCE], ALL_PHASES, 2);
        assert_eq!(
            (run.answers, run.exhausted),
            (Vec::from([("m.py", DEFINITION)]), Vec::new())
        );
    }

    #[test]
    fn the_work_counts_every_stored_path_tried() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let run = stitch_within(&[importer, LIBRARY], &[REFERENCE], ALL_PHASES, 2);
        assert_eq!(run.work, 2);
    }

    #[test]
    fn a_reference_that_runs_out_of_work_leaves_the_answers_of_the_others() {
        let edges = &[
            (REFERENCE, RELAY),
            (RELAY, ROOT),
            (LOCAL_REFERENCE, SCOPE),
            (SCOPE, DEFINITION),
        ];
        let references = [REFERENCE, LOCAL_REFERENCE];
        let run = stitch_within(&[("a.py", edges), LIBRARY], &references, ALL_PHASES, 1);
        assert_eq!(run.answers, [("a.py", DEFINITION)]);
    }

    #[test]
    fn a_walk_whose_stack_grows_without_end_stops_at_the_work_limit() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let growing = ("m.py", &[(ROOT, RELAY), (RELAY, ROOT)][..]);
        let run = stitch_within(&[importer, growing], &[REFERENCE], ALL_PHASES, 8);
        assert_eq!(
            (run.complete, run.exhausted),
            (true, Vec::from([REFERENCE]))
        );
    }

    #[test]
    fn a_walk_at_the_root_tries_no_path_that_needs_symbols_below_the_reference() {
        let importer = ("a.py", &[(REFERENCE, RELAY), (RELAY, ROOT)][..]);
        let deeper = (
            "n.py",
            &[(ROOT, MODULE), (MODULE, INNER), (INNER, MODULE_DEFINITION)][..],
        );
        let run = stitch_within(
            &[importer, LIBRARY, deeper],
            &[REFERENCE],
            ALL_PHASES,
            DEFAULT_MAX_STITCH_WORK,
        );
        assert_eq!(run.work, 2);
    }
}