gollum-ir 0.4.0

//! IR-level representation of planning and reasoning states.

use crate::metadata::IrMetadata;
use crate::term::IrTerm;

/// IR-level representation of a planning or reasoning state.
///
/// Encodes the world state at a given point in a planning or reasoning task.
/// The `facts` vector represents all known true propositions (atoms, structures,
/// lists, etc.) at that state. Optional `metadata` can hold probabilities, plan
/// costs, or confidence scores, useful for probabilistic or cost-sensitive planning.
///
/// States can transition to successor states by applying `IrAction` effects.
///
/// # Example
///
/// ```rust
/// use gollum_ir::{IrState, IrTerm, IrMetadata};
///
/// let initial_state = IrState {
///     facts: vec![
///         IrTerm::Structure {
///             name: "at".into(),
///             args: vec![IrTerm::Atom("robot".into()), IrTerm::Atom("room1".into())],
///         },
///         IrTerm::Structure {
///             name: "connected".into(),
///             args: vec![IrTerm::Atom("room1".into()), IrTerm::Atom("room2".into())],
///         },
///     ],
///     metadata: Some(IrMetadata {
///         probability: Some(1.0),
///         ..IrMetadata::default()
///     }),
/// };
///
/// assert_eq!(initial_state.facts.len(), 2);
/// ```
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct IrState {
    /// Current known facts or fluents in the world.
    pub facts: Vec<IrTerm>,

    /// Optional metadata for state probabilities, costs, or annotations.
    #[serde(default)]
    pub metadata: Option<IrMetadata>,
}

impl IrState {
    /// Create a new empty state with no facts and no metadata.
    pub fn new() -> Self {
        Self { facts: vec![], metadata: None }
    }

    /// Create a state from a vector of facts.
    pub fn with_facts(facts: Vec<IrTerm>) -> Self {
        Self { facts, metadata: None }
    }

    /// Create a state from a vector of facts and metadata.
    pub fn with_facts_and_metadata(facts: Vec<IrTerm>, metadata: IrMetadata) -> Self {
        Self { facts, metadata: Some(metadata) }
    }

    /// Check if a given fact is present in this state.
    pub fn contains(&self, fact: &IrTerm) -> bool {
        self.facts.iter().any(|f| f == fact)
    }

    /// Add a fact to this state.
    pub fn add_fact(&mut self, fact: IrTerm) {
        self.facts.push(fact);
    }

    /// Remove all occurrences of a fact from this state.
    pub fn remove_fact(&mut self, fact: &IrTerm) {
        self.facts.retain(|f| f != fact);
    }

    /// Get the number of facts in this state.
    pub fn fact_count(&self) -> usize {
        self.facts.len()
    }

