conjure_cp_core/rule_engine/

rule.rs

use std::collections::BTreeSet;
use std::fmt::{self, Debug, Display, Formatter};
use std::hash::Hash;
use std::sync::Arc;

use thiserror::Error;

use crate::Model;
use crate::ast::{
    CnfClause, DeclarationKind, DeclarationPtr, Expression, Metadata, Name, SymbolTable,
};

#[derive(Debug, Error)]
pub enum ApplicationError {
    #[error("Rule is not applicable")]
    RuleNotApplicable,

    #[error("Could not calculate the expression domain")]
    DomainError,
}

/// Represents the result of applying a rule to an expression within a model.
///
/// A `RuleEffect` encapsulates the changes made to a model during a rule application.
/// It includes a new expression to replace the original one, an optional top-level constraint
/// to be added to the model, and any updates to the model's symbol table.
///
/// This struct allows for representing side-effects of rule applications, ensuring that
/// all modifications, including symbol table expansions and additional constraints, are
/// accounted for and can be applied to the model consistently.
///
/// # Fields
/// - `new_expression`: The updated [`Expression`] that replaces the original one after applying the rule.
/// - `new_top`: An additional top-level [`Vec<Expression>`] constraint that should be added to the model. If no top-level
///   constraint is needed, this field can be set to an empty vector [`Vec::new()`].
/// - `symbols`: A [`SymbolTable`] containing any new symbol definitions or modifications to be added to the model's
///   symbol table. If no symbols are modified, this field can be set to an empty symbol table.
///
/// # Usage
/// A `RuleEffect` can be created using one of the provided constructors:
/// - [`RuleEffect::new`]: Creates an effect with a new expression, top-level constraint, and symbol modifications.
/// - [`RuleEffect::pure`]: Creates an effect with only a new expression and no side-effects on the symbol table or constraints.
/// - [`RuleEffect::with_symbols`]: Creates an effect with a new expression and symbol table modifications, but no top-level constraint.
/// - [`RuleEffect::with_top`]: Creates an effect with a new expression and a top-level constraint, but no symbol table modifications.
/// - [`RuleEffect::cnf`]: Creates an effect with a new expression, cnf clauses and symbol modifications, but no top-level constraints.
///
/// The `apply` method allows for applying the changes represented by the `RuleEffect` to a [`Model`].
///
/// # Example
/// ```
/// // Need to add an example
/// ```
///
/// # See Also
/// - [`ApplicationResult`]: Represents the result of applying a rule, which may either be a `RuleEffect` or an `ApplicationError`.
/// - [`Model`]: The structure to which the `RuleEffect` changes are applied.
#[non_exhaustive]
#[derive(Clone)]
pub struct RuleEffect {
    pub new_expression: Expression,
    pub new_top: Vec<Expression>,
    pub symbols: SymbolTable,
    pub new_clauses: Vec<CnfClause>,
    /// Shared declarations to update if this effect is selected.
    pub(crate) declaration_updates: Vec<DeclarationUpdate>,
    materialise: Option<DeferredRuleEffect>,
}

/// An in-place update to a shared declaration, applied only after its rule is selected.
#[derive(Clone, Debug)]
pub(crate) struct DeclarationUpdate {
    target: DeclarationPtr,
    replacement_kind: DeclarationKind,
}

impl DeclarationUpdate {
    /// Creates a deferred replacement for a shared declaration's kind.
    fn new(target: DeclarationPtr, replacement_kind: DeclarationKind) -> Self {
        Self {
            target,
            replacement_kind,
        }
    }

    /// Returns the name of the declaration that will be updated.
    pub(crate) fn name(&self) -> Name {
        self.target.name().clone()
    }

    /// Commits the replacement while retaining the declaration's identity.
    pub(crate) fn apply(mut self) {
        let _ = self.target.replace_kind(self.replacement_kind);
    }

    /// Creates an unshared declaration containing the pending replacement.
    fn preview(&self) -> DeclarationPtr {
        DeclarationPtr::new(self.name(), self.replacement_kind.clone())
    }
}

/// Deferred constructor for a concrete rule effect.
type DeferredRuleEffect = Arc<dyn Fn(&SymbolTable) -> RuleEffect + Send + Sync>;

/// The result of applying a rule to an expression.
/// Contains either a set of rule effects or an error.
pub type ApplicationResult = Result<RuleEffect, ApplicationError>;

impl Debug for RuleEffect {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("RuleEffect")
            .field("new_expression", &self.new_expression)
            .field("new_top", &self.new_top)
            .field("symbols", &self.symbols)
            .field("new_clauses", &self.new_clauses)
            .field("declaration_updates", &self.declaration_updates)
            .field("is_deferred", &self.materialise.is_some())
            .finish()
    }
}

impl RuleEffect {
    pub fn new(new_expression: Expression, new_top: Vec<Expression>, symbols: SymbolTable) -> Self {
        Self {
            new_expression,
            new_top,
            symbols,
            new_clauses: Vec::new(),
            declaration_updates: Vec::new(),
            materialise: None,
        }
    }

