linkrs-query 0.1.0

use std::sync::Arc;

use parking_lot::RwLock;

use crate::executor::streaming::chunk::DataChunk;
use crate::executor::streaming::executor::StreamingExecutor;
use crate::executor::streaming::operators::source_operator::OperatorConfig;
use crate::executor::streaming::query_registry::CancelToken;
use crate::executor::streaming::slot::SlotLayout;
use crate::parser::ast::pattern::PathSemantic;
use crate::storage::QueryStorage;
use linkrs_core::error::QueryError;
use linkrs_core::types::expr::Expression;
use linkrs_core::{EdgeDirection, Value};

use super::super::runtime::ExecutionRuntime;
use super::spec::GraphSpec;
use super::visited_set::VisitedSet;

pub(crate) mod expand;
mod expand_buffer;
mod expand_columnar;
mod expand_dispatch;
mod expand_frontier;
mod expand_row_path;
mod expand_seeds;
mod subgraph;
mod traverse;

pub(super) struct ExpandCtx<'a> {
    pub(super) space_name: &'a str,
    pub(super) dst_tag: &'a str,
    pub(super) edge_types: &'a [String],
    pub(super) direction: EdgeDirection,
    pub(super) filter_expr: &'a Option<Expression>,
    pub(super) col_names_template: Vec<String>,
    pub(super) cancel_token: Option<CancelToken>,
    pub(super) path_semantic: Option<PathSemantic>,
    pub(super) edge_required_props: Option<Vec<String>>,
    pub(super) dst_required_props: Option<Vec<String>>,
    pub(super) closed_loop: bool,
    pub(super) skip_rows: bool,
    pub(super) emit_raw_ids: bool,
    pub(super) lightweight_source: bool,
    pub(super) include_empty_paths: bool,
    /// Lazily computed TTL-activeness over the operator's edge types. A plan
    /// fixes the edge-type set, so the flag is constant for the operator and
    /// is computed at most once per `ExpandCtx`.
    pub(super) ttl_active: Option<bool>,
}

impl ExpandCtx<'_> {
    /// Whether any edge type this expand can reach has TTL configured.
    ///
    /// Topology fast paths read raw CSR adjacency and cannot evaluate TTL;
    /// they check this first and degrade to the materializing path when it is
    /// true, so results honor the "expired means absent" contract. The result
    /// is cached per `ExpandCtx` because the edge-type set is a plan-time
    /// constant.
    pub(super) fn edge_ttl_active(
        &mut self,
        reader: &dyn QueryStorage,
    ) -> Result<bool, QueryError> {
        if let Some(active) = self.ttl_active {
            return Ok(active);
        }
        let active = reader.edge_types_ttl_active(self.space_name, self.edge_types)?;
        self.ttl_active = Some(active);
        Ok(active)
    }
}

#[derive(Debug)]
pub enum GraphOperatorKind {
    Expand {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
        filter_expr: Option<Expression>,
    },
    ExpandAll {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
        filter_expr: Option<Expression>,
        col_names: Vec<String>,
        src_vids: Vec<Value>,
        step_limit: u32,
        step_limits: Option<Vec<u32>>,
        count_only: bool,
        emit_raw_ids: bool,
        lightweight_source: bool,
        path_semantic: Option<PathSemantic>,
        edge_required_props: Option<Vec<String>>,
        dst_required_props: Option<Vec<String>>,
        closed_loop: bool,
        skip_rows: bool,
        include_empty_paths: bool,
    },
    Traverse {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
        min_depth: u32,
        max_depth: u32,
        filter_expr: Option<Expression>,
        visited: VisitedSet,
        path_semantic: Option<PathSemantic>,
    },
    TraverseAll {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
        min_depth: u32,
        max_depth: u32,
        filter_expr: Option<Expression>,
        visited: VisitedSet,
        path_semantic: Option<PathSemantic>,
    },
    BiExpand {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
    },
    BiTraverse {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        edge_types: Vec<String>,
        direction: EdgeDirection,
        min_depth: u32,
        max_depth: u32,
        visited: VisitedSet,
    },
    Subgraph {
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        dst_tag: String,
        steps: u32,
        direction: EdgeDirection,
        edge_types: Vec<String>,
    },
}

/// Graph operator.
///
/// Wraps [`GraphOperatorKind`] with the runtime context injected at `open()`.
/// Lifecycle state is owned exclusively by the executor; operators never
/// write it.
#[derive(Debug)]
pub struct GraphOperator {
    pub kind: GraphOperatorKind,
    pub runtime: Option<Arc<ExecutionRuntime>>,
    pub output_layout: Arc<SlotLayout>,
    pub config: OperatorConfig,
}

