fluidattacks-blends-domain 0.40.0

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

use crate::stack::resolve::SymbolStackVariable;

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

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

/// Two partial symbol stacks have no stack in common.
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[non_exhaustive]
pub struct UnifyError;

impl<'a> PartialSymbolStack<'a> {
    /// `symbols`, top first, over `variable`.
    #[must_use]
    pub fn new(symbols: VecDeque<&'a str>, 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 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.symbols.clone(), variable))
    }

    /// This stack with its variable replaced by what `bindings` binds it to.
    #[must_use]
    pub fn apply_bindings(&self, bindings: &PartialBindings<'a>) -> Self {
        self.variable
            .and_then(|variable| bindings.symbols.get(&variable))
            .map_or_else(
                || self.clone(),
                |bound| {
                    let bound = bound.apply_bindings(bindings);
                    Self::new(
                        self.symbols.iter().chain(&bound.symbols).copied().collect(),
                        bound.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 [`UnifyError`] when they have no stack in common, leaving the bindings unchanged.
    pub fn unify(
        &mut self,
        lhs: &PartialSymbolStack<'a>,
        rhs: &PartialSymbolStack<'a>,
    ) -> Result<PartialSymbolStack<'a>, UnifyError> {
        let lhs = lhs.apply_bindings(self);
        let rhs = rhs.apply_bindings(self);
        if lhs
            .symbols
            .iter()
            .zip(&rhs.symbols)
            .any(|(left, right)| left != right)
        {
            return Err(UnifyError);
        }
        let common = lhs.symbols.len().min(rhs.symbols.len());
        let unified = self.join_rests(&lhs, &rhs, common)?.clone();
        Ok(unified.apply_bindings(self))
    }

    fn join_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_rest(rhs.variable, lhs, common).map(|()| lhs),
            (Ordering::Less, _, _) => self.bind_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_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)?;
        let rest = PartialSymbolStack::new(
            longer.symbols.iter().skip(common).copied().collect(),
            longer.variable,
        );
        self.symbols.insert(variable, rest);
        Ok(())
    }
}

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

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

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

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

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

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

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

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

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

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

    #[test]
    fn an_alias_resolves_to_what_its_target_is_bound_to_later() {
        let mut bindings = PartialBindings::new();
        bindings
            .unify(&symbols(&[], Some(2)), &symbols(&[], Some(1)))
            .unwrap();
        bindings
            .unify(&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(&[], Some(2)), &symbols(&[], Some(1)))
            .unwrap();

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

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

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

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

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

    #[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_the_variable_up() {
        assert_eq!(
            symbols(&["a"], Some(1)).with_offset(2),
            Some(symbols(&["a"], Some(3)))
        );
    }

    #[test]
    fn an_offset_past_the_last_number_gives_no_stack() {
        assert_eq!(symbols(&["a"], Some(u32::MAX)).with_offset(1), None);
    }
}