fluidattacks-blends-domain 0.36.0

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

use crate::stack::resolve::{
    PartialBindings, PartialScopeStack, PartialSymbolStack, ScopeStackVariable,
    SymbolStackVariable, TransitionError, UnifyError,
};
use crate::stack::{StackGraph, StackNode};
use crate::NodeId;

/// An edge a partial path follows.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct PathEdge {
    /// The node the edge leaves.
    pub source: NodeId,
    /// The edge's precedence.
    pub precedence: i64,
}

/// A walk from `start` to `end`: the stacks it needs on entry and the stacks it leaves.
#[derive(Clone, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct PartialPath<'a> {
    /// The node the walk starts at.
    pub start: NodeId,
    /// The node the walk ends at.
    pub end: NodeId,
    /// The symbol stack the walk needs on entry.
    pub symbols_pre: PartialSymbolStack<'a>,
    /// The symbol stack the walk leaves.
    pub symbols_post: PartialSymbolStack<'a>,
    /// The scope stack the walk needs on entry.
    pub scopes_pre: PartialScopeStack,
    /// The scope stack the walk leaves.
    pub scopes_post: PartialScopeStack,
    /// The edges followed, in order.
    pub edges: Vec<PathEdge>,
}

/// Why two partial paths do not form one walk.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub enum PathError {
    /// The second path does not start where the first one ends.
    DisjointPaths,
    /// The node where the paths meet does not match the stacks around it.
    Transition(TransitionError),
    /// The stacks one path leaves do not match the stacks the other needs.
    Unify(UnifyError),
    /// The symbol stack left ends in a variable the symbol stack needed lacks.
    UnboundSymbolVariable,
    /// The scope stacks left use a variable the stacks needed lack.
    UnboundScopeVariable,
    /// Renaming the variables apart ran past the last variable number.
    VariableOverflow,
}

impl From<TransitionError> for PathError {
    fn from(error: TransitionError) -> Self {
        Self::Transition(error)
    }
}

impl From<UnifyError> for PathError {
    fn from(error: UnifyError) -> Self {
        Self::Unify(error)
    }
}

impl PartialPath<'_> {
    /// This walk followed by `rhs`, which must start where this one ends.
    ///
    /// # Errors
    ///
    /// Returns the [`PathError`] that keeps the walks apart.
    pub fn concatenate(&self, graph: &StackGraph, rhs: &Self) -> Result<Self, PathError> {
        if self.end != rhs.start {
            return Err(PathError::DisjointPaths);
        }
        let rhs = rhs
            .with_offset(
                self.largest_symbol_variable(),
                self.largest_scope_variable(),
            )
            .ok_or(PathError::VariableOverflow)?;
        let join = graph
            .nodes
            .get(&self.end)
            .ok_or(TransitionError::UnknownNode)?;
        let left = unpush(join, &self.symbols_post)?;
        let (right, right_scopes) = unpop(join, &rhs.symbols_pre, &rhs.scopes_pre)?;
        let mut bindings = PartialBindings::new();
        bindings.unify_symbols(&left, &right)?;
        bindings.unify_scopes(&self.scopes_post, &right_scopes)?;
        let path = Self {
            start: self.start,
            end: rhs.end,
            symbols_pre: self.symbols_pre.apply_bindings(&bindings),
            symbols_post: rhs.symbols_post.apply_bindings(&bindings),
            scopes_pre: self.scopes_pre.apply_bindings(&bindings),
            scopes_post: rhs.scopes_post.apply_bindings(&bindings),
            edges: self.edges.iter().chain(&rhs.edges).copied().collect(),
        };
        path.check_bound_variables()?;
        Ok(path)
    }

    fn largest_symbol_variable(&self) -> u32 {
        self.symbols_pre
            .variable
            .max(self.symbols_post.variable)
            .map_or(0, SymbolStackVariable::get)
    }

    fn largest_scope_variable(&self) -> u32 {
        [
            self.scopes_pre.variable,
            self.scopes_post.variable,
            self.symbols_pre.largest_scope_variable(),
            self.symbols_post.largest_scope_variable(),
        ]
        .into_iter()
        .flatten()
        .max()
        .map_or(0, ScopeStackVariable::get)
    }

    fn with_offset(&self, symbol_offset: u32, scope_offset: u32) -> Option<Self> {
        Some(Self {
            start: self.start,
            end: self.end,
            symbols_pre: self.symbols_pre.with_offset(symbol_offset, scope_offset)?,
            symbols_post: self.symbols_post.with_offset(symbol_offset, scope_offset)?,
            scopes_pre: self.scopes_pre.with_offset(scope_offset)?,
            scopes_post: self.scopes_post.with_offset(scope_offset)?,
            edges: self.edges.clone(),
        })
    }

    fn check_bound_variables(&self) -> Result<(), PathError> {
        if self.symbols_post.variable.is_some()
            && self.symbols_post.variable != self.symbols_pre.variable
        {
            return Err(PathError::UnboundSymbolVariable);
        }
        let needed = scope_variables(&self.symbols_pre, &self.scopes_pre);
        if scope_variables(&self.symbols_post, &self.scopes_post).is_subset(&needed) {
            Ok(())
        } else {
            Err(PathError::UnboundScopeVariable)
        }
    }
}

