fluidattacks-blends-domain 0.41.0

use alloc::collections::{BTreeMap, BTreeSet};
use alloc::vec::Vec;
use core::ops::Bound;

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

/// Where a stored partial path lives: its file and its place among the file's paths.
///
/// An id lives no longer than the borrow of the database that gave it, so that database does
/// not change while the id lives. It names a path only in that database.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
#[non_exhaustive]
pub struct PartialPathId<'a> {
    /// The file whose graph the path walks.
    pub file: &'a str,
    /// The path's place among the file's paths.
    pub index: usize,
}

/// The minimal partial paths of each file, kept with the graph they walk.
#[derive(Clone, Debug, Default)]
pub struct PartialPathDatabase<'a> {
    files: BTreeMap<&'a str, FilePaths<'a>>,
    roots: BTreeMap<Vec<&'a str>, BTreeSet<PartialPathId<'a>>>,
}

#[derive(Clone, Debug)]
struct FilePaths<'a> {
    graph: &'a StackGraph,
    paths: Vec<PartialPath<'a>>,
    exhausted: Vec<NodeId>,
}

impl<'a> PartialPathDatabase<'a> {
    /// Stores the minimal partial paths of `graph` as `file`'s, in place of any it had, trying
    /// at most `max_work` edges from each endpoint.
    pub fn add_file(&mut self, file: &'a str, graph: &'a StackGraph, max_work: usize) {
        self.drop_file(file);
        let found = minimal_partial_paths(graph, max_work);
        let mut paths = found.paths;
        paths.sort_by_key(|path| path.start);
        let stored = FilePaths {
            graph,
            paths,
            exhausted: found.exhausted,
        };
        for (index, needed) in stored.root_paths() {
            self.roots
                .entry(needed)
                .or_default()
                .insert(PartialPathId { file, index });
        }
        self.files.insert(file, stored);
    }

    /// Forgets `file` and its paths.
    pub fn drop_file(&mut self, file: &str) {
        let Some((file, stored)) = self.files.remove_entry(file) else {
            return;
        };
        for (index, needed) in stored.root_paths() {
            let Some(ids) = self.roots.get_mut(&needed) else {
                continue;
            };
            ids.remove(&PartialPathId { file, index });
            if ids.is_empty() {
                self.roots.remove(&needed);
            }
        }
    }

    /// The graph `file`'s paths walk.
    #[must_use]
    pub fn graph(&self, file: &str) -> Option<&'a StackGraph> {
        self.files.get(file).map(|stored| stored.graph)
    }

    /// The path stored under `id`, an id this database gave.
    #[must_use]
    pub fn path(&self, id: PartialPathId<'_>) -> Option<&PartialPath<'a>> {
        self.files.get(id.file)?.paths.get(id.index)
    }

    /// The paths of `file` that start at `node`.
    pub fn starting_at(
        &self,
        file: &str,
        node: NodeId,
    ) -> impl Iterator<Item = (PartialPathId<'_>, &PartialPath<'a>)> + '_ {
        self.files
            .get_key_value(file)
            .into_iter()
            .flat_map(move |(&file, stored)| {
                let first = stored.paths.partition_point(|path| path.start < node);
                stored
                    .paths
                    .iter()
                    .enumerate()
                    .skip(first)
                    .take_while(move |(_, path)| path.start == node)
                    .map(move |(index, path)| (PartialPathId { file, index }, path))
            })
    }

    /// The paths of any file that start at its root and whose needed stack may follow `left`,
    /// a stack left at a root: one is a prefix of the other, and the longer one is only
    /// reachable through the variable of the shorter.
    #[must_use]
    pub fn starting_at_root(
        &self,
        left: &PartialSymbolStack<'a>,
    ) -> Vec<(PartialPathId<'a>, &PartialPath<'a>)> {
        let symbols: Vec<&'a str> = left.symbols.iter().copied().collect();
        let equal = self.roots.get(symbols.as_slice()).into_iter().flatten();
        let shorter = (0..symbols.len())
            .rev()
            .filter_map(|len| symbols.get(..len))
            .filter_map(|prefix| self.roots.get(prefix))
            .flatten();
        let longer = self
            .roots
            .range::<[&'a str], _>((Bound::Excluded(symbols.as_slice()), Bound::Unbounded))
            .take_while(|(needed, _)| left.variable.is_some() && needed.starts_with(&symbols))
            .flat_map(|(_, ids)| ids);
        let with_rest = |id: &PartialPathId<'a>| {
            self.path(*id)
                .is_some_and(|path| path.pre.variable.is_some())
        };
        equal
            .chain(shorter.filter(|id| with_rest(id)))
            .chain(longer)
            .filter_map(|id| Some((*id, self.path(*id)?)))
            .collect()
    }

    /// The endpoints of `file` whose search stopped with edges left to try.
    #[must_use]
    pub fn exhausted(&self, file: &str) -> &[NodeId] {
        self.files
            .get(file)
            .map_or(&[], |stored| stored.exhausted.as_slice())
    }
}

impl<'a> FilePaths<'a> {
    fn root_paths(&self) -> impl Iterator<Item = (usize, Vec<&'a str>)> + '_ {
        self.paths
            .iter()
            .enumerate()
            .filter(|(_, path)| matches!(self.graph.nodes.get(&path.start), Some(StackNode::Root)))
            .map(|(index, path)| (index, path.pre.symbols.iter().copied().collect()))
    }
}

#[cfg(test)]
mod tests {
    use super::{PartialPathDatabase, PartialPathId};
    use crate::stack::resolve::{PartialSymbolStack, SymbolStackVariable, DEFAULT_MAX_WORK};
    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 DEFINITION: NodeId = NodeId(4);
    const RELAY: NodeId = NodeId(5);
    const MODULE: NodeId = NodeId(6);

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

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

