aethellib 0.8.2

Composable text generation primitives over target-specific TOML corpora with provenance tracking.
Documentation
//! # engine
//!
//! Rule execution and composition for text generation.
//!
//! This module provides the runtime that executes named rules against a
//! [`Corpus`] and stores results in a shared generation context.
//!
//! ## Mental Model
//!
//! - A [`Rule`] is one generation step.
//! - The [`Engine`] runs rules in order.
//! - Each rule output is stored in [`GenerationContext::history`] under the
//!   rule name.
//! - Later rules can depend on earlier outputs.
//!
//! ## Key Types
//!
//! - [`Engine`]: orchestrates ordered rule execution
//! - [`Rule`]: trait for executable generation logic
//! - [`InlineRule`]: closure-backed rule adapter
//! - [`GenerationContext`]: shared read state and prior results
//! - [`ComposedValue`]: generated string plus provenance
//!
//! ## Combinators
//!
//! See [`combinators`] for reusable rule builders such as `pick`, `concat`,
//! `chance`, `sequence`, and `weighted_choice`.
//!
//! ## Determinism
//!
//! Deterministic output depends on the RNG you pass to [`Engine::new`].
//! Use a seeded RNG for reproducibility in tests and pipelines.
//!
//! [`Corpus`]: crate::corpus::Corpus
//! [`Rule`]: crate::engine::Rule
//! [`Engine`]: crate::engine::Engine
//! [`GenerationContext`]: crate::engine::GenerationContext
//! [`GenerationContext::history`]: crate::engine::GenerationContext::history
//! [`InlineRule`]: crate::engine::InlineRule
//! [`ComposedValue`]: crate::engine::ComposedValue
//! [`combinators`]: crate::engine::combinators
//! [`Engine::new`]: crate::engine::Engine::new

pub mod combinators;
pub mod error;

use std::collections::HashMap;

use rand::Rng;

use crate::corpus::{Corpus, ValueProvenance};
pub use error::AethelError;

// == ComposedValue ============================================================

#[derive(Debug, Clone)]
/// the final or intermediate result of a generation step.
pub struct ComposedValue {
    pub value: String,
    pub provenance: Vec<ValueProvenance>,
}

impl ComposedValue {
    /// merges two composed values, appending the string and extending the provenance vector.
    pub fn merge(mut self, other: ComposedValue) -> Self {
        self.value.push_str(&other.value);
        self.provenance.extend(other.provenance);
        self
    }
}

// == GenerationContext ========================================================

/// the shared state blackboard passed to every generation rule during execution.
pub struct GenerationContext<'a> {
    /// the corpus available to all rules.
    pub corpus: &'a Corpus,
    /// results of previously executed rules keyed by rule name.
    pub history: HashMap<String, ComposedValue>,
}

impl<'a> GenerationContext<'a> {
    pub fn new(corpus: &'a Corpus) -> Self {
        Self {
            corpus,
            history: HashMap::new(),
        }
    }

    /// returns the result of a previously executed rule, if present.
    pub fn get_previous(&self, key: &str) -> Option<&ComposedValue> {
        self.history.get(key)
    }
}

// == Rule =====================================================================

/// the base interface for all generation logic.
pub trait Rule {
    /// the unique identifier for what this rule generates.
    fn name(&self) -> &str;

    /// executes the generation logic against the current context and rng.
    fn execute<'a>(
        &self,
        ctx: &GenerationContext<'a>,
        rng: &mut dyn Rng,
    ) -> Result<ComposedValue, AethelError>;
}

// == InlineRule ===============================================================

/// wraps a closure as a [`Rule`], allowing inline rule definitions without
/// implementing the trait directly.
pub struct InlineRule<F> {
    name: String,
    logic: F,
}

impl<F> InlineRule<F> {
    pub fn new(name: impl Into<String>, logic: F) -> Self {
        Self {
            name: name.into(),
            logic,
        }
    }
}

impl<F> Rule for InlineRule<F>
where
    F: for<'a> Fn(&GenerationContext<'a>, &mut dyn Rng) -> Result<ComposedValue, AethelError>,
{
    fn name(&self) -> &str {
        &self.name
    }

    fn execute<'a>(
        &self,
        ctx: &GenerationContext<'a>,
        rng: &mut dyn Rng,
    ) -> Result<ComposedValue, AethelError> {
        (self.logic)(ctx, rng)
    }
}

// == Engine ===================================================================

/// orchestrates an ordered sequence of [`Rule`]s against a [`Corpus`].
///
/// determinism depends on the caller supplying a seeded RNG. passing a
/// non-deterministic RNG (e.g. one seeded from system entropy) will produce
/// non-deterministic outputs.
pub struct Engine<'a, R: Rng> {
    corpus: &'a Corpus,
    rng: R,
    rules: Vec<Box<dyn Rule>>,
}