fn unpush<'a>(
    join: &StackNode,
    symbols: &PartialSymbolStack<'a>,
) -> Result<PartialSymbolStack<'a>, TransitionError> {
    match join {
        StackNode::PushSymbol(push) => pop_top(symbols, &push.symbol, false).map(|(rest, _)| rest),
        StackNode::PushScopedSymbol { push, .. } => {
            pop_top(symbols, &push.symbol, true).map(|(rest, _)| rest)
        }
        _ => Ok(symbols.clone()),
    }
}

fn unpop<'a>(
    join: &StackNode,
    symbols: &PartialSymbolStack<'a>,
    scopes: &PartialScopeStack,
) -> Result<(PartialSymbolStack<'a>, PartialScopeStack), TransitionError> {
    match join {
        StackNode::PopSymbol(pop) => {
            pop_top(symbols, &pop.symbol, false).map(|(rest, _)| (rest, scopes.clone()))
        }
        StackNode::PopScopedSymbol(pop) => pop_top(symbols, &pop.symbol, true)
            .map(|(rest, attached)| (rest, attached.unwrap_or_default())),
        _ => Ok((symbols.clone(), scopes.clone())),
    }
}

fn pop_top<'a>(
    symbols: &PartialSymbolStack<'a>,
    symbol: &str,
    scoped: bool,
) -> Result<(PartialSymbolStack<'a>, Option<PartialScopeStack>), TransitionError> {
    let mut rest = symbols.clone();
    let top = rest
        .symbols
        .pop_front()
        .ok_or(TransitionError::EmptySymbolStack)?;
    if top.symbol != symbol {
        return Err(TransitionError::IncorrectPoppedSymbol);
    }
    match (scoped, top.scopes) {
        (true, None) => Err(TransitionError::MissingAttachedScopes),
        (false, Some(_)) => Err(TransitionError::UnexpectedAttachedScopes),
        (_, attached) => Ok((rest, attached)),
    }
}

fn scope_variables(
    symbols: &PartialSymbolStack<'_>,
    scopes: &PartialScopeStack,
) -> BTreeSet<ScopeStackVariable> {
    symbols
        .symbols
        .iter()
        .filter_map(|symbol| symbol.scopes.as_ref()?.variable)
        .chain(scopes.variable)
        .collect()
}

#[cfg(test)]
mod tests {
    use super::{PartialPath, PathEdge, PathError};
    use crate::stack::resolve::{
        PartialScopeStack, PartialScopedSymbol, PartialSymbolStack, ScopeStackVariable,
        SymbolStackVariable, TransitionError, UnifyError,
    };
    use crate::stack::{Definition, StackGraph, StackNode, StackScope, SymbolPop, SymbolPush};
    use crate::NodeId;
    use alloc::borrow::ToOwned;
    use alloc::collections::VecDeque;
    use alloc::vec::Vec;

    const PUSH_X: NodeId = NodeId(10);
    const PUSH_Y: NodeId = NodeId(11);
    const SCOPE: NodeId = NodeId(20);
    const OTHER_SCOPE: NodeId = NodeId(21);
    const POP_X: NodeId = NodeId(30);
    const SCOPED_POP_X: NodeId = NodeId(31);

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

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