impl GraphOperator {
    pub fn bind_runtime(&mut self, runtime: &Arc<ExecutionRuntime>) {
        let storage = runtime.storage.clone();
        let space_name = runtime.query_id().space_name.unwrap_or_default();
        self.runtime = Some(Arc::clone(runtime));
        match &mut self.kind {
            GraphOperatorKind::Expand {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::ExpandAll {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::Traverse {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::TraverseAll {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::BiExpand {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::BiTraverse {
                storage: target_storage,
                space_name: target_space,
                ..
            }
            | GraphOperatorKind::Subgraph {
                storage: target_storage,
                space_name: target_space,
                ..
            } => {
                *target_storage = storage;
                *target_space = space_name;
            }
        }
    }

    pub fn from_spec(
        spec: &GraphSpec,
        storage: Option<Arc<RwLock<dyn QueryStorage>>>,
        space_name: String,
        output_layout: Arc<SlotLayout>,
    ) -> Self {
        let kind = match spec {
            GraphSpec::Expand {
                edge_types,
                direction,
                filter_expr,
                dst_tag,
                ..
            } => GraphOperatorKind::Expand {
                storage: storage.clone(),
                space_name: space_name.clone(),
                dst_tag: dst_tag.clone(),
                edge_types: edge_types.clone(),
                direction: *direction,
                filter_expr: filter_expr.clone(),
            },
            GraphSpec::ExpandAll {
                edge_types,
                direction,
                filter_expr,
                col_names,
                src_vids,
                step_limit,
                step_limits,
                count_only,
                emit_raw_ids,
                lightweight_source,
                path_semantic,
                dst_tag,
                edge_required_props,
                dst_required_props,
                closed_loop,
                skip_rows,
                include_empty_paths,
            } => GraphOperatorKind::ExpandAll {
                storage: storage.clone(),
                space_name: space_name.clone(),
                dst_tag: dst_tag.clone(),
                edge_types: edge_types.clone(),
                direction: *direction,
                filter_expr: filter_expr.clone(),
                col_names: col_names.clone(),
                src_vids: src_vids.clone(),
                step_limit: *step_limit,
                step_limits: step_limits.clone(),
                count_only: *count_only,
                emit_raw_ids: *emit_raw_ids,
                lightweight_source: *lightweight_source,
                path_semantic: path_semantic.clone(),
                edge_required_props: edge_required_props.clone(),
                dst_required_props: dst_required_props.clone(),
                closed_loop: *closed_loop,
                skip_rows: *skip_rows,
                include_empty_paths: *include_empty_paths,
            },
            GraphSpec::Traverse {
                edge_types,
                direction,
                min_depth,
                max_depth,
                filter_expr,
                path_semantic,
                dst_tag,
            } => GraphOperatorKind::Traverse {
                storage: storage.clone(),
                space_name: space_name.clone(),
                dst_tag: dst_tag.clone(),
                edge_types: edge_types.clone(),
                direction: *direction,
                min_depth: *min_depth,
                max_depth: *max_depth,
                filter_expr: filter_expr.clone(),
                visited: VisitedSet::new(),
                path_semantic: path_semantic.clone(),
            },
            GraphSpec::BiExpand {
                edge_types,
                direction,
                dst_tag,
            } => GraphOperatorKind::BiExpand {
                storage: storage.clone(),
                space_name: space_name.clone(),
                dst_tag: dst_tag.clone(),
                edge_types: edge_types.clone(),
                direction: *direction,
            },
            GraphSpec::BiTraverse {
                edge_types,
                direction,
                min_depth,
                max_depth,
                dst_tag,
            } => GraphOperatorKind::BiTraverse {
                storage: storage.clone(),
                space_name: space_name.clone(),
                dst_tag: dst_tag.clone(),
                edge_types: edge_types.clone(),
                direction: *direction,
                min_depth: *min_depth,
                max_depth: *max_depth,
                visited: VisitedSet::new(),
            },
        };
        Self::new(kind, output_layout)
    }

    pub fn new(kind: GraphOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
        Self {
            kind,
            runtime: None,
            output_layout,
            config: OperatorConfig::default(),
        }
    }

    /// Inject the runtime and execution config (called once by the executor
    /// before this operator produces any data).
    pub fn inject_context(
        &mut self,
        runtime: Option<&Arc<ExecutionRuntime>>,
        config: OperatorConfig,
    ) {
        if let Some(rt) = runtime {
            self.runtime = Some(rt.clone());
        }
        self.config = config;
    }

    pub fn open(&mut self, input: &mut StreamingExecutor) -> Result<(), QueryError> {
        match &mut self.kind {
            GraphOperatorKind::Expand { .. }
            | GraphOperatorKind::ExpandAll { .. }
            | GraphOperatorKind::Traverse { .. }
            | GraphOperatorKind::TraverseAll { .. }
            | GraphOperatorKind::BiExpand { .. }
            | GraphOperatorKind::BiTraverse { .. }
            | GraphOperatorKind::Subgraph { .. } => {
                input.open()?;
                Ok(())
            }
        }
    }

    pub fn next(&mut self, input: &mut StreamingExecutor) -> Result<Option<DataChunk>, QueryError> {
        let op: &mut GraphOperator = self;
        if matches!(&op.kind, GraphOperatorKind::Expand { .. }) {
            return expand::handle(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::ExpandAll { .. }) {
            return expand::handle_all(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::Traverse { .. }) {
            return traverse::handle_traverse(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::TraverseAll { .. }) {
            return traverse::handle_traverse_all(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::BiExpand { .. }) {
            return traverse::handle_bi_expand(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::BiTraverse { .. }) {
            return traverse::handle_bi_traverse(op, input);
        }
        if matches!(&op.kind, GraphOperatorKind::Subgraph { .. }) {
            return subgraph::handle(op, input);
        }
        unreachable!("graph_operator::next called for an unknown kind")
    }

    pub fn stop(&mut self) -> Result<(), QueryError> {
        Ok(())
    }

    pub fn close(&mut self) -> Result<(), QueryError> {
        Ok(())
    }

    /// Reset per-run graph state (visited sets) and rewind the input so the
    /// graph operator re-produces the same traversal.
    pub fn reset(&mut self, input: &mut StreamingExecutor) -> Result<bool, QueryError> {
        match &mut self.kind {
            GraphOperatorKind::Traverse { visited, .. }
            | GraphOperatorKind::TraverseAll { visited, .. }
            | GraphOperatorKind::BiTraverse { visited, .. } => {
                *visited = VisitedSet::new();
            }
            GraphOperatorKind::Expand { .. }
            | GraphOperatorKind::ExpandAll { .. }
            | GraphOperatorKind::BiExpand { .. }
            | GraphOperatorKind::Subgraph { .. } => {}
        }
        input.reset()?;
        Ok(false)
    }
}