rust-runes 0.1.0

use crate::bitset::BitSet;
use crate::compiled::{path_from_snippet, CompiledRule, MemoSlot, Node, Path as FactPath};
use crate::facts::{Fact, FactValue};
use crate::functions::{self, Function};
use crate::knowledge_base::KnowledgeBase;
use crate::methods::MethodRegistry;
use crate::operators;
use crate::rule::Rule;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::borrow::Cow;
use std::cell::Cell;
use std::collections::HashMap;
use std::fmt::{self, Write};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering};
use std::time::{Duration, Instant};
use thiserror::Error;

/// Maximum number of rule firings per execution, matching Grule's default.
pub const DEFAULT_MAX_CYCLES: usize = 5000;

/// Errors from executing rules.
#[derive(Error, Debug)]
#[non_exhaustive]
pub enum EngineError {
    /// An expression couldn't be evaluated, e.g. a missing field.
    #[error("Evaluation error: {0}")]
    EvaluationError(String),
    /// A rule referred to a fact that doesn't exist.
    #[error("Unknown variable: {0}")]
    UnknownVariable(String),
    /// An operation was applied to values of the wrong type.
    #[error("Type error: {0}")]
    TypeError(String),
    /// Integer `%` by zero.
    #[error("Division by zero")]
    DivisionByZero,
    /// A rule called a function that is neither built in nor registered.
    #[error("Unknown function: {0}")]
    UnknownFunction(String),
    /// A function or method call failed (wrong arguments, or an error it returned).
    #[error("Function error: {0}")]
    FunctionError(String),
    /// `Retract` named a rule that doesn't exist.
    #[error("Unknown rule: {0}")]
    UnknownRule(String),
    /// Rules kept firing past the cycle limit (see `RuleEngine::set_max_cycles`).
    #[error("Rules fired {0} times without settling; a rule's actions may never make its condition false (consider Retract)")]
    MaxCyclesExceeded(usize),
    /// The execution's cancellation flag was set.
    #[error("Execution was cancelled")]
    Cancelled,
    /// The execution's timeout elapsed.
    #[error("Execution timed out")]
    TimedOut,
    /// An error raised while evaluating or executing a specific rule.
    #[error("Rule '{rule}' failed in its {phase} scope: {source}")]
    InRule {
        /// The rule's name.
        rule: String,
        /// Whether the error came from the rule's condition or its actions.
        phase: RulePhase,
        /// The underlying error.
        #[source]
        source: Box<EngineError>,
    },
}

impl EngineError {
    /// The underlying error, without the rule it occurred in.
    pub fn root(&self) -> &EngineError {
        match self {
            EngineError::InRule { source, .. } => source.root(),
            other => other,
        }
    }

    fn in_rule(self, rule: &Rule, phase: RulePhase) -> EngineError {
        EngineError::InRule {
            rule: rule.name.clone(),
            phase,
            source: Box::new(self),
        }
    }
}

/// Which part of a rule an error came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RulePhase {
    /// The rule's condition.
    When,
    /// The rule's actions.
    Then,
}

impl fmt::Display for RulePhase {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            RulePhase::When => "when",
            RulePhase::Then => "then",
        })
    }
}

/// Receives notifications as the engine runs, like Grule's `GruleEngineListener`.
/// All methods do nothing by default.
pub trait EngineListener: Send + Sync {
    /// A new cycle is starting (cycles are numbered from 1).
    fn begin_cycle(&self, _cycle: u64) {}
    /// A rule's condition was evaluated; `matched` is whether it held.
    fn evaluate_rule(&self, _cycle: u64, _rule: &Rule, _matched: bool) {}
    /// A rule was chosen and its actions are about to run.
    fn execute_rule(&self, _cycle: u64, _rule: &Rule) {}
}

/// What an execution did.
#[derive(Debug, Clone)]
pub struct ExecutionResult {
    /// The rules that fired, in order (a rule may appear more than once).
    pub rules_fired: Vec<String>,
    /// How long the execution took, in milliseconds.
    pub execution_time_ms: u128,
}

impl ExecutionResult {
    /// An empty result.
    pub fn new() -> Self {
        Self {
            rules_fired: Vec::new(),
            execution_time_ms: 0,
        }
    }
}

impl Default for ExecutionResult {
    fn default() -> Self {
        Self::new()
    }
}

/// Per-execution state changed by engine functions like `Retract` and `Complete`.
#[derive(Default)]
struct ExecutionState {
    /// Rules retracted by the current rule's actions
    retracted: Vec<String>,
    completed: bool,
    /// Paths named by `Changed` / `Forget` in the current rule's actions
    changed: Vec<FactPath>,
}

/// Per-execution rule state, by salience position.
struct Agenda {
    /// Rules whose condition result is unknown (always includes rules whose results
    /// can't be reused)
    dirty: BitSet,
    /// Rules known to match
    matching: BitSet,
    retracted: BitSet,
}

/// Memoized results of shared boolean sub-expressions for one execution (see
/// `compiled::Interner`). Uses `Cell`s so evaluation can fill it through `&self`.
struct Memo {
    slots: Vec<Cell<u8>>,
}

impl Memo {
    const UNKNOWN: u8 = 0;

    fn new(slot_count: usize) -> Self {
        Memo {
            slots: vec![Cell::new(Self::UNKNOWN); slot_count],
        }
    }

    /// A memo that never stores anything.
    fn disabled() -> Self {
        Memo { slots: Vec::new() }
    }