    fn graph() -> StackGraph {
        let mut graph = StackGraph::new();
        graph.add_node(PUSH_X, push("x"));
        graph.add_node(PUSH_Y, push("y"));
        graph.add_node(SCOPE, StackNode::Scope(StackScope::default()));
        graph.add_node(OTHER_SCOPE, StackNode::Scope(StackScope::default()));
        graph.add_node(POP_X, StackNode::PopSymbol(pop("x")));
        graph.add_node(SCOPED_POP_X, StackNode::PopScopedSymbol(pop("x")));
        graph
    }

    fn symbols<'a>(names: &[&'a str], variable: Option<u32>) -> PartialSymbolStack<'a> {
        PartialSymbolStack::new(
            names
                .iter()
                .map(|name| PartialScopedSymbol::new(name, None))
                .collect(),
            variable.and_then(SymbolStackVariable::new),
        )
    }

    fn scopes(ids: &[u64], variable: Option<u32>) -> PartialScopeStack {
        PartialScopeStack::new(
            ids.iter().copied().map(NodeId).collect(),
            variable.and_then(ScopeStackVariable::new),
        )
    }

    fn path<'a>(
        start: NodeId,
        end: NodeId,
        pre: PartialSymbolStack<'a>,
        post: PartialSymbolStack<'a>,
    ) -> PartialPath<'a> {
        PartialPath {
            start,
            end,
            symbols_pre: pre,
            symbols_post: post,
            scopes_pre: scopes(&[], Some(1)),
            scopes_post: scopes(&[], Some(1)),
            edges: Vec::new(),
        }
    }

    fn through(start: NodeId, end: NodeId) -> PartialPath<'static> {
        path(start, end, symbols(&[], Some(1)), symbols(&[], Some(1)))
    }

    #[test]
    fn a_reference_path_joined_with_a_definition_path_leaves_the_stacks_as_they_were() {
        let graph = graph();
        let reference = path(
            PUSH_X,
            SCOPE,
            symbols(&[], Some(1)),
            symbols(&["x"], Some(1)),
        );
        let definition = path(
            SCOPE,
            POP_X,
            symbols(&["x"], Some(1)),
            symbols(&[], Some(1)),
        );

        assert_eq!(
            reference.concatenate(&graph, &definition),
            Ok(through(PUSH_X, POP_X))
        );
    }

    #[test]
    fn paths_that_do_not_meet_do_not_concatenate() {
        assert_eq!(
            through(PUSH_X, SCOPE).concatenate(&graph(), &through(OTHER_SCOPE, POP_X)),
            Err(PathError::DisjointPaths)
        );
    }

    #[test]
    fn a_push_where_the_paths_meet_is_counted_once() {
        let graph = graph();
        let lhs = path(
            SCOPE,
            PUSH_X,
            symbols(&[], Some(1)),
            symbols(&["x"], Some(1)),
        );
        let rhs = path(
            PUSH_X,
            PUSH_Y,
            symbols(&[], Some(1)),
            symbols(&["y", "x"], Some(1)),
        );

        assert_eq!(
            lhs.concatenate(&graph, &rhs)
                .map(|joined| joined.symbols_post),
            Ok(symbols(&["y", "x"], Some(1)))
        );
    }

    #[test]
    fn a_pop_where_the_paths_meet_is_required_once() {
        let graph = graph();
        let lhs = path(
            SCOPE,
            POP_X,
            symbols(&["x"], Some(1)),
            symbols(&[], Some(1)),
        );
        let rhs = path(
            POP_X,
            OTHER_SCOPE,
            symbols(&["x"], Some(1)),
            symbols(&[], Some(1)),
        );

        assert_eq!(
            lhs.concatenate(&graph, &rhs)
                .map(|joined| joined.symbols_pre),
            Ok(symbols(&["x"], Some(1)))
        );
    }

    #[test]
    fn a_scoped_pop_where_the_paths_meet_continues_with_the_scopes_its_symbol_carried() {
        let graph = graph();
        let carried = |attached| {
            PartialSymbolStack::new(
                VecDeque::from([PartialScopedSymbol::new("x", Some(attached))]),
                SymbolStackVariable::new(1),
            )
        };
        let lhs = PartialPath {
            scopes_post: scopes(&[5], None),
            ..path(
                SCOPE,
                SCOPED_POP_X,
                carried(scopes(&[5], None)),
                symbols(&[], Some(1)),
            )
        };
        let rhs = PartialPath {
            scopes_post: scopes(&[], Some(2)),
            ..path(
                SCOPED_POP_X,
                OTHER_SCOPE,
                carried(scopes(&[], Some(2))),
                symbols(&[], Some(1)),
            )
        };

        assert_eq!(
            lhs.concatenate(&graph, &rhs)
                .map(|joined| joined.scopes_post),
            Ok(scopes(&[5], None))
        );
    }

    #[test]
    fn a_push_of_another_symbol_where_the_paths_meet_does_not_concatenate() {
        let lhs = path(
            SCOPE,
            PUSH_X,
            symbols(&[], Some(1)),
            symbols(&["y"], Some(1)),
        );

        assert_eq!(
            lhs.concatenate(&graph(), &through(PUSH_X, PUSH_Y)),
            Err(PathError::Transition(
                TransitionError::IncorrectPoppedSymbol
            ))
        );
    }

    #[test]
    fn a_symbol_left_that_the_next_path_does_not_need_does_not_concatenate() {
        let lhs = path(
            PUSH_X,
            SCOPE,
            symbols(&[], Some(1)),
            symbols(&["x"], Some(1)),
        );
        let rhs = path(
            SCOPE,
            OTHER_SCOPE,
            symbols(&["y"], Some(1)),
            symbols(&[], Some(1)),
        );

        assert_eq!(
            lhs.concatenate(&graph(), &rhs),
            Err(PathError::Unify(UnifyError::SymbolStackUnsatisfied))
        );
    }

    #[test]
    fn the_variables_of_the_second_path_are_renamed_apart_from_the_first() {
        let graph = graph();
        let lhs = path(
            PUSH_X,
            SCOPE,
            symbols(&[], Some(1)),
            symbols(&["x"], Some(1)),
        );
        let rhs = path(
            SCOPE,
            OTHER_SCOPE,
            symbols(&[], Some(1)),
            symbols(&["y"], Some(1)),
        );

        assert_eq!(
            lhs.concatenate(&graph, &rhs)
                .map(|joined| joined.symbols_post),
            Ok(symbols(&["y", "x"], Some(1)))
        );
    }

    #[test]
    fn renaming_past_the_last_variable_number_does_not_concatenate() {
        let rhs = path(
            SCOPE,
            OTHER_SCOPE,
            symbols(&[], Some(u32::MAX)),
            symbols(&[], Some(u32::MAX)),
        );

        assert_eq!(
            through(PUSH_X, SCOPE).concatenate(&graph(), &rhs),
            Err(PathError::VariableOverflow)
        );
    }

    #[test]
    fn a_symbol_variable_left_that_was_not_needed_does_not_concatenate() {
        let lhs = path(PUSH_X, SCOPE, symbols(&[], None), symbols(&[], Some(1)));

        assert_eq!(
            lhs.concatenate(&graph(), &through(SCOPE, OTHER_SCOPE)),
            Err(PathError::UnboundSymbolVariable)
        );
    }

    #[test]
    fn a_scope_variable_left_that_was_not_needed_does_not_concatenate() {
        let lhs = PartialPath {
            scopes_pre: scopes(&[], None),
            ..through(PUSH_X, SCOPE)
        };

        assert_eq!(
            lhs.concatenate(&graph(), &through(SCOPE, OTHER_SCOPE)),
            Err(PathError::UnboundScopeVariable)
        );
    }

    #[test]
    fn the_joined_path_follows_the_edges_of_both_in_order() {
        let edge = |source| PathEdge {
            source,
            precedence: 0,
        };
        let lhs = PartialPath {
            edges: Vec::from([edge(PUSH_X)]),
            ..through(PUSH_X, SCOPE)
        };
        let rhs = PartialPath {
            edges: Vec::from([edge(SCOPE)]),
            ..through(SCOPE, OTHER_SCOPE)
        };

        assert_eq!(
            lhs.concatenate(&graph(), &rhs).map(|joined| joined.edges),
            Ok(Vec::from([edge(PUSH_X), edge(SCOPE)]))
        );
    }
}