algeon-query-planning 0.1.2

//! Plan-scope and output-schema identity carried by admitted capability
//! evidence.
//!
//! These facts answer "is the plan about to execute the plan that was
//! admitted?". They deliberately carry no bytes. Allocation limits belong to
//! the query domain's allocator, and a planning-time manifest carries none.

use std::collections::HashSet;

use crate::{
    error::{QueryError, QueryErrorCode, Result},
    plan::{DeferredGraphExecutionScopeIdentity, PlanFingerprint, QueryPlan, SchemaFingerprint},
};

/// One GPU UDTF occurrence in a DataFusion-owned execution graph. Every
/// occurrence materializes host Arrow for a DataFusion consumer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttemptExecutionOccurrence {
    /// Unique within one manifest; names one invocation site in the final
    /// DataFusion plan.
    occurrence_id: String,
    /// Identity of the physical producer, shared by every occurrence of one
    /// shared producer.
    producer_identity: String,
    execution_scope_fingerprint: PlanFingerprint,
    /// Present when the plan reads graph source facts that are filled in only
    /// after Arrow staging. The fingerprint then differs between admission
    /// and execution, and this identity is what matches.
    deferred_execution_scope_identity: Option<DeferredGraphExecutionScopeIdentity>,
    output_schema_fingerprint: SchemaFingerprint,
    /// Occurrence ids that execute before this one. `validate` admits only
    /// ids listed earlier in the manifest, which rules out cycles.
    predecessors: Vec<String>,
}

impl AttemptExecutionOccurrence {
    #[cfg(any(test, feature = "test-hooks"))]
    #[doc(hidden)]
    pub fn host_arrow(
        occurrence_id: impl Into<String>,
        plan: &QueryPlan,
        predecessors: Vec<String>,
        deferred_graph_facts: bool,
    ) -> Self {
        let occurrence_id = occurrence_id.into();
        Self::host_arrow_for_producer(occurrence_id.clone(), occurrence_id, plan, predecessors, deferred_graph_facts)
    }

    /// Creates one invocation occurrence for a producer execution identity.
    /// A shared physical producer may appear more than once in the final plan,
    /// so its stable producer identity is distinct from the occurrence id.
    #[doc(hidden)]
    pub fn host_arrow_for_producer(
        occurrence_id: impl Into<String>,
        producer_identity: impl Into<String>,
        plan: &QueryPlan,
        predecessors: Vec<String>,
        deferred_graph_facts: bool,
    ) -> Self {
        Self {
            occurrence_id: occurrence_id.into(),
            producer_identity: producer_identity.into(),
            execution_scope_fingerprint: plan.execution_fingerprint(),
            deferred_execution_scope_identity: deferred_graph_facts
                .then(|| DeferredGraphExecutionScopeIdentity::for_plan(plan)),
            output_schema_fingerprint: SchemaFingerprint::of(plan.output_schema()),
            predecessors,
        }
    }

    /// Whether `plan` is the admitted plan. The fingerprint matches a plan
    /// unchanged since admission; the deferred identity matches one whose
    /// masked graph source facts were filled in afterwards.
    pub(crate) fn matches_plan(&self, plan: &QueryPlan) -> bool {
        self.execution_scope_fingerprint == plan.execution_fingerprint()
            || self.deferred_execution_scope_identity.as_ref().is_some_and(|expected| expected.matches(plan))
    }

    pub(crate) fn matches_schema(&self, plan: &QueryPlan) -> bool {
        self.output_schema_fingerprint == SchemaFingerprint::of(plan.output_schema())
    }
}

/// Topologically ordered GPU occurrences for one DataFusion invocation.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct AttemptExecutionManifest {
    occurrences: Vec<AttemptExecutionOccurrence>,
}

impl AttemptExecutionManifest {
    #[doc(hidden)]
    pub fn new(occurrences: Vec<AttemptExecutionOccurrence>) -> Result<Self> {
        let manifest = Self { occurrences };
        manifest.validate()?;
        Ok(manifest)
    }

    #[cfg(any(test, feature = "test-hooks"))]
    #[doc(hidden)]
    pub fn occurrences(&self) -> &[AttemptExecutionOccurrence] {
        &self.occurrences
    }

    pub(crate) fn occurrence(&self, occurrence_id: &str) -> Option<&AttemptExecutionOccurrence> {
        self.occurrences.iter().find(|occurrence| occurrence.occurrence_id == occurrence_id)
    }