    fn is_enabled(&self) -> bool {
        !self.slots.is_empty()
    }

    fn get(&self, slot: MemoSlot) -> Option<bool> {
        match self.slots[slot].get() {
            Self::UNKNOWN => None,
            value => Some(value == 2),
        }
    }

    fn set(&self, slot: MemoSlot, value: bool) {
        self.slots[slot].set(if value { 2 } else { 1 });
    }

    fn forget(&self, slot: MemoSlot) {
        self.slots[slot].set(Self::UNKNOWN);
    }
}

/// One step of a path into a fact, e.g. `.customer` or `[0]`.
enum Key {
    Field(String),
    Index(i64),
}

/// Runs rules against facts, as Grule's engine does.
///
/// An engine holds a set of rules (its [`KnowledgeBase`]) plus registered functions,
/// methods and listeners. [`execute`](Self::execute) takes `&self`, so one engine can be
/// shared between threads and run concurrently.
pub struct RuleEngine {
    knowledge_base: KnowledgeBase,
    functions: HashMap<String, Box<Function>>,
    methods: MethodRegistry,
    max_cycles: usize,
    return_err_on_failed_rule_evaluation: bool,
    listeners: Vec<Box<dyn EngineListener>>,
    condition_caching: bool,
}

impl RuleEngine {
    /// An engine with no rules, Grule's default cycle limit (5000), and failed
    /// conditions counting as false.
    pub fn new() -> Self {
        Self {
            knowledge_base: KnowledgeBase::new(),
            functions: HashMap::new(),
            methods: MethodRegistry::default(),
            max_cycles: DEFAULT_MAX_CYCLES,
            return_err_on_failed_rule_evaluation: false,
            listeners: Vec::new(),
            condition_caching: true,
        }
    }

    /// Adds a rule. Fails if a rule with the same name exists.
    pub fn add_rule(&mut self, rule: Rule) -> Result<(), String> {
        self.knowledge_base.add_rule(rule)
    }

    /// Parses GRL text and adds all of its rules. Returns how many were added.
    pub fn add_rules_from_str(&mut self, grl: &str) -> Result<usize, String> {
        self.knowledge_base.add_rules_from_str(grl)
    }

    /// Loads all rules from a `.grl` file. Returns how many were added.
    pub fn add_rules_from_file(&mut self, path: impl AsRef<Path>) -> Result<usize, String> {
        self.knowledge_base.add_rules_from_file(path)
    }

    /// Loads all `.grl` files under `dir`, recursively. Returns how many rules were added.
    pub fn add_rules_from_dir(&mut self, dir: impl AsRef<Path>) -> Result<usize, String> {
        self.knowledge_base.add_rules_from_dir(dir)
    }

    /// Sets the maximum number of rule firings per execution before it fails
    /// with `EngineError::MaxCyclesExceeded`.
    pub fn set_max_cycles(&mut self, max_cycles: usize) {
        self.max_cycles = max_cycles;
    }

    /// By default, as in Grule, a condition that fails to evaluate (a missing field,
    /// a type error, ...) is logged and treated as false. When enabled, execution
    /// stops with the error instead.
    pub fn set_return_err_on_failed_rule_evaluation(&mut self, enabled: bool) {
        self.return_err_on_failed_rule_evaluation = enabled;
    }

    /// Reusing condition results between cycles is on by default and never changes
    /// which rules fire; turning it off is only useful for comparison and debugging.
    pub fn set_condition_caching(&mut self, enabled: bool) {
        self.condition_caching = enabled;
    }

    /// Adds a listener notified as executions run. With any listener attached, every
    /// rule is evaluated every cycle (as in Grule) so that listeners see each evaluation.
    pub fn add_listener(&mut self, listener: impl EngineListener + 'static) {
        self.listeners.push(Box::new(listener));
    }

    /// Makes `function` callable from rules as `name(args...)`.
    pub fn register_function<F>(&mut self, name: &str, function: F) -> Result<(), String>
    where
        F: Fn(&[FactValue]) -> Result<FactValue, String> + Send + Sync + 'static,
    {
        if functions::is_reserved(name) {
            return Err(format!("'{}' is a built-in function", name));
        }
        if self.functions.contains_key(name) {
            return Err(format!("Function '{}' is already registered", name));
        }
        self.functions.insert(name.to_string(), Box::new(function));
        Ok(())
    }

    /// Makes `method` callable from rules as `receiver.name(args...)` on any value
    /// that has no built-in method of that name. This is the equivalent of calling a
    /// method on a Go fact in Grule: when called as a statement in a rule's `then`
    /// scope on a fact path (`Order.AddItem(x);`), the method receives the stored
    /// value and its changes persist. Anywhere else it receives a copy.
    pub fn register_method<F>(&mut self, name: &str, method: F) -> Result<(), String>
    where
        F: Fn(&mut FactValue, &[FactValue]) -> Result<FactValue, String> + Send + Sync + 'static,
    {
        self.methods.register(name, Box::new(method))
    }

