aimx/values/

node.rs

//! Workflow node handle with lazy loading and shared access.
//!
//! Provides [`Node`], a clonable handle to a workflow that loads on first use and
//! shares an immutable snapshot across readers. Mutation is performed on owned
//! `Workflow` values and committed via [`Node::set_workflow`].
//! External locking is required for write coordination.

use crate::{
    aim::{Workflow, WorkflowLike, Writer, WriterLike},
    expressions::Reference,
    inference::{Inference, Pattern, parse_inference},
};
use nom::IResult;
use std::{
    fmt,
    sync::{Arc, RwLock},
};
use anyhow::{Result, anyhow};

/// Internal state for a workflow node.
///
/// `Unloaded` stores the [`Reference`]. `Loaded` stores the shared [`Workflow`].
#[derive(Debug, Clone)]
enum NodeState {
    Pending { inference: Arc<Inference> },
    /// Workflow is not loaded; holds the [`Reference`] used to load it.
    Unloaded { reference: Arc<Reference>, inference: Arc<Inference> },
    /// Workflow is loaded; holds the shared [`Workflow`] snapshot.
    Loaded { workflow: Arc<Workflow>, inference: Arc<Inference> },
}

/// A clonable handle to a workflow with lazy loading and shared reads.
#[derive(Debug, Clone)]
pub struct Node {
    /// Internal state protected by read-write lock
    inner: Arc<RwLock<NodeState>>,
}

impl Default for Node {
    fn default() -> Self {
        let inference = Inference::new(Pattern::Evaluate, None);
        Self::new(Arc::new(inference))
    }
}

impl Node {
    /// Creates a new `Node` in the `Pending` state with the given [`Inference`].
    pub fn new(inference: Arc<Inference>) -> Self {
        Node {
            inner: Arc::new(RwLock::new(NodeState::Pending {
                inference, 
            })),
        }
    }

    pub fn init_new(reference: Arc<Reference>, inference: Arc<Inference>) -> Self {
        let node = Self::new(inference);
        node.init(reference);
        node
    }

    pub fn init_default(reference: Arc<Reference>) -> Self {
        let inference = Inference::new(Pattern::Evaluate, None);
        let node = Self::new(Arc::new(inference));
        node.init(reference);
        node
    }

    pub fn parse(input: &str) -> Result<Self> {
        match parse_node(input) {
            Ok((_, node)) => Ok(node),
            Err(e) => Err(anyhow!("Parse error: {}", e)),
        }
    } 

    /// Initializes the `Node` in the `Unloaded` state with the given [`Reference`].
    pub fn init(&self, reference: Arc<Reference>) {
        // Acquire read lock first
        let read_guard = self.inner.read()
            .expect("Lock poisoned in init() node read.");

        if let NodeState::Pending { inference: _ } = &*read_guard {
            // Need to load, so release read lock and acquire write lock
            drop(read_guard);
            let mut write_guard = self.inner.write()
                .expect("Lock poisoned in init() node write.");
            // Double-check after acquiring write lock
            if let NodeState::Pending { inference} = &*write_guard {
                *write_guard = NodeState::Unloaded { 
                    reference,
                    inference: inference.clone(),
                };
            }
        }
    }

pub fn reference(&self) -> Result<Arc<Reference>> {
       // Acquire read lock first
         let read_guard = self.inner.read()
            .expect("Lock poisoned in reference() node read.");

         match &*read_guard {
             NodeState::Pending { inference: _ } => Err(anyhow!("Node not initialized")),
             NodeState::Unloaded { reference, inference: _ } => Ok(reference.clone()),
             NodeState::Loaded { workflow, inference: _ } => Ok(workflow.locator()),
         }
    }

    /// Returns the inference characteristic for this workflow.
    ///
    /// The model provides an abstraction to decouple inference from reliance on a
    /// specific underlying model. The model performance characteristic reflects one
    /// of six specializations:
    ///
    /// * Fast - optimized for speed and low latency
    /// * Standard - balanced performance for general use
    /// * Thinking - optimized for complex reasoning
    /// * Extraction - optimized for a large context window
    /// * Instruct - optimized for closely following instructions
    /// * Coder - optimized for programming tasks
    /// 
    /// The pattern provides an abstraction for different kinds of workflow tasks.
    /// 
    /// * Evaluate - non-agentic workflow evaluation
    /// * Inference - general inference pattern
    /// * Search - search extraction pattern
    /// * Summarize - map-reduce pattern
    /// * Compose - ReAct composer pattern for accurate annotation based document editing (Crate markdown)
    /// * Debate - debate pattern where two agents each make an assessment; a third judge scores resolution.
pub fn inference(&self) -> Arc<Inference> {
       // Acquire read lock first
         let read_guard = self.inner.read()
            .expect("Lock poisoned in inference() node read.");

         match &*read_guard {
             NodeState::Pending { inference }
             | NodeState::Unloaded { reference: _, inference }
             | NodeState::Loaded { workflow: _, inference } => {
                 inference.clone()
             }
         }
    }