    /// Represents an effect with no side effects on the model.
    pub fn pure(new_expression: Expression) -> Self {
        Self {
            new_expression,
            new_top: Vec::new(),
            symbols: SymbolTable::new(),
            new_clauses: Vec::new(),
            declaration_updates: Vec::new(),
            materialise: None,
        }
    }

    /// Represents an effect that also modifies the symbol table.
    pub fn with_symbols(new_expression: Expression, symbols: SymbolTable) -> Self {
        Self {
            new_expression,
            new_top: Vec::new(),
            symbols,
            new_clauses: Vec::new(),
            declaration_updates: Vec::new(),
            materialise: None,
        }
    }

    /// Represents an effect that also adds a top-level constraint to the model.
    pub fn with_top(new_expression: Expression, new_top: Vec<Expression>) -> Self {
        Self {
            new_expression,
            new_top,
            symbols: SymbolTable::new(),
            new_clauses: Vec::new(),
            declaration_updates: Vec::new(),
            materialise: None,
        }
    }

    /// Represents an effect that also adds clauses to the model.
    pub fn cnf(
        new_expression: Expression,
        new_clauses: Vec<CnfClause>,
        symbols: SymbolTable,
    ) -> Self {
        Self {
            new_expression,
            new_top: Vec::new(),
            symbols,
            new_clauses,
            declaration_updates: Vec::new(),
            materialise: None,
        }
    }

    /// Defers constructing a concrete effect until the rewriter chooses to apply this rule.
    ///
    /// This is intended for rule effects that allocate fresh names or otherwise depend on global
    /// model state. Applicability checks can return a deferred effect without consuming those
    /// effects; the rewriter calls [`RuleEffect::materialise`] only for the selected rule.
    pub fn deferred(
        materialise: impl Fn(&SymbolTable) -> RuleEffect + Send + Sync + 'static,
    ) -> Self {
        Self {
            new_expression: Expression::Root(Metadata::new(), Vec::new()),
            new_top: Vec::new(),
            symbols: SymbolTable::new(),
            new_clauses: Vec::new(),
            declaration_updates: Vec::new(),
            materialise: Some(Arc::new(materialise)),
        }
    }

    /// Returns the concrete effect for the current symbol table.
    ///
    /// This consumes the selected effect: cloning a concrete effect can duplicate its expression,
    /// top-level constraints, clauses, and speculative symbol table. Before returning, the symbol
    /// snapshot is reduced to the bindings that the effect actually changes.
    pub fn materialise(mut self, symbols: &SymbolTable) -> Self {
        if let Some(materialise) = self.materialise.take() {
            return materialise(symbols).materialise(symbols);
        }

        self.symbols.retain_local_changes_from(symbols);
        self
    }

    pub(crate) fn is_deferred(&self) -> bool {
        self.materialise.is_some()
    }

    /// Adds declaration replacements that are committed only if this effect is selected.
    pub fn with_declaration_updates(
        mut self,
        updates: impl IntoIterator<Item = (DeclarationPtr, DeclarationKind)>,
    ) -> Self {
        self.declaration_updates.extend(
            updates
                .into_iter()
                .map(|(target, kind)| DeclarationUpdate::new(target, kind)),
        );
        self
    }

    /// Iterates over the names changed by deferred declaration updates.
    pub fn updated_declaration_names(&self) -> impl Iterator<Item = Name> + '_ {
        self.declaration_updates.iter().map(DeclarationUpdate::name)
    }

    /// Applies pending declaration replacements to a detached symbol-table preview.
    pub(crate) fn preview_declaration_updates(&self, symbols: &mut SymbolTable) {
        for update in &self.declaration_updates {
            symbols.update_insert(update.preview());
        }
    }

    /// Applies side-effects (e.g. symbol table updates)
    pub fn apply(self, model: &mut Model) {
        debug_assert!(
            self.materialise.is_none(),
            "deferred rule effects must be materialised before being applied"
        );
        for update in self.declaration_updates {
            update.apply();
        }
        model.symbols_mut().extend(self.symbols); // Add new assignments to the symbol table
        model.add_constraints(self.new_top);
        model.add_clauses(self.new_clauses);
    }

    /// Gets symbols added by this effect.
    ///
    /// Walks this effect's own symbols rather than diffing two whole tables. The rewriter asks
    /// this once per applied rule, and most effects are pure, so diffing made every rewrite cost
    /// a clone and a sort of the entire model symbol table.
    pub fn added_symbols(&self, initial_symbols: &SymbolTable) -> BTreeSet<Name> {
        self.symbols
            .iter_local()
            .filter(|(name, _)| initial_symbols.lookup_local(name).is_none())
            .map(|(name, _)| name.clone())
            .collect()
    }

    /// Gets symbols changed by this effect.
    ///
    /// Returns a list of tuples of (name, domain before effect, domain after effect), ordered by
    /// name.
    ///
    /// Walks this effect's symbols for the same reason as [`RuleEffect::added_symbols`]: a symbol
    /// can only have changed if this effect carries its new value.
    pub fn changed_symbols(
        &self,
        initial_symbols: &SymbolTable,
    ) -> Vec<(Name, DeclarationPtr, DeclarationPtr)> {
        let mut changes: Vec<(Name, DeclarationPtr, DeclarationPtr)> = self
            .symbols
            .iter_local()
            .filter_map(|(name, new_value)| {
                let initial_value = initial_symbols.lookup_local(name)?;
                (!new_value.content_eq(&initial_value))
                    .then(|| (name.clone(), initial_value.clone(), new_value.clone()))
            })
            .collect();
        changes.sort_by(|(lhs, ..), (rhs, ..)| lhs.cmp(rhs));
        changes
    }
}