    /// Like `register_method`, but the method works on your own type `T`, the
    /// counterpart of a method on a Go fact struct. Around each call the receiver is
    /// converted into `T` and back with serde (so `T` is typically the type the fact was
    /// created from with `Fact::from_serialize`), the arguments are decoded into `A`,
    /// and the result `R` is converted into a fact value.
    ///
    /// `A` is a tuple for fixed arguments (`(f64,)`, `(String, i64)`), `()` for none,
    /// or a `Vec` for any number. Return `Ok(())` for no result (it becomes `nil`).
    ///
    /// The same name can be registered for several types; a call goes to the first
    /// whose type the receiver converts into, then to an untyped method of that name.
    /// Fields the receiver has that `T` doesn't are kept.
    ///
    /// ```
    /// # use rust_runes::RuleEngine;
    /// # #[derive(serde::Serialize, serde::Deserialize)]
    /// # struct Loan { amount: f64, escrow: f64 }
    /// let mut engine = RuleEngine::new();
    /// engine.register_typed_method("CalculateEscrow", |loan: &mut Loan, (): ()| {
    ///     loan.escrow = loan.amount * 0.01;
    ///     Ok(())
    /// })?;
    /// engine.register_typed_method("Payment", |loan: &mut Loan, (rate,): (f64,)| {
    ///     Ok(loan.amount * rate / 12.0)
    /// })?;
    /// // GRL: then Loan.CalculateEscrow(); Loan.Monthly = Loan.Payment(0.06);
    /// # Ok::<(), String>(())
    /// ```
    pub fn register_typed_method<T, A, R, F>(&mut self, name: &str, method: F) -> Result<(), String>
    where
        T: Serialize + DeserializeOwned + 'static,
        A: DeserializeOwned + 'static,
        R: Serialize + 'static,
        F: Fn(&mut T, A) -> Result<R, String> + Send + Sync + 'static,
    {
        self.methods.register_typed(name, method)
    }

    /// Runs rules until no condition holds, as Grule does. Each cycle evaluates every
    /// rule's condition, then fires the highest-salience rule among those that matched
    /// (ties go to the rule added first). Rules removed with `Retract` stay out for the
    /// rest of this execution, and `Complete()` ends it once the current rule's actions
    /// have run.
    ///
    /// Two optimizations keep this from re-evaluating everything every cycle, without
    /// changing which rules fire:
    /// - A condition's result is reused until a rule's actions (or `Changed`/`Forget`)
    ///   touch a fact path it reads. Conditions using `Now()`, logging or user
    ///   functions/methods are always re-evaluated.
    /// - Once the cycle's highest-salience match is found, lower rules are skipped
    ///   unless evaluating them has observable effects: a listener is attached,
    ///   evaluation errors are returned, or the condition logs or calls user code.
    pub fn execute(
        &self,
        facts: &mut HashMap<String, Fact>,
    ) -> Result<ExecutionResult, EngineError> {
        self.run(facts, &Cancellation::default())
    }

    /// Like `execute`, but stops with `EngineError::Cancelled` once `cancel` is set,
    /// checked before each rule evaluation. Set it from another thread.
    pub fn execute_with_cancel(
        &self,
        facts: &mut HashMap<String, Fact>,
        cancel: &AtomicBool,
    ) -> Result<ExecutionResult, EngineError> {
        let cancellation = Cancellation {
            flag: Some(cancel),
            deadline: None,
        };
        self.run(facts, &cancellation)
    }

    /// Like `execute`, but stops with `EngineError::TimedOut` once `timeout` has
    /// elapsed, checked before each rule evaluation (like a Go context deadline).
    /// A single slow function call can still overrun it.
    pub fn execute_with_timeout(
        &self,
        facts: &mut HashMap<String, Fact>,
        timeout: Duration,
    ) -> Result<ExecutionResult, EngineError> {
        let cancellation = Cancellation {
            flag: None,
            deadline: Some(Instant::now() + timeout),
        };
        self.run(facts, &cancellation)
    }

    fn run(
        &self,
        facts: &mut HashMap<String, Fact>,
        cancel: &Cancellation,
    ) -> Result<ExecutionResult, EngineError> {
        let start_time = Instant::now();
        let mut result = ExecutionResult::new();
        let mut state = ExecutionState::default();
        let kb = &self.knowledge_base;
        let n = kb.len();

        // Rule state by salience position: `dirty` rules have no known result;
        // `agenda` holds rules known to match
        let mut agenda = Agenda {
            dirty: BitSet::new(n, true),
            matching: BitSet::new(n, false),
            retracted: BitSet::new(n, false),
        };
        let memo = Memo::new(if self.condition_caching {
            kb.interner().slot_count()
        } else {
            0
        });

        // Every evaluation is observable with a listener or when errors are returned,
        // so evaluate every rule each cycle as Grule does
        let full_evaluation =
            !self.listeners.is_empty() || self.return_err_on_failed_rule_evaluation;
        // Otherwise only conditions with effects must run every cycle
        let mut always = BitSet::new(n, false);
        let mut always_positions = Vec::new();
        if !full_evaluation && !kb.all_conditions_skippable(!self.methods.is_empty()) {
            for position in 0..n {
                if !self.is_skippable(kb.entry_at(position).2) {
                    always.insert(position);
                    always_positions.push(position);
                    agenda.dirty.remove(position);
                }
            }
        }

        let mut cycle: u64 = 0;
        loop {
            cancel.check()?;
            self.listeners.iter().for_each(|l| l.begin_cycle(cycle + 1));

            let selected = if full_evaluation {
                self.select_by_full_evaluation(cycle + 1, &mut agenda, facts, &memo, cancel)?
            } else {
                self.select_incrementally(&mut agenda, &always_positions, facts, &memo, cancel)?
            };

            let Some(position) = selected else { break };
            cycle += 1;
            if cycle > self.max_cycles as u64 {
                return Err(EngineError::MaxCyclesExceeded(self.max_cycles));
            }

            let (_, rule, compiled) = kb.entry_at(position);
            self.listeners
                .iter()
                .for_each(|l| l.execute_rule(cycle, rule));
            for action in &compiled.then {
                self.execute_action(action, facts, &mut state)
                    .map_err(|e| e.in_rule(rule, RulePhase::Then))?;
            }
            result.rules_fired.push(rule.name.clone());
            if state.completed {
                break;
            }

            for name in state.retracted.drain(..) {
                let position = kb.position_of(kb.index_of(&name).expect("checked by Retract"));
                agenda.retracted.insert(position);
                agenda.matching.remove(position);
                agenda.dirty.remove(position);
            }
            for written in compiled.writes.iter().chain(&state.changed) {
                self.invalidate(written, &mut agenda, &always, &memo);
            }
            state.changed.clear();
        }

        result.execution_time_ms = start_time.elapsed().as_millis();
        Ok(result)
    }