    /// Gets the workflow for read-only access, loading it with `Workflow::load_new` if needed.
    ///
    /// Uses a double-checked lock: fast `RwLock` read when loaded, upgrades to
    /// write lock to perform a single load on transition from `Unloaded`.
    /// Panics on poisoned locks.
    pub fn get_workflow(&self) -> Result<Arc<Workflow>> {
        // Acquire read lock first
        let read_guard = self.inner.read()
            .expect("Lock poisoned in get_workflow() node read.");

        match &*read_guard {
            NodeState::Pending { inference: _ } => {
                Err(anyhow!("Node not initialized"))
            }
            NodeState::Loaded { workflow, inference: _ } => {
                Ok(workflow.clone())
            }
            NodeState::Unloaded { reference: _, inference: _ } => {
                // Need to load, so release read lock and acquire write lock
                drop(read_guard);
                let mut write_guard = self.inner.write()
                    .expect("Lock poisoned in get_workflow() node write.");

                // Double-check after acquiring write lock
                match &*write_guard {
                    NodeState::Pending { inference: _ } => {
                        Err(anyhow!("Node not initialized"))
                    }
                    NodeState::Loaded { workflow, inference: _ } => {
                        Ok(workflow.clone())
                    }
                    NodeState::Unloaded { reference, inference } => {
                        // Load the workflow
                        let workflow = Workflow::open(reference.clone())?;
                        let arc_workflow = Arc::new(workflow);

                        // Update state
                        *write_guard = NodeState::Loaded {
                            workflow: arc_workflow.clone(),
                            inference: inference.clone(),
                        };

                        Ok(arc_workflow)
                    }
                }
            }
        }
    }

    /// Returns the workflow as `Arc<dyn WorkflowLike>`.
    pub fn get_workflow_like(&self) -> Result<Arc<dyn WorkflowLike>> {
        let workflow_like = self.get_workflow()?;
        Ok(workflow_like as Arc<dyn WorkflowLike>)
    }

    /// Returns an owned [`Workflow`] for mutation.
    pub fn get_workflow_mut(&self) -> Result<Workflow> {
        // Fork a copy of the read-only workflow
        let workflow = (*self.get_workflow()?).clone();
        Ok(workflow)
    }

    /// Atomically replace the old workflow with the new workflow.
    ///
    /// Sets state to `Loaded` with a new shared snapshot. Panics on poisoned lock.
    pub fn set_workflow(&self, workflow: Workflow) {
        let mut write_guard = self.inner.write()
            .expect("Lock poisoned in set_workflow() node write.");
        let inference: Arc<Inference> = match &*write_guard {
            NodeState::Pending { inference }
            | NodeState::Unloaded { reference: _, inference }
            | NodeState::Loaded { workflow: _, inference } => {
                inference.clone()
            }
        };
        *write_guard = NodeState::Loaded {
            workflow: Arc::new(workflow),
            inference,
        };
    }

pub fn compact(&self) {
        // Check if loaded
        let read_guard = self.inner.read()
            .expect("Lock poisoned in compact() node read.");

        if let NodeState::Loaded { workflow: _, inference: _ } = &*read_guard {
            drop(read_guard);                
            let mut write_guard = self.inner.write()
                .expect("Lock poisoned in compact() node write.");

            // Double-check pattern
            if let NodeState::Loaded { workflow , inference} = &*write_guard {
                let workflow = workflow.clone();
                if !workflow.is_touched() {
                    let reference = workflow.locator();
                    *write_guard = NodeState::Unloaded {
                        reference,
                        inference: inference.clone(),
                    };
                }
            }
        }
    }

    /// Recursively saves all workflows and instances within this node's workflow hierarchy.
    /// 
    /// This method saves the current workflow if it has pending changes, then recursively
    /// saves all child nodes and instances. It handles both loaded and unloaded workflows
    /// to ensure comprehensive state persistence.
    /// 
    /// # Returns
    /// 
    /// `Ok(())` if all save operations complete successfully, or an error if any
    /// workflow or instance fails to save.
    /// 
    /// # Behavior
    /// 
    /// - Saves the current workflow if it has pending changes (`is_touched()`)
    /// - Recursively saves all child nodes in the workflow
    /// - Saves all instances within the workflow
    /// - Handles unloaded workflows by checking if they need saving
/// - Fails fast on any save error to maintain consistency
 pub fn save_all(&self) -> Result<()> {
        // Check if loaded
        let read_guard = self.inner.read()
            .expect("Lock poisoned in save_all() node read.");

        if let NodeState::Loaded { workflow, inference } = &*read_guard {
            let workflow = (*workflow).clone();
            let inference = inference.clone();
            
            // Always save the workflow if it has changes
            if workflow.is_touched() {
                let mut mutable_workflow = (*workflow).clone();
                mutable_workflow.save()?;
                
                // Recursively save child nodes and instances
                for rule in mutable_workflow.iter_rules() {
                    if let Some(node) = rule.get_node() {
                        node.save_all()?;
                    // Also handle instances within this workflow
                    } else if let Some(instance) = rule.get_instance() {
                        instance.save_all()?;
                    }
                }
                
                // Update the node with the saved workflow
                drop(read_guard);
                let mut write_guard = self.inner.write()
                    .expect("Lock poisoned in save_all() node write.");
                *write_guard = NodeState::Loaded {
                    workflow: Arc::new(mutable_workflow),
                    inference,
                };
            }
        }
        Ok(())
    }

    pub fn print(&self, writer: &mut Writer) {
        let read_guard = self.inner.read()
            .expect("Lock poisoned in print() node read.");
        match &*read_guard {
            NodeState::Pending { inference }
            | NodeState::Unloaded { reference: _, inference }
            | NodeState::Loaded { workflow: _, inference } => {
                inference.print(writer);
            }
        }
    }
}

pub fn parse_node(input: &str) -> IResult<&str, Node> {
    let (input, opt_inference) = parse_inference(input)?;
    let inference = match opt_inference {
        Some(inference) => inference,
        None => Inference::new(Pattern::Evaluate, None),
    };
    Ok((input, Node::new(Arc::new(inference))))
}

impl PartialEq for Node {
    /// Equality is based on shared internal `Arc` identity, not workflow content.
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.inner, &other.inner)
    }
}

impl WriterLike for Node {
    fn write(&self, writer: &mut Writer) {
        self.print(writer);
    }
}

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