    #[doc(hidden)]
    pub fn occurrence_for_plan(&self, plan: &QueryPlan) -> Option<&AttemptExecutionOccurrence> {
        self.occurrences.iter().find(|occurrence| occurrence.matches_plan(plan) && occurrence.matches_schema(plan))
    }

    /// Validates the exact host-materializing occurrence about to execute and
    /// returns its already-validated predecessor identities.
    #[doc(hidden)]
    #[cfg(test)]
    pub fn validate_host_occurrence(&self, occurrence_id: &str, plan: &QueryPlan) -> Result<Vec<String>> {
        self.validate_host_occurrence_for_producer(occurrence_id, occurrence_id, plan)
    }

    /// Validates the concrete invocation occurrence and its stable producer
    /// identity, then returns its predecessor occurrence ids.
    #[doc(hidden)]
    pub fn validate_host_occurrence_for_producer(
        &self,
        occurrence_id: &str,
        producer_identity: &str,
        plan: &QueryPlan,
    ) -> Result<Vec<String>> {
        let occurrence = self
            .occurrence(occurrence_id)
            .ok_or_else(|| manifest_error("execution_occurrence_not_admitted", Some(occurrence_id)))?;
        if occurrence.producer_identity != producer_identity {
            return Err(manifest_error("execution_occurrence_producer_mismatch", Some(occurrence_id)));
        }
        if !occurrence.matches_plan(plan) {
            return Err(manifest_error("execution_occurrence_plan_mismatch", Some(occurrence_id)));
        }
        if !occurrence.matches_schema(plan) {
            return Err(manifest_error("execution_occurrence_schema_mismatch", Some(occurrence_id)));
        }
        Ok(occurrence.predecessors.clone())
    }

    /// Validates that a producer execution identity has at least one matching
    /// host occurrence in this invocation. Runtime claims the concrete unique
    /// occurrence when execution begins.
    #[doc(hidden)]
    pub fn validate_host_producer(&self, producer_identity: &str, plan: &QueryPlan) -> Result<()> {
        let candidates = self
            .occurrences
            .iter()
            .filter(|occurrence| occurrence.producer_identity == producer_identity)
            .collect::<Vec<_>>();
        // Checks go from coarse to fine so the `reason` fact names the first
        // identity that failed: the producer, then its plan, then its schema.
        if candidates.is_empty() {
            return Err(manifest_error("execution_producer_not_admitted", Some(producer_identity)));
        }
        if !candidates.iter().any(|occurrence| occurrence.matches_plan(plan)) {
            return Err(manifest_error("execution_occurrence_plan_mismatch", Some(producer_identity)));
        }
        if !candidates.iter().any(|occurrence| occurrence.matches_plan(plan) && occurrence.matches_schema(plan)) {
            return Err(manifest_error("execution_occurrence_schema_mismatch", Some(producer_identity)));
        }
        Ok(())
    }

    fn validate(&self) -> Result<()> {
        let mut seen = HashSet::with_capacity(self.occurrences.len());
        for occurrence in &self.occurrences {
            if occurrence.occurrence_id.is_empty() {
                return Err(manifest_error("missing_execution_occurrence_id", None));
            }
            if occurrence.producer_identity.is_empty() {
                return Err(manifest_error("missing_execution_producer_identity", Some(&occurrence.occurrence_id)));
            }
            if !seen.insert(occurrence.occurrence_id.as_str()) {
                return Err(manifest_error("duplicate_execution_occurrence_id", Some(&occurrence.occurrence_id)));
            }
            for predecessor in &occurrence.predecessors {
                if predecessor == &occurrence.occurrence_id || !seen.contains(predecessor.as_str()) {
                    return Err(manifest_error("invalid_execution_predecessor", Some(&occurrence.occurrence_id)));
                }
            }
        }
        Ok(())
    }
}

/// One producer occurrence bound to an admitted consumer scope.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttemptProducerOccurrence {
    occurrence_id: String,
    execution_scope_fingerprint: PlanFingerprint,
    deferred_execution_scope_identity: Option<DeferredGraphExecutionScopeIdentity>,
    output_schema_fingerprint: SchemaFingerprint,
}

impl AttemptProducerOccurrence {
    #[doc(hidden)]
    pub fn new(
        occurrence_id: impl Into<String>,
        execution_scope_fingerprint: PlanFingerprint,
        output_schema_fingerprint: SchemaFingerprint,
    ) -> Self {
        Self {
            occurrence_id: occurrence_id.into(),
            execution_scope_fingerprint,
            deferred_execution_scope_identity: None,
            output_schema_fingerprint,
        }
    }