    /// Finds the highest-salience matching rule by evaluating every rule (using known
    /// results where still valid) and notifying listeners of each.
    fn select_by_full_evaluation(
        &self,
        cycle: u64,
        agenda: &mut Agenda,
        facts: &HashMap<String, Fact>,
        memo: &Memo,
        cancel: &Cancellation,
    ) -> Result<Option<usize>, EngineError> {
        let mut selected = None;
        for position in 0..self.knowledge_base.len() {
            cancel.check()?;
            if agenda.retracted.contains(position) {
                continue;
            }
            let (_, rule, compiled) = self.knowledge_base.entry_at(position);
            let matched = if agenda.dirty.contains(position) || !self.is_cacheable(compiled) {
                self.evaluate_into(position, rule, compiled, agenda, facts, memo)?
            } else {
                agenda.matching.contains(position)
            };
            self.listeners
                .iter()
                .for_each(|l| l.evaluate_rule(cycle, rule, matched));
            if matched && selected.is_none() {
                selected = Some(position);
            }
        }
        Ok(selected)
    }

    /// Finds the highest-salience matching rule while evaluating as little as possible:
    /// the best known match is the first rule on the agenda, so only rules above it
    /// whose results are unknown need evaluating (plus conditions with effects).
    fn select_incrementally(
        &self,
        agenda: &mut Agenda,
        always: &[usize],
        facts: &HashMap<String, Fact>,
        memo: &Memo,
        cancel: &Cancellation,
    ) -> Result<Option<usize>, EngineError> {
        let mut best = agenda.matching.next_from(0);

        for &position in always {
            if agenda.retracted.contains(position) {
                continue;
            }
            cancel.check()?;
            let (_, rule, compiled) = self.knowledge_base.entry_at(position);
            if self.test_rule(rule, compiled, facts, memo)? && best.is_none_or(|b| position < b) {
                best = Some(position);
            }
        }

        let mut cursor = 0;
        while let Some(position) = agenda.dirty.next_from(cursor) {
            if best.is_some_and(|b| position >= b) {
                break;
            }
            cursor = position + 1;
            cancel.check()?;
            let (_, rule, compiled) = self.knowledge_base.entry_at(position);
            if self.evaluate_into(position, rule, compiled, agenda, facts, memo)? {
                best = Some(position);
                break;
            }
        }
        Ok(best)
    }

    /// Evaluates a rule's condition and, if its result can be reused, records it.
    fn evaluate_into(
        &self,
        position: usize,
        rule: &Rule,
        compiled: &CompiledRule,
        agenda: &mut Agenda,
        facts: &HashMap<String, Fact>,
        memo: &Memo,
    ) -> Result<bool, EngineError> {
        let matched = self.test_rule(rule, compiled, facts, memo)?;
        if self.is_cacheable(compiled) {
            agenda.dirty.remove(position);
            if matched {
                agenda.matching.insert(position);
            }
        }
        Ok(matched)
    }

    /// Whether a rule's condition result can be reused until its inputs change.
    fn is_cacheable(&self, compiled: &CompiledRule) -> bool {
        self.condition_caching && compiled.deterministic && !self.calls_user_methods(compiled)
    }

    /// Whether a rule's condition can go unevaluated without anyone noticing.
    fn is_skippable(&self, compiled: &CompiledRule) -> bool {
        !compiled.has_effects && !self.calls_user_methods(compiled)
    }

    fn calls_user_methods(&self, compiled: &CompiledRule) -> bool {
        self.any_user_method(&compiled.condition_methods)
    }

    fn any_user_method(&self, names: &[String]) -> bool {
        !self.methods.is_empty() && names.iter().any(|name| self.methods.contains(name))
    }

    /// Forgets results that depend on anything under `written`: the conditions that
    /// read it go back to unknown, and so do memoized sub-expressions that read it.
    fn invalidate(&self, written: &[String], agenda: &mut Agenda, always: &BitSet, memo: &Memo) {
        let kb = &self.knowledge_base;
        kb.readers().overlapping(written, |index| {
            let position = kb.position_of(index);
            if !agenda.retracted.contains(position) && !always.contains(position) {
                agenda.matching.remove(position);
                agenda.dirty.insert(position);
            }
        });
        if memo.is_enabled() {
            kb.interner()
                .slot_readers()
                .overlapping(written, |slot| memo.forget(slot));
        }
    }

