fluidattacks-blends-domain 0.36.0

use alloc::collections::{BTreeMap, VecDeque};
use core::cmp::Ordering;

use crate::stack::resolve::{ScopeStackVariable, SymbolStackVariable};
use crate::NodeId;

/// A scope stack pattern: known scopes, top first, over an optional unknown rest.
#[derive(Clone, PartialEq, Eq, Debug, Default)]
#[non_exhaustive]
pub struct PartialScopeStack {
    /// The known scopes, top first.
    pub scopes: VecDeque<NodeId>,
    /// The unknown rest below them; `None` when the stack ends there.
    pub variable: Option<ScopeStackVariable>,
}

/// A symbol in a partial symbol stack.
#[derive(Clone, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct PartialScopedSymbol<'a> {
    /// The symbol's name.
    pub symbol: &'a str,
    /// The attached scope stack pattern; `None` for a plain symbol.
    pub scopes: Option<PartialScopeStack>,
}

/// A symbol stack pattern: known symbols, top first, over an optional unknown rest.
#[derive(Clone, PartialEq, Eq, Debug, Default)]
#[non_exhaustive]
pub struct PartialSymbolStack<'a> {
    /// The known symbols, top first.
    pub symbols: VecDeque<PartialScopedSymbol<'a>>,
    /// The unknown rest below them; `None` when the stack ends there.
    pub variable: Option<SymbolStackVariable>,
}

/// What the variables of unified partial stacks stand for.
#[derive(Clone, PartialEq, Eq, Debug, Default)]
#[non_exhaustive]
pub struct PartialBindings<'a> {
    /// The symbol stack each symbol stack variable stands for.
    pub symbols: BTreeMap<SymbolStackVariable, PartialSymbolStack<'a>>,
    /// The scope stack each scope stack variable stands for.
    pub scopes: BTreeMap<ScopeStackVariable, PartialScopeStack>,
}

/// Why two partial stacks have no stack in common.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub enum UnifyError {
    /// The symbol stacks disagree.
    SymbolStackUnsatisfied,
    /// The scope stacks disagree.
    ScopeStackUnsatisfied,
}

impl PartialScopeStack {
    /// `scopes`, top first, over `variable`.
    #[must_use]
    pub fn new(scopes: VecDeque<NodeId>, variable: Option<ScopeStackVariable>) -> Self {
        Self { scopes, variable }
    }

    /// Only the unknown rest `variable`.
    #[must_use]
    pub fn from_variable(variable: ScopeStackVariable) -> Self {
        Self::new(VecDeque::new(), Some(variable))
    }

    /// This stack with its variable renumbered `offset` places up; `None` past the last number.
    #[must_use]
    pub fn with_offset(&self, offset: u32) -> Option<Self> {
        let variable = self.variable.map_or(Some(None), |variable| {
            variable.with_offset(offset).map(Some)
        })?;
        Some(Self::new(self.scopes.clone(), variable))
    }

    /// This stack with its variable replaced by what `bindings` binds it to.
    #[must_use]
    pub fn apply_bindings(&self, bindings: &PartialBindings<'_>) -> Self {
        self.variable
            .and_then(|variable| bindings.scopes.get(&variable))
            .map_or_else(
                || self.clone(),
                |bound| {
                    let bound = bound.apply_bindings(bindings);
                    Self::new(
                        self.scopes.iter().chain(&bound.scopes).copied().collect(),
                        bound.variable,
                    )
                },
            )
    }
}

impl<'a> PartialScopedSymbol<'a> {
    /// `symbol` with the attached `scopes`, if any.
    #[must_use]
    pub fn new(symbol: &'a str, scopes: Option<PartialScopeStack>) -> Self {
        Self { symbol, scopes }
    }

    /// This symbol with its attached scope variable renumbered `offset` places up.
    #[must_use]
    pub fn with_offset(&self, offset: u32) -> Option<Self> {
        let scopes = self
            .scopes
            .as_ref()
            .map_or(Some(None), |scopes| scopes.with_offset(offset).map(Some))?;
        Some(Self::new(self.symbol, scopes))
    }