/// The function type used in a [`Rule`].
pub type RuleFn = fn(&Expression, &SymbolTable) -> ApplicationResult;

/// Atomic expression subvariants that can be used as rule prefilters.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum AtomKind {
    /// An `Atomic` expression containing a literal.
    Literal,
    /// An `Atomic` expression containing a reference.
    Reference,
}

/// A complete rule prefilter alternative.
///
/// A rule matches when any of its prefilter alternatives matches the focused expression. This keeps
/// compound filters such as `And / Comprehension` paired instead of forming a cross product with
/// other alternatives in the same rule declaration.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum RulePrefilter {
    /// Focused expression must have this `Expression` variant.
    Variant(usize),
    /// Focused expression must have an immediate child with this `Expression` variant.
    Child { child: usize },
    /// Focused expression must have this variant and an immediate child with `child`'s variant.
    VariantChild { variant: usize, child: usize },
    /// Focused expression must be an `Atomic` expression with this atomic subvariant.
    Atom(AtomKind),
    /// Focused expression must have an immediate `Atomic` child with this atomic subvariant.
    ChildAtom(AtomKind),
}

/// State changes that can invalidate a failed rule application.
///
/// Most rules may become applicable when either their focused expression or the symbol table
/// changes. A small number of Root rules use the expression only as an entry point and decide
/// applicability entirely from declarations. Remembering their failed applications until the
/// symbol table changes avoids repeatedly scanning every declaration after unrelated expression
/// rewrites.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum RuleFailureInvalidation {
    /// Reconsider a failed application after any relevant expression or symbol-table change.
    #[default]
    ExpressionOrSymbols,
    /// Reconsider a failed application only after the symbol table changes.
    SymbolsOnly,
}

/**
 * A rule with a name, application function, and rule sets.
 *
 * # Fields
 * - `name` The name of the rule.
 * - `application` The function to apply the rule.
 * - `rule_sets` A list of rule set names and priorities that this rule is a part of. This is used to populate rulesets at runtime.
 */
#[derive(Clone, Debug)]
pub struct Rule<'a> {
    pub name: &'a str,
    pub application: RuleFn,
    pub rule_sets: &'a [(&'a str, u16)], // (name, priority). At runtime, we add the rule to rulesets
    /// Complete prefilter alternatives this rule applies to, or `None` for universal rules.
    pub prefilters: Option<&'static [RulePrefilter]>,
    /// Which state changes can make a failed application become applicable.
    pub failure_invalidation: RuleFailureInvalidation,
}

impl<'a> Rule<'a> {
    pub const fn new(
        name: &'a str,
        application: RuleFn,
        rule_sets: &'a [(&'static str, u16)],
    ) -> Self {
        Self {
            name,
            application,
            rule_sets,
            prefilters: None,
            failure_invalidation: RuleFailureInvalidation::ExpressionOrSymbols,
        }
    }

    pub fn apply(&self, expr: &Expression, symbols: &SymbolTable) -> ApplicationResult {
        (self.application)(expr, symbols)
    }
}

impl Display for Rule<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name)
    }
}

impl PartialEq for Rule<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.name == other.name
    }
}

impl Eq for Rule<'_> {}

impl Hash for Rule<'_> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.name.hash(state);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{DecisionVariable, Domain, Range, Reference};
    use crate::rule_engine::rewriter_common::{
        snapshot_symbols_after_effect, snapshot_variable_declarations,
    };

    #[test]
    fn declaration_updates_are_applied_only_when_effect_is_committed() {
        let old_domain = Domain::int(vec![Range::Bounded(1, 3)]);
        let new_domain = Domain::int(vec![Range::Bounded(1, 2)]);
        let mut model = Model::new(Default::default());
        let declaration = DeclarationPtr::new_find(Name::user("x"), old_domain.clone());
        let reference = Reference::new(declaration.clone());
        model.symbols_mut().insert(declaration.clone()).unwrap();

        let effect = RuleEffect::pure(Expression::from(1)).with_declaration_updates([(
            declaration,
            DeclarationKind::Find(DecisionVariable::new(new_domain.clone())),
        )]);

        let before = snapshot_variable_declarations(&model.symbols());
        let after = snapshot_symbols_after_effect(&model.symbols(), &effect);
        assert_eq!(before.get(&Name::user("x")).unwrap(), "find x: int(1..3)");
        assert_eq!(after.get(&Name::user("x")).unwrap(), "find x: int(1..2)");
        assert_eq!(reference.domain(), Some(old_domain));
        effect.apply(&mut model);
        assert_eq!(reference.domain(), Some(new_domain));
    }
}