    #[doc(hidden)]
    pub fn with_deferred_graph_identity(mut self, identity: Option<DeferredGraphExecutionScopeIdentity>) -> Self {
        self.deferred_execution_scope_identity = identity;
        self
    }

    #[doc(hidden)]
    pub fn occurrence_id(&self) -> &str {
        &self.occurrence_id
    }

    #[doc(hidden)]
    pub fn execution_scope_fingerprint(&self) -> PlanFingerprint {
        self.execution_scope_fingerprint
    }

    #[doc(hidden)]
    pub fn output_schema_fingerprint(&self) -> SchemaFingerprint {
        self.output_schema_fingerprint
    }

    #[doc(hidden)]
    pub fn matches_execution_plan(&self, plan: &QueryPlan) -> bool {
        self.execution_scope_fingerprint == plan.execution_fingerprint()
            || self.deferred_execution_scope_identity.as_ref().is_some_and(|identity| identity.matches(plan))
    }
}

/// Immutable scope and output-schema identity for one admitted plan fragment.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AttemptExecutionScope {
    execution_scope_fingerprint: PlanFingerprint,
    output_schema_fingerprint: SchemaFingerprint,
    /// Shared-resident producers whose output this consumer scope reads.
    /// Pushed while admitted evidence is assembled; execution later claims
    /// each one by its occurrence id, which `validate_manifest_core` requires
    /// to be unique.
    producer_occurrences: Vec<AttemptProducerOccurrence>,
}

impl AttemptExecutionScope {
    /// Scope of `plan`, keyed like its fingerprint by the caller-visible
    /// output: a host sink and a deferred host sink of the same schema share
    /// one scope.
    pub(crate) fn for_query_plan(plan: &QueryPlan) -> Self {
        Self {
            execution_scope_fingerprint: plan.execution_fingerprint(),
            output_schema_fingerprint: SchemaFingerprint::of(plan.caller_visible_output_schema()),
            producer_occurrences: Vec::new(),
        }
    }

    #[doc(hidden)]
    pub fn execution_scope_fingerprint(&self) -> PlanFingerprint {
        self.execution_scope_fingerprint
    }

    #[doc(hidden)]
    pub fn output_schema_fingerprint(&self) -> SchemaFingerprint {
        self.output_schema_fingerprint
    }

    #[doc(hidden)]
    pub fn producer_occurrences(&self) -> &[AttemptProducerOccurrence] {
        &self.producer_occurrences
    }

    pub(crate) fn push_producer_occurrence(&mut self, occurrence: AttemptProducerOccurrence) {
        self.producer_occurrences.push(occurrence);
    }

    #[doc(hidden)]
    pub fn matches(&self, plan: &QueryPlan) -> bool {
        self.execution_scope_fingerprint == plan.execution_fingerprint()
    }

    pub(crate) fn matches_execution_scope_fingerprint(&self, execution_scope_fingerprint: PlanFingerprint) -> bool {
        self.execution_scope_fingerprint == execution_scope_fingerprint
    }

    pub(crate) fn producer_occurrence_scope_count(&self, execution_scope_fingerprint: PlanFingerprint) -> usize {
        self.producer_occurrences
            .iter()
            .filter(|occurrence| occurrence.execution_scope_fingerprint == execution_scope_fingerprint)
            .count()
    }

    /// Rejects a manifest whose occurrence identity is incomplete or
    /// ambiguous. A duplicate or empty occurrence id would let one admitted
    /// scope stand in for another at execution time.
    #[doc(hidden)]
    pub fn validate_manifest_core(&self) -> Result<()> {
        let mut occurrence_ids = HashSet::with_capacity(self.producer_occurrences.len());
        for occurrence in &self.producer_occurrences {
            if occurrence.occurrence_id.is_empty() {
                return Err(manifest_error("missing_producer_occurrence_id", None));
            }
            if !occurrence_ids.insert(&occurrence.occurrence_id) {
                return Err(manifest_error("duplicate_producer_occurrence_id", Some(&occurrence.occurrence_id)));
            }
        }
        Ok(())
    }
}