    /// Check if this state is empty (no facts).
    pub fn is_empty(&self) -> bool {
        self.facts.is_empty()
    }
}

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

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

    fn atom(s: &str) -> IrTerm {
        IrTerm::Atom(s.into())
    }

    fn structure(name: &str, args: Vec<IrTerm>) -> IrTerm {
        IrTerm::Structure { name: name.into(), args }
    }

    #[test]
    fn test_new_state_is_empty() {
        let state = IrState::new();
        assert!(state.is_empty());
        assert_eq!(state.fact_count(), 0);
        assert!(state.metadata.is_none());
    }

    #[test]
    fn test_with_facts_creates_state() {
        let facts = vec![atom("at"), atom("connected")];
        let state = IrState::with_facts(facts.clone());
        assert_eq!(state.facts, facts);
        assert_eq!(state.fact_count(), 2);
        assert!(state.metadata.is_none());
    }

    #[test]
    fn test_with_facts_and_metadata() {
        let facts = vec![atom("at")];
        let metadata = IrMetadata { probability: Some(0.8), ..IrMetadata::default() };
        let state = IrState::with_facts_and_metadata(facts.clone(), metadata.clone());
        assert_eq!(state.facts, facts);
        assert_eq!(state.metadata, Some(metadata));
    }

    #[test]
    fn test_contains_fact_true() {
        let fact = atom("at");
        let state = IrState::with_facts(vec![fact.clone()]);
        assert!(state.contains(&fact));
    }

    #[test]
    fn test_contains_fact_false() {
        let state = IrState::with_facts(vec![atom("at")]);
        assert!(!state.contains(&atom("connected")));
    }

    #[test]
    fn test_contains_structure() {
        let fact = structure("at", vec![atom("robot"), atom("room1")]);
        let state = IrState::with_facts(vec![fact.clone()]);
        assert!(state.contains(&fact));
    }

    #[test]
    fn test_add_fact() {
        let mut state = IrState::new();
        let fact = atom("at");
        state.add_fact(fact.clone());
        assert_eq!(state.fact_count(), 1);
        assert!(state.contains(&fact));
    }

    #[test]
    fn test_add_multiple_facts() {
        let mut state = IrState::new();
        state.add_fact(atom("at"));
        state.add_fact(atom("connected"));
        state.add_fact(atom("clear"));
        assert_eq!(state.fact_count(), 3);
    }

    #[test]
    fn test_remove_fact() {
        let fact = atom("at");
        let mut state = IrState::with_facts(vec![fact.clone(), atom("connected")]);
        state.remove_fact(&fact);
        assert!(!state.contains(&fact));
        assert_eq!(state.fact_count(), 1);
    }

    #[test]
    fn test_remove_all_occurrences() {
        let fact = atom("at");
        let mut state = IrState::with_facts(vec![fact.clone(), atom("connected"), fact.clone()]);
        state.remove_fact(&fact);
        assert!(!state.contains(&fact));
        assert_eq!(state.fact_count(), 1);
    }

    #[test]
    fn test_remove_nonexistent_fact() {
        let mut state = IrState::with_facts(vec![atom("at")]);
        state.remove_fact(&atom("connected"));
        assert_eq!(state.fact_count(), 1);
    }

    #[test]
    fn test_default_is_empty() {
        let state = IrState::default();
        assert!(state.is_empty());
    }

    #[test]
    fn test_clone_and_eq() {
        let facts = vec![atom("at"), atom("connected")];
        let state1 = IrState::with_facts(facts);
        let state2 = state1.clone();
        assert_eq!(state1, state2);
    }

    #[test]
    fn test_serde_roundtrip_empty() {
        let state = IrState::new();
        let s = ron::to_string(&state).unwrap();
        let back: IrState = ron::from_str(&s).unwrap();
        assert_eq!(state, back);
    }

    #[test]
    fn test_serde_roundtrip_with_facts() {
        let facts = vec![
            atom("at"),
            structure("connected", vec![atom("a"), atom("b")]),
        ];
        let state = IrState::with_facts(facts);
        let s = ron::to_string(&state).unwrap();
        let back: IrState = ron::from_str(&s).unwrap();
        assert_eq!(state.facts, back.facts);
    }

    #[test]
    fn test_serde_roundtrip_with_metadata() {
        let metadata = IrMetadata { probability: Some(0.75), ..IrMetadata::default() };
        let state = IrState::with_facts_and_metadata(vec![atom("at")], metadata.clone());
        let s = ron::to_string(&state).unwrap();
        let back: IrState = ron::from_str(&s).unwrap();
        assert_eq!(state, back);
    }

    #[test]
    fn test_example_from_spec() {
        let state = IrState {
            facts: vec![
                structure("at", vec![atom("robot"), atom("room1")]),
                structure("connected", vec![atom("room1"), atom("room2")]),
            ],
            metadata: Some(IrMetadata { probability: Some(1.0), ..IrMetadata::default() }),
        };
        assert_eq!(state.fact_count(), 2);
        assert!(state
            .contains(&structure("at", vec![atom("robot"), atom("room1")])));
    }
}