    /// Returns the rules whose conditions currently hold, highest salience first,
    /// without running any of them (Grule's `FetchMatchingRules`).
    pub fn fetch_matching_rules(
        &self,
        facts: &HashMap<String, Fact>,
    ) -> Result<Vec<&Rule>, EngineError> {
        let kb = &self.knowledge_base;
        let memo = Memo::new(if self.condition_caching {
            kb.interner().slot_count()
        } else {
            0
        });
        let mut matching = Vec::new();
        for position in 0..kb.len() {
            let (_, rule, compiled) = kb.entry_at(position);
            if self.test_rule(rule, compiled, facts, &memo)? {
                matching.push(rule);
            }
        }
        Ok(matching)
    }

    /// Evaluates a rule's condition, applying the failed-evaluation policy.
    fn test_rule(
        &self,
        rule: &Rule,
        compiled: &CompiledRule,
        facts: &HashMap<String, Fact>,
        memo: &Memo,
    ) -> Result<bool, EngineError> {
        match self.evaluate_condition(&compiled.when, facts, "rule condition", memo) {
            Ok(matched) => Ok(matched),
            Err(e) if self.return_err_on_failed_rule_evaluation => {
                Err(e.in_rule(rule, RulePhase::When))
            }
            Err(e) => {
                log::warn!("Failed testing condition for rule {}: {}", rule.name, e);
                Ok(false)
            }
        }
    }

    fn evaluate_condition(
        &self,
        node: &Node,
        facts: &HashMap<String, Fact>,
        operation: &str,
        memo: &Memo,
    ) -> Result<bool, EngineError> {
        operators::boolean(&*self.eval(node, facts, memo)?, operation)
    }

    /// Evaluates without memoization, for actions: facts change between statements.
    fn evaluate_expression(
        &self,
        node: &Node,
        facts: &HashMap<String, Fact>,
    ) -> std::result::Result<FactValue, EngineError> {
        self.eval(node, facts, &Memo::disabled())
            .map(Cow::into_owned)
    }