    /// This symbol with its attached scopes bound by `bindings`.
    #[must_use]
    pub fn apply_bindings(&self, bindings: &PartialBindings<'_>) -> Self {
        Self::new(
            self.symbol,
            self.scopes
                .as_ref()
                .map(|scopes| scopes.apply_bindings(bindings)),
        )
    }
}

impl<'a> PartialSymbolStack<'a> {
    /// `symbols`, top first, over `variable`.
    #[must_use]
    pub fn new(
        symbols: VecDeque<PartialScopedSymbol<'a>>,
        variable: Option<SymbolStackVariable>,
    ) -> Self {
        Self { symbols, variable }
    }

    /// Only the unknown rest `variable`.
    #[must_use]
    pub fn from_variable(variable: SymbolStackVariable) -> Self {
        Self::new(VecDeque::new(), Some(variable))
    }

    /// This stack with its symbol and attached scope variables renumbered up.
    #[must_use]
    pub fn with_offset(&self, symbol_offset: u32, scope_offset: u32) -> Option<Self> {
        let symbols = self
            .symbols
            .iter()
            .map(|symbol| symbol.with_offset(scope_offset))
            .collect::<Option<_>>()?;
        let variable = self.variable.map_or(Some(None), |variable| {
            variable.with_offset(symbol_offset).map(Some)
        })?;
        Some(Self::new(symbols, variable))
    }

    /// The highest scope stack variable among the attached scopes.
    #[must_use]
    pub fn largest_scope_variable(&self) -> Option<ScopeStackVariable> {
        self.symbols
            .iter()
            .filter_map(|symbol| symbol.scopes.as_ref()?.variable)
            .max()
    }

    /// This stack with every variable replaced by what `bindings` binds it to.
    #[must_use]
    pub fn apply_bindings(&self, bindings: &PartialBindings<'a>) -> Self {
        let mut symbols: VecDeque<_> = self
            .symbols
            .iter()
            .map(|symbol| symbol.apply_bindings(bindings))
            .collect();
        match self
            .variable
            .and_then(|variable| bindings.symbols.get(&variable))
        {
            Some(bound) => {
                let bound = bound.apply_bindings(bindings);
                symbols.extend(bound.symbols);
                Self::new(symbols, bound.variable)
            }
            None => Self::new(symbols, self.variable),
        }
    }
}