impl<'a, R: Rng> Engine<'a, R> {
    pub fn new(corpus: &'a Corpus, rng: R) -> Self {
        Self {
            corpus,
            rng,
            rules: Vec::new(),
        }
    }

    /// appends a rule to the end of the pipeline.
    pub fn with_rule(mut self, rule: impl Rule + 'static) -> Self {
        self.rules.push(Box::new(rule));
        self
    }

    /// runs every rule in pipeline order, storing each result in the context
    /// history, and returns the populated context.
    pub fn generate(mut self) -> Result<GenerationContext<'a>, AethelError> {
        let mut ctx = GenerationContext::new(self.corpus);

        for rule in self.rules {
            let name = rule.name().to_string();
            let result = rule.execute(&ctx, &mut self.rng)?;
            ctx.history.insert(name, result);
        }

        Ok(ctx)
    }
}

#[cfg(test)]
mod tests {
    use rand::{Rng, SeedableRng, rngs::StdRng};

    use crate::corpus::{Corpus, ValueProvenance};

    use super::{AethelError, ComposedValue, Engine, GenerationContext, InlineRule, Rule};

    fn test_corpus() -> Corpus {
        let raw = r#"
[header]
title = "engine test"
target = "weapon"

[name]
first = ["ash", "birch"]
"#;

        Corpus::builder("weapon")
            .add_str("engine-test", raw)
            .build()
            .expect("corpus should build")
    }

    fn provenance(source_id: &str) -> ValueProvenance {
        ValueProvenance {
            source_id: source_id.to_string(),
            document_title: "doc".to_string(),
            section: "name".to_string(),
            field: "first".to_string(),
        }
    }

    #[test]
    fn composed_value_merge_appends_string_and_provenance() {
        let left = ComposedValue {
            value: "A".to_string(),
            provenance: vec![provenance("left")],
        };
        let right = ComposedValue {
            value: "B".to_string(),
            provenance: vec![provenance("right")],
        };

        let merged = left.merge(right);

        assert_eq!(merged.value, "AB");
        assert_eq!(merged.provenance.len(), 2);
        assert_eq!(merged.provenance[0].source_id, "left");
        assert_eq!(merged.provenance[1].source_id, "right");
    }

    #[test]
    fn generation_context_get_previous_returns_inserted_value() {
        let corpus = test_corpus();
        let mut ctx = GenerationContext::new(&corpus);

        ctx.history.insert(
            "greeting".to_string(),
            ComposedValue {
                value: "hello".to_string(),
                provenance: vec![],
            },
        );

        let found = ctx.get_previous("greeting").expect("key should exist");
        assert_eq!(found.value, "hello");
        assert!(ctx.get_previous("missing").is_none());
    }

    #[test]
    fn inline_rule_exposes_name_and_executes_logic() {
        let corpus = test_corpus();
        let ctx = GenerationContext::new(&corpus);
        let mut rng = StdRng::seed_from_u64(7);

        let rule = InlineRule::new(
            "fixed",
            |_ctx: &GenerationContext<'_>,
             _rng: &mut dyn Rng|
             -> Result<ComposedValue, AethelError> {
                Ok(ComposedValue {
                    value: "ok".to_string(),
                    provenance: vec![],
                })
            },
        );

        assert_eq!(rule.name(), "fixed");
        let output = rule.execute(&ctx, &mut rng).expect("rule should execute");
        assert_eq!(output.value, "ok");
    }

    #[test]
    fn engine_generate_runs_rules_and_tracks_history() {
        let corpus = test_corpus();
        let rng = StdRng::seed_from_u64(42);

        let first = InlineRule::new(
            "first",
            |_ctx: &GenerationContext<'_>,
             _rng: &mut dyn Rng|
             -> Result<ComposedValue, AethelError> {
                Ok(ComposedValue {
                    value: "A".to_string(),
                    provenance: vec![],
                })
            },
        );
        let second = InlineRule::new(
            "second",
            |ctx: &GenerationContext<'_>,
             _rng: &mut dyn Rng|
             -> Result<ComposedValue, AethelError> {
                let base = ctx
                    .get_previous("first")
                    .ok_or_else(|| AethelError::MissingDependency("first".to_string()))?;
                Ok(ComposedValue {
                    value: format!("{}B", base.value),
                    provenance: vec![],
                })
            },
        );

        let ctx = Engine::new(&corpus, rng)
            .with_rule(first)
            .with_rule(second)
            .generate()
            .expect("engine should generate");

        assert_eq!(ctx.history.len(), 2);
        assert_eq!(ctx.get_previous("first").expect("first exists").value, "A");
        assert_eq!(
            ctx.get_previous("second").expect("second exists").value,
            "AB"
        );
    }
}