    /// Returns a memoized boolean, or computes and memoizes it.
    fn memoized<'f>(
        &self,
        slot: Option<MemoSlot>,
        memo: &Memo,
        compute: impl FnOnce() -> Result<bool, EngineError>,
    ) -> Result<Cow<'f, FactValue>, EngineError> {
        let slot = slot.filter(|&slot| {
            memo.is_enabled()
                && !self.any_user_method(&self.knowledge_base.interner().slot_info(slot).methods)
        });
        if let Some(value) = slot.and_then(|slot| memo.get(slot)) {
            return Ok(Cow::Owned(FactValue::Boolean(value)));
        }
        let value = compute()?;
        if let Some(slot) = slot {
            memo.set(slot, value);
        }
        Ok(Cow::Owned(FactValue::Boolean(value)))
    }

    /// Evaluates `node`, borrowing from `facts` (and from the rule's constants) where
    /// possible so that reading `order.customer.age` doesn't copy `order`.
    fn eval<'f>(
        &self,
        node: &'f Node,
        facts: &'f HashMap<String, Fact>,
        memo: &Memo,
    ) -> Result<Cow<'f, FactValue>, EngineError> {
        match node {
            Node::Const(value) => Ok(Cow::Borrowed(value)),

            Node::Variable(name) => facts
                .get(name)
                .map(|fact| Cow::Borrowed(&fact.value))
                .ok_or_else(|| EngineError::UnknownVariable(name.clone())),

            Node::Field(base, field) => project(self.eval(base, facts, memo)?, |obj| match obj {
                FactValue::Object(obj) => obj.get(field).ok_or_else(|| {
                    EngineError::EvaluationError(format!("Field '{}' not found", field))
                }),
                other => Err(EngineError::TypeError(format!(
                    "Cannot access field '{}' on {}",
                    field,
                    other.type_name()
                ))),
            }),

            Node::Index(base, index_node) => {
                let index = self.eval(index_node, facts, memo)?;
                project(self.eval(base, facts, memo)?, |container| {
                    match (container, &*index) {
                        (FactValue::Array(items), FactValue::Integer(n)) => {
                            Ok(&items[array_index(*n, items.len())?])
                        }
                        (FactValue::Object(obj), FactValue::String(key)) => {
                            obj.get(key).ok_or_else(|| {
                                EngineError::EvaluationError(format!("Key '{}' not found", key))
                            })
                        }
                        // Maps with integer or boolean keys store them as strings
                        (
                            FactValue::Object(obj),
                            key @ (FactValue::Integer(_) | FactValue::Boolean(_)),
                        ) => {
                            let key = operators::to_text(key);
                            obj.get(&key).ok_or_else(|| {
                                EngineError::EvaluationError(format!("Key '{}' not found", key))
                            })
                        }
                        _ => Err(EngineError::TypeError(format!(
                            "Cannot index {} with {}",
                            container.type_name(),
                            index.type_name()
                        ))),
                    }
                })
            }

            Node::Builtin {
                name,
                function,
                args,
            } => function(name, &self.evaluate_args(args, facts, memo)?)
                .map(Cow::Owned)
                .map_err(EngineError::FunctionError),

            Node::EngineCall(name, _) => Err(EngineError::EvaluationError(format!(
                "{}() can only be used as a statement in a rule's then block",
                name
            ))),

            Node::UserCall(name, args) => {
                let args = self.evaluate_args(args, facts, memo)?;
                let function = self
                    .functions
                    .get(name)
                    .ok_or_else(|| EngineError::UnknownFunction(name.clone()))?;
                function(&args)
                    .map(Cow::Owned)
                    .map_err(|e| EngineError::FunctionError(format!("{}(): {}", name, e)))
            }

            Node::Method(target_node, name, args) => {
                let target = self.eval(target_node, facts, memo)?;
                let args = self.evaluate_args(args, facts, memo)?;
                if let Some(result) = functions::call_method(&target, name, &args) {
                    return result.map(Cow::Owned).map_err(EngineError::FunctionError);
                }
                // Methods that may change their receiver run on a copy here
                let mut receiver = target.into_owned();
                self.call_method_mut(&mut receiver, name, &args)
                    .map(Cow::Owned)
            }

            Node::Binary(op, left, right) => Ok(Cow::Owned(op(
                &*self.eval(left, facts, memo)?,
                &*self.eval(right, facts, memo)?,
            )?)),
            Node::Compare {
                test,
                left,
                right,
                memo: slot,
            } => self.memoized(*slot, memo, || {
                let ordering = operators::compare(
                    &*self.eval(left, facts, memo)?,
                    &*self.eval(right, facts, memo)?,
                )?;
                Ok(ordering.is_some_and(test))
            }),
            Node::Equal {
                negate,
                left,
                right,
                memo: slot,
            } => self.memoized(*slot, memo, || {
                let equal = operators::equal(
                    &*self.eval(left, facts, memo)?,
                    &*self.eval(right, facts, memo)?,
                )?;
                Ok(equal != *negate)
            }),

            // && and || short-circuit, so the right side is only evaluated when needed
            Node::And {
                left,
                right,
                memo: slot,
            } => self.memoized(*slot, memo, || {
                Ok(self.evaluate_condition(left, facts, "logical AND", memo)?
                    && self.evaluate_condition(right, facts, "logical AND", memo)?)
            }),
            Node::Or {
                left,
                right,
                memo: slot,
            } => self.memoized(*slot, memo, || {
                Ok(self.evaluate_condition(left, facts, "logical OR", memo)?
                    || self.evaluate_condition(right, facts, "logical OR", memo)?)
            }),
            Node::Not { inner, memo: slot } => self.memoized(*slot, memo, || {
                Ok(!self.evaluate_condition(inner, facts, "negation", memo)?)
            }),

            Node::Assign(..) => Err(EngineError::EvaluationError(
                "Assignment cannot be used as a value".to_string(),
            )),
        }
    }

    fn evaluate_args(
        &self,
        args: &[Node],
        facts: &HashMap<String, Fact>,
        memo: &Memo,
    ) -> Result<Vec<FactValue>, EngineError> {
        args.iter()
            .map(|arg| self.eval(arg, facts, memo).map(Cow::into_owned))
            .collect()
    }

    /// Calls a method that may change its receiver: a built-in like `Append`,
    /// or a user-registered method.
    fn call_method_mut(
        &self,
        receiver: &mut FactValue,
        name: &str,
        args: &[FactValue],
    ) -> Result<FactValue, EngineError> {
        if let Some(result) = functions::call_method_mut(receiver, name, args) {
            return result.map_err(EngineError::FunctionError);
        }
        let type_name = receiver.type_name();
        match self.methods.call(receiver, name, args) {
            Some(result) => {
                result.map_err(|e| EngineError::FunctionError(format!("{}(): {}", name, e)))
            }
            None => Err(EngineError::FunctionError(format!(
                "Unknown method '{}' on {}",
                name, type_name
            ))),
        }
    }

    fn execute_action(
        &self,
        action: &Node,
        facts: &mut HashMap<String, Fact>,
        state: &mut ExecutionState,
    ) -> std::result::Result<(), EngineError> {
        match action {
            Node::Assign(target, value_node) => {
                let value = self.evaluate_expression(value_node, facts)?;
                self.assign(target, value, facts)
            }

            Node::EngineCall(name, args) if name == "Retract" => {
                let [arg] = args.as_slice() else {
                    return Err(EngineError::FunctionError(format!(
                        "Retract() takes 1 argument(s), got {}",
                        args.len()
                    )));
                };
                let rule_name = match self.evaluate_expression(arg, facts)? {
                    FactValue::String(s) => s,
                    other => {
                        return Err(EngineError::TypeError(format!(
                            "Retract() expects a rule name, got {}",
                            other.type_name()
                        )))
                    }
                };
                if self.knowledge_base.get_rule(&rule_name).is_none() {
                    return Err(EngineError::UnknownRule(rule_name));
                }
                state.retracted.push(rule_name);
                Ok(())
            }

            Node::EngineCall(_, args) => {
                // Complete()
                if !args.is_empty() {
                    return Err(EngineError::FunctionError(format!(
                        "Complete() takes 0 argument(s), got {}",
                        args.len()
                    )));
                }
                state.completed = true;
                Ok(())
            }

            // Changed("Fact.X") / Forget("Fact.X"): something changed that the engine
            // couldn't see, so conditions reading it must be re-evaluated
            Node::Builtin { name, args, .. } if name == "Changed" || name == "Forget" => {
                self.evaluate_expression(action, facts)?;
                if let FactValue::String(snippet) = self.evaluate_expression(&args[0], facts)? {
                    state.changed.push(path_from_snippet(&snippet));
                }
                Ok(())
            }

            // A method called on a fact path operates on the stored value, so
            // `Order.Items.Append(x);` changes the fact
            Node::Method(target, name, args) if target.is_assignable() => {
                let args = self.evaluate_args(args, facts, &Memo::disabled())?;
                let current = self.eval(target, facts, &Memo::disabled())?;
                if let Some(result) = functions::call_method(&current, name, &args) {
                    // A built-in that never changes its receiver
                    return result.map(|_| ()).map_err(EngineError::FunctionError);
                }
                let receiver = self.resolve_mut(target, facts)?;
                self.call_method_mut(receiver, name, &args).map(|_| ())
            }

            // Called for their effect; the result is discarded
            Node::Builtin { .. } | Node::UserCall(..) | Node::Method(..) => {
                self.evaluate_expression(action, facts)?;
                Ok(())
            }

            _ => Err(EngineError::EvaluationError(
                "Invalid action expression".to_string(),
            )),
        }
    }

    /// Assigns `value` to a variable, or to a field/element reached through it.
    fn assign(
        &self,
        target: &Node,
        value: FactValue,
        facts: &mut HashMap<String, Fact>,
    ) -> Result<(), EngineError> {
        let (root, keys) = self.target_path(target, facts)?;

        let Some((last, parent_keys)) = keys.split_last() else {
            facts.insert(root.clone(), Fact::new(root, value));
            return Ok(());
        };

        // Existing fields and elements keep their type, as Go struct fields do
        let slot = |current: Option<&FactValue>, value| match current {
            Some(current) => operators::coerce_assignment(current, value).map_err(|e| {
                EngineError::TypeError(format!("{}: {}", describe_path(&root, &keys), e))
            }),
            None => Ok(value),
        };
        match (walk_mut(facts, &root, &keys, parent_keys.len())?, last) {
            (FactValue::Object(obj), key @ (Key::Field(_) | Key::Index(_))) => {
                let field = object_key(key);
                let value = slot(obj.get(&field), value)?;
                obj.insert(field, value);
                Ok(())
            }
            (FactValue::Array(items), Key::Index(n)) => {
                let i = array_index(*n, items.len())?;
                items[i] = slot(Some(&items[i]), value)?;
                Ok(())
            }
            (other, key) => Err(key_mismatch(other, key, &describe_path(&root, parent_keys))),
        }
    }

    /// Finds the stored value a path like `order.items[0]` refers to.
    fn resolve_mut<'f>(
        &self,
        target: &Node,
        facts: &'f mut HashMap<String, Fact>,
    ) -> Result<&'f mut FactValue, EngineError> {
        let (root, keys) = self.target_path(target, facts)?;
        walk_mut(facts, &root, &keys, keys.len())
    }

    /// Splits a path like `order.items[i].price` into its root variable and the keys
    /// leading from it. Index expressions are evaluated here, before facts are
    /// borrowed mutably.
    fn target_path(
        &self,
        target: &Node,
        facts: &HashMap<String, Fact>,
    ) -> Result<(String, Vec<Key>), EngineError> {
        match target {
            Node::Variable(name) => Ok((name.clone(), Vec::new())),
            Node::Field(inner, field) => {
                let (root, mut keys) = self.target_path(inner, facts)?;
                keys.push(Key::Field(field.clone()));
                Ok((root, keys))
            }
            Node::Index(inner, index) => {
                let (root, mut keys) = self.target_path(inner, facts)?;
                keys.push(match self.evaluate_expression(index, facts)? {
                    FactValue::Integer(n) => Key::Index(n),
                    FactValue::String(s) => Key::Field(s),
                    FactValue::Boolean(b) => Key::Field(b.to_string()),
                    other => {
                        return Err(EngineError::TypeError(format!(
                            "Cannot index with {}",
                            other.type_name()
                        )))
                    }
                });
                Ok((root, keys))
            }
            _ => Err(EngineError::EvaluationError(
                "Invalid assignment target".to_string(),
            )),
        }
    }

    /// The engine's rules.
    pub fn get_knowledge_base(&self) -> &KnowledgeBase {
        &self.knowledge_base
    }
}