impl<'a> PartialBindings<'a> {
    /// No variable bound.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// The symbol stack both `lhs` and `rhs` match, binding their variables.
    ///
    /// # Errors
    ///
    /// Returns the [`UnifyError`] that keeps them apart, leaving the bindings unchanged.
    pub fn unify_symbols(
        &mut self,
        lhs: &PartialSymbolStack<'a>,
        rhs: &PartialSymbolStack<'a>,
    ) -> Result<PartialSymbolStack<'a>, UnifyError> {
        let mut trial = self.clone();
        let unified = trial.join_symbols(lhs, rhs)?.apply_bindings(&trial);
        *self = trial;
        Ok(unified)
    }

    /// The scope stack both `lhs` and `rhs` match, binding their variables.
    ///
    /// # Errors
    ///
    /// Returns the [`UnifyError`] that keeps them apart, leaving the bindings unchanged.
    pub fn unify_scopes(
        &mut self,
        lhs: &PartialScopeStack,
        rhs: &PartialScopeStack,
    ) -> Result<PartialScopeStack, UnifyError> {
        let mut trial = self.clone();
        let unified = trial.join_scopes(lhs, rhs)?;
        *self = trial;
        Ok(unified)
    }

    fn join_symbols(
        &mut self,
        lhs: &PartialSymbolStack<'a>,
        rhs: &PartialSymbolStack<'a>,
    ) -> Result<PartialSymbolStack<'a>, UnifyError> {
        let lhs = lhs.apply_bindings(self);
        let rhs = rhs.apply_bindings(self);
        let mut symbols = lhs
            .symbols
            .iter()
            .zip(&rhs.symbols)
            .map(|(left, right)| self.join_symbol(left, right))
            .collect::<Result<VecDeque<_>, _>>()?;
        let common = symbols.len();
        let longer = self.join_symbol_rests(&lhs, &rhs, common)?;
        symbols.extend(longer.symbols.iter().skip(common).cloned());
        Ok(PartialSymbolStack::new(symbols, longer.variable))
    }

    fn join_symbol(
        &mut self,
        lhs: &PartialScopedSymbol<'a>,
        rhs: &PartialScopedSymbol<'a>,
    ) -> Result<PartialScopedSymbol<'a>, UnifyError> {
        if lhs.symbol != rhs.symbol {
            return Err(UnifyError::SymbolStackUnsatisfied);
        }
        let scopes = match (&lhs.scopes, &rhs.scopes) {
            (Some(left), Some(right)) => Some(self.join_scopes(left, right)?),
            (None, None) => None,
            (Some(_), None) | (None, Some(_)) => return Err(UnifyError::SymbolStackUnsatisfied),
        };
        Ok(PartialScopedSymbol::new(lhs.symbol, scopes))
    }

    fn join_symbol_rests<'s>(
        &mut self,
        lhs: &'s PartialSymbolStack<'a>,
        rhs: &'s PartialSymbolStack<'a>,
        common: usize,
    ) -> Result<&'s PartialSymbolStack<'a>, UnifyError> {
        match (
            lhs.symbols.len().cmp(&rhs.symbols.len()),
            lhs.variable,
            rhs.variable,
        ) {
            (Ordering::Greater, _, _) => self
                .bind_symbol_rest(rhs.variable, lhs, common)
                .map(|()| lhs),
            (Ordering::Less, _, _) => self
                .bind_symbol_rest(lhs.variable, rhs, common)
                .map(|()| rhs),
            (Ordering::Equal, Some(left), Some(right)) if left != right => {
                self.symbols
                    .insert(right, PartialSymbolStack::from_variable(left));
                Ok(lhs)
            }
            (Ordering::Equal, None, Some(right)) => {
                self.symbols.insert(right, PartialSymbolStack::default());
                Ok(lhs)
            }
            (Ordering::Equal, Some(left), None) => {
                self.symbols.insert(left, PartialSymbolStack::default());
                Ok(rhs)
            }
            (Ordering::Equal, _, _) => Ok(lhs),
        }
    }

    fn bind_symbol_rest(
        &mut self,
        variable: Option<SymbolStackVariable>,
        longer: &PartialSymbolStack<'a>,
        common: usize,
    ) -> Result<(), UnifyError> {
        let variable = variable
            .filter(|variable| longer.variable != Some(*variable))
            .ok_or(UnifyError::SymbolStackUnsatisfied)?;
        let rest = PartialSymbolStack::new(
            longer.symbols.iter().skip(common).cloned().collect(),
            longer.variable,
        );
        self.symbols.insert(variable, rest);
        Ok(())
    }

    fn join_scopes(
        &mut self,
        lhs: &PartialScopeStack,
        rhs: &PartialScopeStack,
    ) -> Result<PartialScopeStack, UnifyError> {
        let lhs = lhs.apply_bindings(self);
        let rhs = rhs.apply_bindings(self);
        if lhs
            .scopes
            .iter()
            .zip(&rhs.scopes)
            .any(|(left, right)| left != right)
        {
            return Err(UnifyError::ScopeStackUnsatisfied);
        }
        let common = lhs.scopes.len().min(rhs.scopes.len());
        self.join_scope_rests(&lhs, &rhs, common).cloned()
    }

    fn join_scope_rests<'s>(
        &mut self,
        lhs: &'s PartialScopeStack,
        rhs: &'s PartialScopeStack,
        common: usize,
    ) -> Result<&'s PartialScopeStack, UnifyError> {
        match (
            lhs.scopes.len().cmp(&rhs.scopes.len()),
            lhs.variable,
            rhs.variable,
        ) {
            (Ordering::Greater, _, _) => self
                .bind_scope_rest(rhs.variable, lhs, common)
                .map(|()| lhs),
            (Ordering::Less, _, _) => self
                .bind_scope_rest(lhs.variable, rhs, common)
                .map(|()| rhs),
            (Ordering::Equal, Some(left), Some(right)) if left != right => {
                self.scopes
                    .insert(right, PartialScopeStack::from_variable(left));
                Ok(lhs)
            }
            (Ordering::Equal, None, Some(right)) => {
                self.scopes.insert(right, PartialScopeStack::default());
                Ok(lhs)
            }
            (Ordering::Equal, Some(left), None) => {
                self.scopes.insert(left, PartialScopeStack::default());
                Ok(rhs)
            }
            (Ordering::Equal, _, _) => Ok(lhs),
        }
    }

    fn bind_scope_rest(
        &mut self,
        variable: Option<ScopeStackVariable>,
        longer: &PartialScopeStack,
        common: usize,
    ) -> Result<(), UnifyError> {
        let variable = variable
            .filter(|variable| longer.variable != Some(*variable))
            .ok_or(UnifyError::ScopeStackUnsatisfied)?;
        let rest = PartialScopeStack::new(
            longer.scopes.iter().skip(common).copied().collect(),
            longer.variable,
        );
        self.scopes.insert(variable, rest);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::{
        PartialBindings, PartialScopeStack, PartialScopedSymbol, PartialSymbolStack, UnifyError,
    };
    use crate::stack::resolve::{ScopeStackVariable, SymbolStackVariable};
    use crate::NodeId;
    use alloc::collections::{BTreeMap, VecDeque};

    fn symbol_variable(id: u32) -> SymbolStackVariable {
        SymbolStackVariable::new(id).unwrap()
    }

    fn scope_variable(id: u32) -> ScopeStackVariable {
        ScopeStackVariable::new(id).unwrap()
    }

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

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

    fn scoped(name: &str, attached: PartialScopeStack) -> PartialSymbolStack<'_> {
        PartialSymbolStack::new(
            VecDeque::from([PartialScopedSymbol::new(name, Some(attached))]),
            None,
        )
    }

    fn unify<'a>(
        lhs: &PartialSymbolStack<'a>,
        rhs: &PartialSymbolStack<'a>,
    ) -> (
        Result<PartialSymbolStack<'a>, UnifyError>,
        PartialBindings<'a>,
    ) {
        let mut bindings = PartialBindings::new();
        let unified = bindings.unify_symbols(lhs, rhs);
        (unified, bindings)
    }

    #[test]
    fn identical_stacks_unify_to_themselves_without_bindings() {
        let stack = symbols(&["a", "b"], Some(1));

        assert_eq!(
            unify(&stack, &stack),
            (Ok(stack.clone()), PartialBindings::new())
        );
    }

    #[test]
    fn different_symbols_do_not_unify() {
        assert_eq!(
            unify(&symbols(&["a"], Some(1)), &symbols(&["b"], Some(2))).0,
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn a_variable_binds_the_symbols_the_other_stack_has_below_the_common_prefix() {
        let (unified, bindings) = unify(&symbols(&["a"], Some(1)), &symbols(&["a", "b"], Some(2)));

        assert_eq!(unified, Ok(symbols(&["a", "b"], Some(2))));
        assert_eq!(
            bindings.symbols,
            BTreeMap::from([(symbol_variable(1), symbols(&["b"], Some(2)))])
        );
    }

    #[test]
    fn a_longer_left_stack_binds_the_right_variable() {
        let (unified, bindings) = unify(&symbols(&["a", "b"], None), &symbols(&["a"], Some(1)));

        assert_eq!(unified, Ok(symbols(&["a", "b"], None)));
        assert_eq!(
            bindings.symbols,
            BTreeMap::from([(symbol_variable(1), symbols(&["b"], None))])
        );
    }

    #[test]
    fn two_variables_over_the_same_symbols_bind_the_right_one_to_the_left_one() {
        let (unified, bindings) = unify(&symbols(&["a"], Some(1)), &symbols(&["a"], Some(2)));

        assert_eq!(unified, Ok(symbols(&["a"], Some(1))));
        assert_eq!(
            bindings.symbols,
            BTreeMap::from([(symbol_variable(2), symbols(&[], Some(1)))])
        );
    }

    #[test]
    fn a_variable_facing_the_end_of_a_stack_binds_the_empty_stack() {
        let (unified, bindings) = unify(&symbols(&["a"], Some(1)), &symbols(&["a"], None));

        assert_eq!(unified, Ok(symbols(&["a"], None)));
        assert_eq!(
            bindings.symbols,
            BTreeMap::from([(symbol_variable(1), PartialSymbolStack::default())])
        );
    }

    #[test]
    fn a_stack_without_variable_does_not_unify_with_a_longer_one() {
        assert_eq!(
            unify(&symbols(&["a"], None), &symbols(&["a", "b"], None)).0,
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn a_variable_cannot_stand_for_symbols_over_itself() {
        assert_eq!(
            unify(&symbols(&[], Some(1)), &symbols(&["a"], Some(1))).0,
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn a_failed_unification_keeps_the_earlier_bindings() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_symbols(&symbols(&[], Some(1)), &symbols(&["b"], None))
            .unwrap();
        let before = bindings.clone();

        assert_eq!(
            bindings.unify_symbols(&symbols(&["a"], Some(2)), &symbols(&["c"], Some(1))),
            Err(UnifyError::SymbolStackUnsatisfied)
        );
        assert_eq!(bindings, before);
    }

    #[test]
    fn a_variable_bound_twice_must_agree_with_its_first_binding() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_symbols(&symbols(&[], Some(1)), &symbols(&["b"], None))
            .unwrap();

        assert_eq!(
            bindings.unify_symbols(&symbols(&[], Some(1)), &symbols(&["c"], None)),
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn an_alias_resolves_to_what_its_target_is_bound_to_later() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_symbols(&symbols(&[], Some(2)), &symbols(&[], Some(1)))
            .unwrap();
        bindings
            .unify_symbols(&symbols(&[], Some(2)), &symbols(&["a"], None))
            .unwrap();

        assert_eq!(
            symbols(&[], Some(1)).apply_bindings(&bindings),
            symbols(&["a"], None)
        );
    }

    #[test]
    fn a_variable_cannot_stand_for_symbols_over_its_alias() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_symbols(&symbols(&[], Some(2)), &symbols(&[], Some(1)))
            .unwrap();

        assert_eq!(
            bindings.unify_symbols(&symbols(&[], Some(2)), &symbols(&["a"], Some(1))),
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn a_variable_bound_through_an_alias_must_agree_with_its_first_binding() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_symbols(&symbols(&[], Some(2)), &symbols(&[], Some(1)))
            .unwrap();
        bindings
            .unify_symbols(&symbols(&[], Some(2)), &symbols(&["a"], None))
            .unwrap();

        assert_eq!(
            bindings.unify_symbols(&symbols(&[], Some(1)), &symbols(&["b"], None)),
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn a_scoped_symbol_does_not_unify_with_a_plain_one() {
        assert_eq!(
            unify(&scoped("a", scopes(&[1], None)), &symbols(&["a"], None)).0,
            Err(UnifyError::SymbolStackUnsatisfied)
        );
    }

    #[test]
    fn attached_scopes_unify_by_binding_their_variables() {
        let (unified, bindings) = unify(
            &scoped("a", scopes(&[1], Some(1))),
            &scoped("a", scopes(&[1, 2], None)),
        );

        assert_eq!(unified, Ok(scoped("a", scopes(&[1, 2], None))));
        assert_eq!(
            bindings.scopes,
            BTreeMap::from([(scope_variable(1), scopes(&[2], None))])
        );
    }

    #[test]
    fn different_attached_scopes_do_not_unify() {
        assert_eq!(
            unify(
                &scoped("a", scopes(&[1], Some(1))),
                &scoped("a", scopes(&[2], Some(2)))
            )
            .0,
            Err(UnifyError::ScopeStackUnsatisfied)
        );
    }

    #[test]
    fn a_scope_variable_binds_the_scopes_below_the_common_prefix() {
        let mut bindings = PartialBindings::new();

        assert_eq!(
            bindings.unify_scopes(&scopes(&[1, 2], Some(2)), &scopes(&[1], Some(1))),
            Ok(scopes(&[1, 2], Some(2)))
        );
        assert_eq!(
            bindings.scopes,
            BTreeMap::from([(scope_variable(1), scopes(&[2], Some(2)))])
        );
    }

    #[test]
    fn a_scope_variable_cannot_stand_for_scopes_over_itself() {
        assert_eq!(
            PartialBindings::new().unify_scopes(&scopes(&[], Some(1)), &scopes(&[1], Some(1))),
            Err(UnifyError::ScopeStackUnsatisfied)
        );
    }

    #[test]
    fn a_scope_alias_resolves_to_what_its_target_is_bound_to_later() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_scopes(&scopes(&[], Some(2)), &scopes(&[], Some(1)))
            .unwrap();
        bindings
            .unify_scopes(&scopes(&[], Some(2)), &scopes(&[3], None))
            .unwrap();

        assert_eq!(
            scopes(&[], Some(1)).apply_bindings(&bindings),
            scopes(&[3], None)
        );
    }

    #[test]
    fn a_scope_variable_cannot_stand_for_scopes_over_its_alias() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify_scopes(&scopes(&[], Some(2)), &scopes(&[], Some(1)))
            .unwrap();

        assert_eq!(
            bindings.unify_scopes(&scopes(&[], Some(2)), &scopes(&[3], Some(1))),
            Err(UnifyError::ScopeStackUnsatisfied)
        );
    }

    #[test]
    fn applying_bindings_replaces_the_variable_with_its_stack() {
        let mut bindings = PartialBindings::new();
        bindings
            .symbols
            .insert(symbol_variable(1), symbols(&["b"], Some(2)));

        assert_eq!(
            symbols(&["a"], Some(1)).apply_bindings(&bindings),
            symbols(&["a", "b"], Some(2))
        );
    }

    #[test]
    fn applying_bindings_replaces_the_variables_of_attached_scopes() {
        let mut bindings = PartialBindings::new();
        bindings
            .scopes
            .insert(scope_variable(1), scopes(&[2], None));

        assert_eq!(
            scoped("a", scopes(&[1], Some(1))).apply_bindings(&bindings),
            scoped("a", scopes(&[1, 2], None))
        );
    }

    #[test]
    fn an_unbound_variable_stays_after_applying_bindings() {
        let stack = symbols(&["a"], Some(1));

        assert_eq!(stack.apply_bindings(&PartialBindings::new()), stack);
    }

    #[test]
    fn an_offset_renumbers_symbol_and_attached_scope_variables_separately() {
        let mut stack = scoped("a", scopes(&[1], Some(1)));
        stack.variable = Some(symbol_variable(1));
        let mut expected = scoped("a", scopes(&[1], Some(4)));
        expected.variable = Some(symbol_variable(3));

        assert_eq!(stack.with_offset(2, 3), Some(expected));
    }

    #[test]
    fn the_largest_scope_variable_is_the_highest_among_attached_scopes() {
        let mut stack = scoped("a", scopes(&[], Some(2)));
        stack
            .symbols
            .push_back(PartialScopedSymbol::new("b", Some(scopes(&[], Some(5)))));
        stack.symbols.push_back(PartialScopedSymbol::new("c", None));

        assert_eq!(stack.largest_scope_variable(), Some(scope_variable(5)));
    }
}