fn manifest_error(reason: &'static str, occurrence_id: Option<&str>) -> QueryError {
    let error = QueryError::new(
        QueryErrorCode::InternalInvariantViolation,
        "query_engine.admission.execution_scope",
        "an admitted attempt lacks complete plan-scope identity",
    )
    .with_fact("reason", reason)
    .with_fact("validation_stage", "manifest");
    match occurrence_id {
        Some(occurrence_id) => error.with_fact("occurrence_id", occurrence_id),
        None => error,
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow::datatypes::{DataType, Field, Schema};

    use super::*;
    use crate::{
        plan::QueryPlanBuilder,
        source::{GpuResidentSource, SourceNode},
    };

    fn scope_with_occurrences(occurrences: Vec<AttemptProducerOccurrence>) -> AttemptExecutionScope {
        let mut scope = AttemptExecutionScope::for_query_plan(&resident_plan("consumer", None));
        scope.producer_occurrences = occurrences;
        scope
    }

    fn producer_occurrence(occurrence_id: &str, producer: &str) -> AttemptProducerOccurrence {
        let plan = resident_plan(producer, None);
        AttemptProducerOccurrence::new(
            occurrence_id,
            plan.execution_fingerprint(),
            SchemaFingerprint::of(plan.output_schema()),
        )
    }

    #[test]
    fn occurrence_manifest_rejects_duplicate_runtime_identity() {
        let scope = scope_with_occurrences(vec![
            producer_occurrence("left", "producer-a"),
            producer_occurrence("left", "producer-b"),
        ]);

        let error = scope.validate_manifest_core().expect_err("duplicate occurrence identity must be ambiguous");
        assert_eq!(error.fact_value("reason"), Some("duplicate_producer_occurrence_id"));
        assert_eq!(error.fact_value("occurrence_id"), Some("left"));
    }

    #[test]
    fn occurrence_manifest_rejects_missing_occurrence_id() {
        let error = scope_with_occurrences(vec![producer_occurrence("", "producer")])
            .validate_manifest_core()
            .expect_err("an occurrence without an id must reject");
        assert_eq!(error.fact_value("reason"), Some("missing_producer_occurrence_id"));
        assert_eq!(error.fact_value("validation_stage"), Some("manifest"));
    }

    fn resident_plan(name: &str, row_bound: Option<u64>) -> QueryPlan {
        let schema = Arc::new(Schema::new(vec![Field::new("vertex", DataType::Int64, false)]));
        let mut source = GpuResidentSource::new(name, schema);
        if let Some(row_bound) = row_bound {
            source = source.with_row_count_upper_bound(row_bound);
        }
        let mut builder = QueryPlanBuilder::new();
        let root = builder.add_source(SourceNode::GpuResident(source));
        builder.finish_gpu_resident(root).unwrap()
    }

    #[test]
    fn execution_manifest_rejects_duplicate_and_forward_predecessors() {
        let plan = resident_plan("edges", None);
        let duplicate = AttemptExecutionManifest::new(vec![
            AttemptExecutionOccurrence::host_arrow("gpu-0", &plan, vec![], true),
            AttemptExecutionOccurrence::host_arrow("gpu-0", &plan, vec![], true),
        ])
        .expect_err("duplicate occurrence ids must reject");
        assert_eq!(duplicate.fact_value("reason"), Some("duplicate_execution_occurrence_id"));

        let forward = AttemptExecutionManifest::new(vec![AttemptExecutionOccurrence::host_arrow(
            "gpu-1",
            &plan,
            vec!["gpu-0".to_string()],
            true,
        )])
        .expect_err("predecessors must already exist in topological order");
        assert_eq!(forward.fact_value("reason"), Some("invalid_execution_predecessor"));
    }

    #[test]
    fn host_occurrence_validation_requires_id_plan_schema_and_declared_refinement() {
        let admitted = resident_plan("edges", None);
        let manifest = AttemptExecutionManifest::new(vec![AttemptExecutionOccurrence::host_arrow(
            "gpu-0",
            &admitted,
            vec![],
            true,
        )])
        .unwrap();

        assert!(manifest.validate_host_occurrence("gpu-0", &resident_plan("edges", Some(12))).is_ok());
        let wrong_id = manifest.validate_host_occurrence("gpu-1", &admitted).unwrap_err();
        assert_eq!(wrong_id.fact_value("reason"), Some("execution_occurrence_not_admitted"));
        let wrong_plan = manifest.validate_host_occurrence("gpu-0", &resident_plan("other", Some(12))).unwrap_err();
        assert_eq!(wrong_plan.fact_value("reason"), Some("execution_occurrence_plan_mismatch"));
    }
}