impl Default for RuleEngine {
    fn default() -> Self {
        Self::new()
    }
}

/// When to stop an execution early: a flag set by the caller, a deadline, or both.
#[derive(Default)]
struct Cancellation<'a> {
    flag: Option<&'a AtomicBool>,
    deadline: Option<Instant>,
}

impl Cancellation<'_> {
    fn check(&self) -> Result<(), EngineError> {
        if self
            .flag
            .is_some_and(|flag| flag.load(AtomicOrdering::Relaxed))
        {
            return Err(EngineError::Cancelled);
        }
        if self
            .deadline
            .is_some_and(|deadline| Instant::now() >= deadline)
        {
            return Err(EngineError::TimedOut);
        }
        Ok(())
    }
}

/// Follows the first `depth` keys from the fact `root`, returning the value reached.
/// Every step must already exist.
fn walk_mut<'f>(
    facts: &'f mut HashMap<String, Fact>,
    root: &str,
    keys: &[Key],
    depth: usize,
) -> Result<&'f mut FactValue, EngineError> {
    let fact = facts
        .get_mut(root)
        .ok_or_else(|| EngineError::UnknownVariable(root.to_string()))?;
    let mut current = &mut fact.value;
    for (step, key) in keys[..depth].iter().enumerate() {
        current = match (current, key) {
            (FactValue::Object(obj), key @ (Key::Field(_) | Key::Index(_))) => {
                obj.get_mut(&object_key(key)).ok_or_else(|| {
                    EngineError::EvaluationError(format!(
                        "Field '{}' not found",
                        describe_path(root, &keys[..=step])
                    ))
                })?
            }
            (FactValue::Array(items), Key::Index(n)) => {
                let i = array_index(*n, items.len())?;
                &mut items[i]
            }
            (other, key) => {
                return Err(key_mismatch(
                    other,
                    key,
                    &describe_path(root, &keys[..step]),
                ))
            }
        };
    }
    Ok(current)
}

