fluidattacks-blends-domain 0.40.0

use alloc::collections::BTreeMap;
use alloc::vec::Vec;

use crate::stack::resolve::{minimal_partial_paths, PartialPath};
use crate::stack::StackGraph;
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>>,
}

#[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) {
        let found = minimal_partial_paths(graph, max_work);
        let mut paths = found.paths;
        paths.sort_by_key(|path| path.start);
        let exhausted = found.exhausted;
        self.files.insert(
            file,
            FilePaths {
                graph,
                paths,
                exhausted,
            },
        );
    }

    /// Forgets `file` and its paths.
    pub fn drop_file(&mut self, file: &str) {
        self.files.remove(file);
    }

    /// 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 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())
    }
}

#[cfg(test)]
mod tests {
    use super::{PartialPathDatabase, PartialPathId};
    use crate::stack::resolve::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);

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

    fn node(id: NodeId) -> StackNode {
        match id {
            ROOT => StackNode::Root,
            REFERENCE => push("x", true),
            RELAY => push("m", false),
            DEFINITION => StackNode::PopSymbol(SymbolPop {
                symbol: "x".to_owned(),
                precedence: 0,
                definition: Some(Definition::default()),
                import: 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()
    }

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