    fn node(id: NodeId) -> StackNode {
        match id {
            ROOT => StackNode::Root,
            REFERENCE => push("x", true),
            RELAY => push("m", false),
            DEFINITION => pop("x", Some(Definition::default())),
            MODULE => pop("m", None),
            _ => 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 ends(database: &PartialPathDatabase<'_>, file: &str, start: NodeId) -> Vec<NodeId> {
        database
            .starting_at(file, start)
            .map(|(_, path)| path.end)
            .collect()
    }

    fn reached<'a>(
        database: &PartialPathDatabase<'a>,
        symbols: &[&'a str],
        with_variable: bool,
    ) -> Vec<(&'a str, NodeId)> {
        let variable = with_variable.then_some(SymbolStackVariable::INITIAL);
        let left = PartialSymbolStack::new(symbols.iter().copied().collect(), variable);
        database
            .starting_at_root(&left)
            .into_iter()
            .map(|(id, path)| (id.file, path.end))
            .collect()
    }

    fn module_graph() -> StackGraph {
        graph(&[(ROOT, MODULE), (MODULE, DEFINITION)])
    }

    #[test]
    fn a_stack_with_the_symbols_a_root_path_needs_reaches_it() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(
            reached(&database, &["m", "x"], false),
            [("a.py", DEFINITION)]
        );
    }

    #[test]
    fn a_stack_with_more_symbols_reaches_a_root_path_that_leaves_the_rest() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(
            reached(&database, &["m", "x", "y"], false),
            [("a.py", DEFINITION)]
        );
    }

    #[test]
    fn a_stack_with_a_variable_reaches_a_root_path_that_needs_more_symbols() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(reached(&database, &["m"], true), [("a.py", DEFINITION)]);
    }

    #[test]
    fn a_stack_without_a_variable_does_not_reach_a_root_path_that_needs_more_symbols() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(reached(&database, &["m"], false), []);
    }

    #[test]
    fn a_stack_with_other_symbols_reaches_no_root_path() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(reached(&database, &["n", "x"], true), []);
    }

    #[test]
    fn a_stack_left_at_a_root_reaches_the_root_paths_of_every_file() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
        database.add_file("b.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(
            reached(&database, &["m", "x"], false),
            [("a.py", DEFINITION), ("b.py", DEFINITION)]
        );
    }

    #[test]
    fn a_path_that_starts_elsewhere_is_not_reached_from_the_root() {
        let graph = graph(&[(ROOT, MODULE), (MODULE, DEFINITION), (REFERENCE, ROOT)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(reached(&database, &[], true), [("a.py", DEFINITION)]);
    }

    #[test]
    fn a_dropped_file_is_not_reached_from_the_root() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
        database.drop_file("a.py");

        assert_eq!(reached(&database, &["m", "x"], false), []);
    }

    #[test]
    fn a_file_stored_again_is_reached_from_the_root_once() {
        let graph = module_graph();
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(
            reached(&database, &["m", "x"], false),
            [("a.py", DEFINITION)]
        );
    }

    #[test]
    fn the_paths_starting_at_a_node_leave_out_those_starting_elsewhere() {
        let graph = graph(&[(ROOT, DEFINITION), (REFERENCE, SCOPE), (SCOPE, ROOT)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(ends(&database, "a.py", REFERENCE), [ROOT]);
    }

    #[test]
    fn a_node_id_of_one_file_does_not_reach_the_paths_of_another() {
        let local = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
        let other = graph(&[(REFERENCE, SCOPE), (SCOPE, ROOT)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &local, DEFAULT_MAX_WORK);
        database.add_file("b.py", &other, DEFAULT_MAX_WORK);

        assert_eq!(ends(&database, "a.py", REFERENCE), [DEFINITION]);
    }

    #[test]
    fn a_file_never_stored_has_no_paths() {
        let database = PartialPathDatabase::default();

        assert_eq!(ends(&database, "a.py", REFERENCE), []);
    }

    #[test]
    fn the_id_of_a_path_gives_back_that_path() {
        let graph = graph(&[(ROOT, DEFINITION), (REFERENCE, SCOPE), (SCOPE, ROOT)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
        let found = database.starting_at("a.py", REFERENCE);

        assert_eq!(
            found
                .map(|(id, path)| database.path(id) == Some(path))
                .collect::<Vec<_>>(),
            [true]
        );
    }

    #[test]
    fn storing_a_file_again_replaces_its_paths() {
        let first = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
        let second = graph(&[(REFERENCE, SCOPE), (SCOPE, ROOT)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &first, DEFAULT_MAX_WORK);
        database.add_file("a.py", &second, DEFAULT_MAX_WORK);

        assert_eq!(ends(&database, "a.py", REFERENCE), [ROOT]);
    }

    #[test]
    fn the_id_of_a_dropped_file_gives_no_path() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);
        let id = PartialPathId {
            file: "a.py",
            index: 0,
        };
        database.drop_file("a.py");

        assert_eq!(database.path(id), None);
    }

    #[test]
    fn a_file_keeps_the_graph_its_paths_walk() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, DEFINITION)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, DEFAULT_MAX_WORK);

        assert_eq!(database.graph("a.py"), Some(&graph));
    }

    #[test]
    fn a_file_keeps_the_endpoints_whose_search_ran_out_of_work() {
        let graph = graph(&[(REFERENCE, SCOPE), (SCOPE, RELAY), (RELAY, SCOPE)]);
        let mut database = PartialPathDatabase::default();
        database.add_file("a.py", &graph, 100);

        assert_eq!(database.exhausted("a.py"), [REFERENCE]);
    }
}