/// Narrows a value to a part of it (a field or element), borrowing when the value
/// itself is borrowed and cloning just the part when it was computed.
fn project<'f>(
    value: Cow<'f, FactValue>,
    part: impl for<'v> FnOnce(&'v FactValue) -> Result<&'v FactValue, EngineError>,
) -> Result<Cow<'f, FactValue>, EngineError> {
    match value {
        Cow::Borrowed(value) => part(value).map(Cow::Borrowed),
        Cow::Owned(value) => part(&value).map(|p| Cow::Owned(p.clone())),
    }
}

/// The key an object is accessed with: fields by name, and numeric indexes by their
/// decimal form, which is how serde stores maps with integer keys.
fn object_key(key: &Key) -> String {
    match key {
        Key::Field(field) => field.clone(),
        Key::Index(n) => n.to_string(),
    }
}

fn array_index(n: i64, len: usize) -> Result<usize, EngineError> {
    if n < 0 {
        Err(EngineError::EvaluationError(format!(
            "Invalid array index: {}",
            n
        )))
    } else if n as usize >= len {
        Err(EngineError::EvaluationError(format!(
            "Index {} out of bounds for array of length {}",
            n, len
        )))
    } else {
        Ok(n as usize)
    }
}

fn describe_path(root: &str, keys: &[Key]) -> String {
    let mut path = root.to_string();
    for key in keys {
        match key {
            Key::Field(field) => write!(path, ".{}", field),
            Key::Index(n) => write!(path, "[{}]", n),
        }
        .unwrap();
    }
    path
}

fn key_mismatch(value: &FactValue, key: &Key, path: &str) -> EngineError {
    EngineError::TypeError(match key {
        Key::Field(field) => format!(
            "Cannot access field '{}' on {} '{}'",
            field,
            value.type_name(),
            path
        ),
        Key::Index(n) => format!("Cannot index {} '{}' with [{}]", value.type_name(), path, n),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::GrlParser;

    fn obj(fields: Vec<(&str, FactValue)>) -> FactValue {
        FactValue::Object(
            fields
                .into_iter()
                .map(|(k, v)| (k.to_string(), v))
                .collect(),
        )
    }

    fn order_facts() -> HashMap<String, Fact> {
        let order = obj(vec![
            (
                "customer",
                obj(vec![
                    ("age", FactValue::Integer(30)),
                    (
                        "address",
                        obj(vec![("city", FactValue::String("Porto".to_string()))]),
                    ),
                ]),
            ),
            ("total", FactValue::Integer(100)),
        ]);
        HashMap::from([("order".to_string(), Fact::new("order".to_string(), order))])
    }

    fn run(grl: &str, facts: &mut HashMap<String, Fact>) -> Result<ExecutionResult, EngineError> {
        let mut engine = RuleEngine::new();
        engine
            .add_rule(GrlParser::new().parse_rule(grl).unwrap())
            .unwrap();
        engine.execute(facts)
    }

    fn customer_field<'a>(facts: &'a HashMap<String, Fact>, field: &str) -> Option<&'a FactValue> {
        facts["order"]
            .get_field("customer")?
            .as_object()?
            .get(field)
    }

    #[test]
    fn test_nested_assignment_updates_deep_field() {
        let mut facts = order_facts();
        run(
            r#"rule R { when order.customer.address.city == "Porto"
                        then order.customer.address.city = "Lisbon"; order.customer.age += 1; }"#,
            &mut facts,
        )
        .unwrap();

        assert_eq!(
            customer_field(&facts, "address")
                .unwrap()
                .as_object()
                .unwrap()["city"],
            FactValue::String("Lisbon".to_string())
        );
        assert_eq!(customer_field(&facts, "age"), Some(&FactValue::Integer(31)));
    }

    #[test]
    fn test_nested_assignment_creates_leaf_field() {
        let mut facts = order_facts();
        run(
            r#"rule R { when true then order.customer.vip = true; Retract("R"); }"#,
            &mut facts,
        )
        .unwrap();
        assert_eq!(
            customer_field(&facts, "vip"),
            Some(&FactValue::Boolean(true))
        );
    }

    #[test]
    fn test_nested_assignment_missing_intermediate_field() {
        let mut facts = order_facts();
        let err = run(
            "rule R { when true then order.shipping.city = \"X\"; }",
            &mut facts,
        )
        .unwrap_err();
        assert!(err.to_string().contains("order.shipping"), "{}", err);
    }

    #[test]
    fn test_nested_assignment_through_non_object() {
        let mut facts = order_facts();
        let err = run(
            "rule R { when true then order.total.cents = 5; }",
            &mut facts,
        )
        .unwrap_err();
        assert!(matches!(err.root(), EngineError::TypeError(_)), "{}", err);
        assert!(err.to_string().contains("order.total"), "{}", err);
    }

    #[test]
    fn test_nested_assignment_unknown_root() {
        let mut facts = order_facts();
        let err = run("rule R { when true then invoice.total = 1; }", &mut facts).unwrap_err();
        assert!(matches!(err.root(), EngineError::UnknownVariable(name) if name == "invoice"));
    }
}