open-standards/joule-code/joule-code-rs/crates/joule-code-cascade/src/lib.rs

//! The Joule Code cascade (SPEC §1.1, §5, `spec/cascade.md`).
//!
//! A [`Cascade`] dispatches a [`Request`] through its [`Stage`]s in ascending
//! tier order and returns the first [`Resolution`] produced — the
//! stop-at-first-close rule. The agent schema that implements the tiers is
//! JouleClaw; this crate fixes the *binding* of that cascade to coding
//! requests, including the integer-only, deterministic tier-selection that the
//! `cascade_vectors.json` conformance suite (C4) checks.

use std::collections::HashMap;
use std::sync::Mutex;

use joule_code_cargo::CodingOp;
use joule_code_receipt::{Energy, Provenance, Receipt};
use joule_code_schema::Tier;

/// A coding request entering the cascade.
pub struct Request {
    /// The content address of the tree state the request is against.
    pub prior: String,
    /// The natural-language or structured request.
    pub intent: String,
}

/// The outcome of a stage closing a request (SPEC §5).
#[derive(Clone)]
pub struct Resolution {
    /// The produced operation, or `None` for a non-mutating answer.
    pub op: Option<CodingOp>,
    /// A non-mutating answer (e.g. a query), or `None`.
    pub answer: Option<String>,
    /// The energy receipt; `receipt.tier` is the closing tier.
    pub receipt: Receipt,
}

/// One tier's implementation (SPEC §5).
pub trait Stage {
    /// The tier this stage serves.
    fn tier(&self) -> Tier;
    /// Attempt to close the request. `Some` means closed at `self.tier()`.
    fn try_resolve(&self, req: &Request) -> Option<Resolution>;
}

/// An ordered set of stages forming the energy-tiered cascade.
///
/// Invariant: stages are held in strictly ascending [`Tier`] order, so calling
/// them in order and returning the first `Some` is exactly the stop-at-first-
/// close rule. Construct with [`Cascade::new`], which sorts and de-duplicates by
/// tier.
pub struct Cascade {
    stages: Vec<Box<dyn Stage>>,
}

impl Cascade {
    /// Build a cascade from stages in any order. Stages are sorted by tier;
    /// if two stages claim the same tier, the first one given wins (a runtime
    /// has one implementation per tier).
    pub fn new(stages: impl IntoIterator<Item = Box<dyn Stage>>) -> Self {
        let mut stages: Vec<Box<dyn Stage>> = stages.into_iter().collect();
        stages.sort_by_key(|s| s.tier());
        stages.dedup_by_key(|s| s.tier());
        Cascade { stages }
    }

    /// The tiers present, in ascending order.
    pub fn tiers(&self) -> Vec<Tier> {
        self.stages.iter().map(|s| s.tier()).collect()
    }

    /// Dispatch: call stages in tier order, return the first that closes the
    /// request. A higher tier is never consulted once a lower one closes
    /// (SPEC §1.1). Returns `None` if no stage closes the request.
    pub fn resolve(&self, req: &Request) -> Option<Resolution> {
        for stage in &self.stages {
            if let Some(res) = stage.try_resolve(req) {
                debug_assert_eq!(res.receipt.tier, stage.tier());
                return Some(res);
            }
        }
        None
    }

    /// Pure tier-selection used by the C4 conformance suite: given, for each
    /// present tier, whether its stage *would* close the request, return the
    /// tier that actually closes it under stop-at-first-close. Deterministic
    /// and integer-only.
    pub fn select_closing_tier(present: &[Tier], closes: impl Fn(Tier) -> bool) -> Option<Tier> {
        let mut tiers = present.to_vec();
        tiers.sort();
        tiers.dedup();
        tiers.into_iter().find(|&t| closes(t))
    }
}

/// A reference **L0 (Cache)** stage: a content-addressed cache of resolved
/// coding operations keyed by `(prior, intent)`. A hit closes the request at L0
/// (picojoule class) without consulting any higher tier — the cascade's cheapest
/// path (SPEC §1.1). Interior-mutable so a runtime can populate it after a
/// higher tier closes a request. A third party may use this or substitute its
/// own L0 stage.
#[derive(Default)]
pub struct L0Cache {
    entries: Mutex<HashMap<(String, String), CodingOp>>,
}

impl L0Cache {
    /// An empty cache.
    pub fn new() -> Self {
        Self::default()
    }

    /// Record the op that closed `(prior, intent)` so an identical later request
    /// closes at L0.
    pub fn insert(&self, prior: String, intent: String, op: CodingOp) {
        if let Ok(mut entries) = self.entries.lock() {
            entries.insert((prior, intent), op);
        }
    }

    /// Number of cached operations.
    pub fn len(&self) -> usize {
        self.entries.lock().map(|e| e.len()).unwrap_or(0)
    }

    /// Whether the cache is empty.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl Stage for L0Cache {
    fn tier(&self) -> Tier {
        Tier::L0
    }

    fn try_resolve(&self, req: &Request) -> Option<Resolution> {
        let op = self
            .entries
            .lock()
            .ok()?
            .get(&(req.prior.clone(), req.intent.clone()))
            .cloned()?;
        let op_id = op.op_id().ok()?;
        Some(Resolution {
            op: Some(op),
            answer: None,
            // A cache hit performs no inference; report a picojoule-class figure
            // honestly tagged as an estimate (1 µJ is the smallest unit the
            // receipt carries; the true cost is far below it).
            receipt: Receipt {
                op_id,
                tier: Tier::L0,
                energy: Energy {
                    measured_uj: None,
                    estimated_uj: 1,
                    provenance: Provenance::Estimator,
                },
                sequence: 0,
                nonce: [0u8; 8],
            },
        })
    }
}

