bb-compiler 0.3.5

//! Pair every NodeProto carrying a `deadline_ns` attribute with an
//! upstream `DeadlineCheck` syscall gate. The check writes to a new
//! sibling `<node>#__trigger_deadline` slot so the protected node
//! fires only after both its original payload AND the deadline gate.
//!
//! Idempotent — re-running on an already-gated graph is a no-op.

use crate::error::CompileError;
use bb_ir::proto::onnx::{AttributeProto, GraphProto, NodeProto, StringStringEntryProto};
use bb_ir::syscall_ids::{
    ATTR_DEADLINE_NS, OP_DEADLINE_CHECK as DEADLINE_OP_TYPE, SYSCALL_DOMAIN as DEADLINE_DOMAIN,
};

/// Idempotence stamp - once a node has been paired with a
/// `DeadlineCheck`, this metadata entry stops the pass from
/// re-inserting on a second run.
pub const GATED_KEY: &str = "ai.bytesandbrains.deadline_gated";

/// Suffix appended to a protected node's name to form the sibling
/// trigger slot that the inserted `DeadlineCheck` writes.
pub const TRIGGER_DEADLINE_SUFFIX: &str = "#__trigger_deadline";

/// Insert a `DeadlineCheck` upstream of every `deadline_ns`-bearing
/// node in `sub_graph`. Pure (the input graph is mutated in place
/// but the function has no side effects beyond that).
///
/// The `DeadlineCheck` writes to a sibling input slot
/// `<protected>#__trigger_deadline`; the engine fires the protected
/// node only after every input (original payloads plus the new
/// trigger slot) has been filled.
pub fn insert_async_deadlines(sub_graph: &mut GraphProto) -> Result<(), CompileError> {
    let mut gates: Vec<NodeProto> = Vec::new();

    for node in sub_graph.node.iter_mut() {
        if metadata_value(node, GATED_KEY).is_some() {
            continue;
        }
        let Some(deadline_ns) = read_deadline(node) else {
            continue;
        };
        // The gate observes the protected node's first input as a
        // proxy for "this node is preparing to fire". When the node
        // has no inputs there is nothing to observe, so skip.
        let Some(trigger_proxy) = node.input.first().cloned() else {
            continue;
        };

        let trigger_slot = format!("{}{TRIGGER_DEADLINE_SUFFIX}", node.name);
        gates.push(build_deadline_check_node(
            &node.name,
            &trigger_proxy,
            &trigger_slot,
            deadline_ns,
        ));

        // Append the sibling deadline-trigger slot; existing inputs
        // stay untouched.
        node.input.push(trigger_slot);
        set_metadata(&mut node.metadata_props, GATED_KEY, "true");
    }

    sub_graph.node.extend(gates);
    Ok(())
}

fn build_deadline_check_node(
    source_name: &str,
    trigger_input: &str,
    gate_output: &str,
    deadline_ns: i64,
) -> NodeProto {
    NodeProto {
        op_type: DEADLINE_OP_TYPE.to_string(),
        domain: DEADLINE_DOMAIN.to_string(),
        name: format!("DeadlineCheck@{source_name}"),
        input: vec![trigger_input.to_string()],
        output: vec![gate_output.to_string()],
        attribute: vec![AttributeProto {
            name: ATTR_DEADLINE_NS.to_string(),
            i: deadline_ns,
            r#type: bb_ir::proto::onnx::attribute_proto::AttributeType::Int as i32,
            ..Default::default()
        }],
        metadata_props: vec![
            StringStringEntryProto {
                key: "ai.bytesandbrains.deadline_source".to_string(),
                value: source_name.to_string(),
            },
            // Mark the inserted gate as already-gated so a second
            // pass run doesn't try to gate its own previous output.
            StringStringEntryProto {
                key: GATED_KEY.to_string(),
                value: "true".to_string(),
            },
        ],
        ..Default::default()
    }
}

fn read_deadline(node: &NodeProto) -> Option<i64> {
    node.attribute
        .iter()
        .find(|a| a.name == ATTR_DEADLINE_NS)
        .map(|a| a.i)
}

fn metadata_value(node: &NodeProto, key: &str) -> Option<String> {
    node.metadata_props
        .iter()
        .find(|p| p.key == key)
        .map(|p| p.value.clone())
}

fn set_metadata(props: &mut Vec<StringStringEntryProto>, key: &str, value: &str) {
    if let Some(existing) = props.iter_mut().find(|p| p.key == key) {
        existing.value = value.to_string();
        return;
    }
    props.push(StringStringEntryProto {
        key: key.to_string(),
        value: value.to_string(),
    });
}

// ──── inventory-published GateContract ─────────────────

/// Insertion contract for the deadline gate. Every node carrying
/// `deadline_ns` must have the sibling `__trigger_deadline` input
/// slot AND a matching `DeadlineCheck` writing to that slot —
/// presence-only checks would miss partial runs.
struct DeadlineCheckContract;

impl crate::gate_contract::GateContract for DeadlineCheckContract {
    fn name(&self) -> &'static str {
        "DeadlineCheck"
    }

    fn assert_inserted(
        &self,
        sub_graph: &bb_ir::proto::onnx::GraphProto,
    ) -> Result<(), CompileError> {
        for node in &sub_graph.node {
            // The DeadlineCheck gate itself carries `deadline_ns`
            // (it propagates the value to the runtime check) AND
            // is marked GATED — it never needs a sibling trigger
            // slot of its own. Skip it.
            if node.op_type == DEADLINE_OP_TYPE || metadata_value(node, GATED_KEY).is_some() {
                continue;
            }
            if node.attribute.iter().any(|a| a.name == ATTR_DEADLINE_NS)
                && !node
                    .input
                    .iter()
                    .any(|i| i.ends_with(TRIGGER_DEADLINE_SUFFIX))
            {
                return Err(CompileError::RuntimeIncomplete {
                    missing: format!(
                        "DeadlineCheck not inserted upstream of `{}` (no sibling `{TRIGGER_DEADLINE_SUFFIX}` input)",
                        node.name,
                    ),
                });
            }
        }
        Ok(())
    }
}

static DEADLINE_CHECK_CONTRACT: DeadlineCheckContract = DeadlineCheckContract;

bb_ir::registry::inventory::submit! {
    crate::gate_contract::GateContractRegistration {
        contract: &DEADLINE_CHECK_CONTRACT,
    }
}