/// A stage backed by a closure. Lets a runtime express any tier — `L1` Lawful
/// (deterministic refactors), `L2` Embed (retrieval), `L3` Model (a local model
/// such as the reference consumer's vLLM provider) — by supplying a resolver,
/// without this crate depending on a specific runtime or model. The closure
/// returns `Some(Resolution)` when it closes the request, `None` to fall
/// through to the next tier. This is how a consumer binds the JouleClaw agent
/// schema (or any other) to the cascade.
pub struct FnStage<F> {
    tier: Tier,
    resolve: F,
}

impl<F> FnStage<F>
where
    F: Fn(&Request) -> Option<Resolution> + Send + Sync,
{
    /// A stage serving `tier`, closing requests via `resolve`.
    pub fn new(tier: Tier, resolve: F) -> Self {
        Self { tier, resolve }
    }
}

impl<F> Stage for FnStage<F>
where
    F: Fn(&Request) -> Option<Resolution> + Send + Sync,
{
    fn tier(&self) -> Tier {
        self.tier
    }

    fn try_resolve(&self, req: &Request) -> Option<Resolution> {
        (self.resolve)(req)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use joule_code_cargo::{Anchor, Edit};
    use joule_code_receipt::{Energy, Provenance};
    use joule_code_schema::{ContentId, ALG_ED25519};

    fn sample_op(prior: &str) -> CodingOp {
        CodingOp {
            alg: ALG_ED25519,
            prior: ContentId(prior.into()),
            edits: vec![Edit {
                path: "src/lib.rs".into(),
                anchor: Anchor::ByteRange { start: 0, end: 0 },
                replacement: b"pub fn add() {}".to_vec(),
            }],
            author: [0u8; 32],
            intent: "add add()".into(),
            tier: Tier::L3,
        }
    }

    #[test]
    fn fn_stage_l3_closes_when_lower_tiers_miss() {
        // An L0 cache (empty) in front of an L3 model stage: L0 misses, L3
        // closes. Demonstrates a model tier expressed via FnStage without any
        // runtime dependency.
        let l0 = L0Cache::new();
        let l3 = FnStage::new(Tier::L3, |req: &Request| {
            Some(Resolution {
                op: Some(sample_op(&req.prior)),
                answer: None,
                receipt: Receipt {
                    op_id: ContentId("b3:fromL3".into()),
                    tier: Tier::L3,
                    energy: Energy {
                        measured_uj: None,
                        estimated_uj: 2_000_000,
                        provenance: Provenance::Estimator,
                    },
                    sequence: 0,
                    nonce: [0u8; 8],
                },
            })
        });
        let cascade = Cascade::new([
            Box::new(l0) as Box<dyn Stage>,
            Box::new(l3) as Box<dyn Stage>,
        ]);
        let res = cascade
            .resolve(&Request {
                prior: "b3:AAAA".into(),
                intent: "write a function".into(),
            })
            .expect("resolved");
        assert_eq!(res.receipt.tier, Tier::L3);
    }

    #[test]
    fn l0_cache_miss_then_hit_closes_at_l0() {
        let cache = L0Cache::new();
        let req = Request {
            prior: "b3:AAAA".into(),
            intent: "add add()".into(),
        };
        // Miss: nothing cached yet.
        assert!(cache.try_resolve(&req).is_none());
        // A higher tier closed the request; record its op.
        cache.insert(req.prior.clone(), req.intent.clone(), sample_op("b3:AAAA"));
        // Hit: identical request now closes at L0, picojoule class.
        let res = cache.try_resolve(&req).expect("cache hit");
        assert_eq!(res.receipt.tier, Tier::L0);
        assert_eq!(res.receipt.energy.estimated_uj, 1);
        assert!(res.op.is_some());
    }

    #[test]
    fn cascade_stops_at_l0_when_cache_hits() {
        // An L0 cache (populated) in front of nothing else: a hit means the
        // cascade closes at L0 without any higher tier.
        let cache = L0Cache::new();
        cache.insert("b3:AAAA".into(), "add add()".into(), sample_op("b3:AAAA"));
        let cascade = Cascade::new([Box::new(cache) as Box<dyn Stage>]);
        let res = cascade
            .resolve(&Request {
                prior: "b3:AAAA".into(),
                intent: "add add()".into(),
            })
            .expect("resolved");
        assert_eq!(res.receipt.tier, Tier::L0);
    }

    #[test]
    fn stops_at_first_closing_tier() {
        let present = [Tier::L0, Tier::L2, Tier::L3, Tier::L4];
        let closing = Cascade::select_closing_tier(&present, |t| {
            matches!(t, Tier::L3 | Tier::L4)
        });
        assert_eq!(closing, Some(Tier::L3));
    }

    #[test]
    fn l4_only_runtime_is_valid() {
        let closing = Cascade::select_closing_tier(&[Tier::L4], |_| true);
        assert_eq!(closing, Some(Tier::L4));
    }

    #[test]
    fn nothing_closes_returns_none() {
        let present = [Tier::L0, Tier::L2, Tier::L3];
        assert_eq!(Cascade::select_closing_tier(&present, |_| false), None);
    }
}