uni_query/query/df_graph/

locy_explain.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2024-2026 Dragonscale Team

//! EXPLAIN RULE derivation tree construction.
//!
//! Ported from `uni-locy/src/orchestrator/explain.rs`. Uses `DerivedFactSource`
//! instead of `CypherExecutor`. Uses `RowStore` for row-based fact storage.
//!
//! Implements Mode A (provenance-based, uses DerivationTracker recorded during fixpoint)
//! with fallback to Mode B (re-execution) when tracker has no entries for the rule.

use std::collections::{HashMap, HashSet};
use std::sync::RwLock;

use uni_common::Value;
use uni_cypher::locy_ast::{ExplainRule, RuleCondition};
use uni_locy::types::CompiledRule;
use uni_locy::{CompiledProgram, DerivationNode, LocyConfig, LocyError, LocyStats, Row};

use super::locy_delta::{
    KeyTuple, RowStore, extract_cypher_conditions, extract_key, resolve_clause_with_is_refs,
};

use super::locy_eval::{eval_expr, record_batches_to_locy_rows};
use super::locy_slg::SLGResolver;
use super::locy_traits::DerivedFactSource;

/// Input dependency recorded for a derived fact: the source IS-ref rule and
/// the hash of the source fact.
#[derive(Clone)]
pub struct DerivationInput {
    pub is_ref_rule: String,
    pub fact_hash: Vec<u8>,
}

/// Provenance record for a single derived fact, recorded during fixpoint iteration.
#[derive(Clone)]
pub struct DerivationEntry {
    /// Name of the rule that derived this fact.
    pub rule_name: String,
    /// Index of the clause within the rule that produced this fact.
    pub clause_index: usize,
    /// Hashes of IS-ref input facts (populated when IS-ref tracking is available).
    pub inputs: Vec<DerivationInput>,
    /// ALONG column values captured at derivation time.
    pub along_values: HashMap<String, Value>,
    /// Fixpoint iteration number when the fact was first derived.
    pub iteration: usize,
    /// Full fact row stored for Mode A filtering/display.
    pub fact_row: Row,
}

/// Tracks provenance of derived facts recorded during fixpoint iteration.
///
/// Enables Mode A (provenance-based) EXPLAIN without re-execution.
/// First-derivation-wins semantics: once a fact hash is recorded, subsequent
/// iterations do not overwrite it.
pub struct DerivationTracker {
    entries: RwLock<HashMap<Vec<u8>, DerivationEntry>>,
}

impl DerivationTracker {
    pub fn new() -> Self {
        Self {
            entries: RwLock::new(HashMap::new()),
        }
    }

    /// Record a derivation entry. First-derivation-wins: if the hash is already
    /// present, the existing entry is kept.
    pub fn record(&self, fact_hash: Vec<u8>, entry: DerivationEntry) {
        if let Ok(mut guard) = self.entries.write() {
            guard.entry(fact_hash).or_insert(entry);
        }
    }

    /// Look up the derivation entry for a fact hash.
    pub fn lookup(&self, fact_hash: &[u8]) -> Option<DerivationEntry> {
        self.entries.read().ok()?.get(fact_hash).cloned()
    }

    /// Get all entries for a specific rule name.
    pub fn entries_for_rule(&self, rule_name: &str) -> Vec<(Vec<u8>, DerivationEntry)> {
        match self.entries.read() {
            Ok(guard) => guard
                .iter()
                .filter(|(_, e)| e.rule_name == rule_name)
                .map(|(k, v)| (k.clone(), v.clone()))
                .collect(),
            Err(_) => vec![],
        }
    }
}

impl Default for DerivationTracker {
    fn default() -> Self {
        Self::new()
    }
}

/// Set of (rule_name, key_tuple) to detect cycles during recursive derivation (Mode B).
type VisitedSet = HashSet<(String, KeyTuple)>;

/// Build a derivation tree for a rule, showing how each fact was derived.
///
/// Tries Mode A (provenance-based, uses DerivationTracker) first when a tracker is
/// provided and has entries for the rule.  Falls through to Mode B (re-execution)
/// when Mode A cannot produce a result.
pub async fn explain_rule(
    query: &ExplainRule,
    program: &CompiledProgram,
    fact_source: &dyn DerivedFactSource,
    config: &LocyConfig,
    derived_store: &mut RowStore,
    stats: &mut LocyStats,
    tracker: Option<&DerivationTracker>,
) -> Result<DerivationNode, LocyError> {
    // Mode A: provenance-based (no re-execution required).
    // Falls through to Mode B when tracker is absent or has no matching entries.
    if let Some(Ok(node)) = tracker.map(|t| explain_rule_mode_a(query, program, t, derived_store)) {
        return Ok(node);
    }

    // Mode B: re-execution fallback
    explain_rule_mode_b(query, program, fact_source, config, derived_store, stats).await
}

/// Mode A: build derivation tree using recorded provenance from the fixpoint loop.
///
/// Returns `Err` when no tracker entries exist for the rule (signals Mode B fallback).
fn explain_rule_mode_a(
    query: &ExplainRule,
    program: &CompiledProgram,
    tracker: &DerivationTracker,
    derived_store: &RowStore,
) -> Result<DerivationNode, LocyError> {
    let rule_name = query.rule_name.to_string();
    let rule = program
        .rule_catalog
        .get(&rule_name)
        .ok_or_else(|| LocyError::EvaluationError {
            message: format!("rule '{}' not found for EXPLAIN RULE (Mode A)", rule_name),
        })?;

    let tracker_entries = tracker.entries_for_rule(&rule_name);
    if tracker_entries.is_empty() {
        return Err(LocyError::EvaluationError {
            message: format!("no tracker entries for rule '{rule_name}' (falling back to Mode B)"),
        });
    }

    // Filter tracker entries by WHERE expression
    let matching_entries: Vec<_> = tracker_entries
        .into_iter()
        .filter(|(_, entry)| {
            eval_expr(&query.where_expr, &entry.fact_row)
                .map(|v| v.as_bool().unwrap_or(false))
                .unwrap_or(false)
        })
        .collect();

    if matching_entries.is_empty() {
        return Err(LocyError::EvaluationError {
            message: format!("no tracker entries match WHERE clause for rule '{rule_name}'"),
        });
    }

    // Also check orchestrator store for any additional facts not in tracker
    // (e.g., facts derived outside the native fixpoint path)
    let orch_facts = derived_store
        .get(&rule_name)
        .map(|r| r.rows.clone())
        .unwrap_or_default();
    let _ = orch_facts; // available for future use

    let mut root = DerivationNode {
        rule: rule_name.clone(),
        clause_index: 0,
        priority: rule.priority,
        bindings: HashMap::new(),
        along_values: HashMap::new(),
        children: Vec::new(),
        graph_fact: None,
    };

    for (_, entry) in matching_entries {
        let along_values = extract_along_values(&entry.fact_row, rule);
        let clause_priority = rule
            .clauses
            .get(entry.clause_index)
            .and_then(|c| c.priority);
        let node = DerivationNode {
            rule: rule_name.clone(),
            clause_index: entry.clause_index,
            priority: clause_priority.or(rule.priority),
            bindings: entry.fact_row.clone(),
            along_values,
            // Mode A: children not tracked (inputs list is reserved for future recursion)
            children: vec![],
            graph_fact: Some(format!(
                "[iter={}] {}",
                entry.iteration,
                format_graph_fact(&entry.fact_row)
            )),
        };
        root.children.push(node);
    }

    Ok(root)
}

/// Mode B: re-execution fallback — re-executes clause queries to find which
/// clause produced each matching fact, then recurses into IS references.
async fn explain_rule_mode_b(
    query: &ExplainRule,
    program: &CompiledProgram,
    fact_source: &dyn DerivedFactSource,
    config: &LocyConfig,
    derived_store: &mut RowStore,
    stats: &mut LocyStats,
) -> Result<DerivationNode, LocyError> {
    let rule_name = query.rule_name.to_string();
    let rule = program
        .rule_catalog
        .get(&rule_name)
        .ok_or_else(|| LocyError::EvaluationError {
            message: format!("rule '{}' not found for EXPLAIN RULE", rule_name),
        })?;

    let key_columns: Vec<String> = rule
        .yield_schema
        .iter()
        .filter(|c| c.is_key)
        .map(|c| c.name.clone())
        .collect();

    // Re-evaluate the rule via SLG to obtain rows with full node objects and properties.
    // The native fixpoint's orch_store has VID-only integers that fail property-based
    // WHERE filters (e.g. a.name = 'A') — we need actual Value::Node values here.
    {
        let mut fresh_store = RowStore::new();
        let slg_start = std::time::Instant::now();
        let mut resolver =
            SLGResolver::new(program, fact_source, config, &mut fresh_store, slg_start);
        resolver.resolve_goal(&rule_name, &HashMap::new()).await?;
        stats.queries_executed += resolver.stats.queries_executed;
        // Merge full-node facts into derived_store so IS-ref lookups in
        // build_derivation_node also get proper node objects (not VIDs).
        for (name, relation) in fresh_store {
            derived_store.insert(name, relation);
        }
    }

    // Get all derived facts for this rule (now populated with full node objects)
    let facts = derived_store
        .get(&rule_name)
        .map(|r| r.rows.clone())
        .unwrap_or_default();

    // Filter facts by WHERE expression
    let filtered: Vec<Row> = facts
        .into_iter()
        .filter(|row| {
            eval_expr(&query.where_expr, row)
                .map(|v| v.as_bool().unwrap_or(false))
                .unwrap_or(false)
        })
        .collect();

    // Build derivation tree root
    let mut root = DerivationNode {
        rule: rule_name.clone(),
        clause_index: 0,
        priority: rule.priority,
        bindings: HashMap::new(),
        along_values: HashMap::new(),
        children: Vec::new(),
        graph_fact: None,
    };

    // For each matching fact, recursively build a derivation node
    for fact in &filtered {
        let mut visited = VisitedSet::new();
        let node = build_derivation_node(
            &rule_name,
            fact,
            &key_columns,
            program,
            fact_source,
            derived_store,
            stats,
            &mut visited,
            config.max_explain_depth,
        )
        .await?;
        root.children.push(node);
    }

    Ok(root)
}

/// Recursively build a derivation node for a single fact of a rule.
///
/// Finds which clause produced this fact, extracts ALONG values,
/// and recurses into IS reference dependencies.
#[allow(clippy::too_many_arguments)]
fn build_derivation_node<'a>(
    rule_name: &'a str,
    fact: &'a Row,
    key_columns: &'a [String],
    program: &'a CompiledProgram,
    fact_source: &'a dyn DerivedFactSource,
    derived_store: &'a mut RowStore,
    stats: &'a mut LocyStats,
    visited: &'a mut VisitedSet,
    max_depth: usize,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<DerivationNode, LocyError>> + 'a>> {
    Box::pin(async move {
        let rule =
            program
                .rule_catalog
                .get(rule_name)
                .ok_or_else(|| LocyError::EvaluationError {
                    message: format!("rule '{}' not found during EXPLAIN", rule_name),
                })?;

        let key_tuple = extract_key(fact, key_columns);
        let visit_key = (rule_name.to_string(), key_tuple);

        // Cycle detection
        if !visited.insert(visit_key.clone()) || max_depth == 0 {
            return Ok(DerivationNode {
                rule: rule_name.to_string(),
                clause_index: 0,
                priority: rule.priority,
                bindings: fact.clone(),
                along_values: extract_along_values(fact, rule),
                children: Vec::new(),
                graph_fact: Some("(cycle)".to_string()),
            });
        }

        // All yield columns (key + non-key) for exact row matching
        let yield_columns: Vec<String> = rule.yield_schema.iter().map(|c| c.name.clone()).collect();

        // Try each clause to find the one that produced this fact
        for (clause_idx, clause) in rule.clauses.iter().enumerate() {
            let has_is_refs = clause
                .where_conditions
                .iter()
                .any(|c| matches!(c, RuleCondition::IsReference(_)));
            let has_along = !clause.along.is_empty();

            let resolved = if has_is_refs || has_along {
                let rows = resolve_clause_with_is_refs(clause, fact_source, derived_store).await?;
                stats.queries_executed += 1;
                rows
            } else {
                let cypher_conditions = extract_cypher_conditions(&clause.where_conditions);
                let raw_batches = fact_source
                    .execute_pattern(&clause.match_pattern, &cypher_conditions)
                    .await?;
                stats.queries_executed += 1;
                record_batches_to_locy_rows(&raw_batches)
            };

            // Match on ALL yield columns (not just key columns) to find exact row
            let matching_row = resolved
                .iter()
                .find(|row| yield_columns.iter().all(|k| row.get(k) == fact.get(k)));

            if let Some(evidence_row) = matching_row {
                let along_values = extract_along_values(fact, rule);

                // Build children by recursing into IS references
                let mut children = Vec::new();
                for cond in &clause.where_conditions {
                    if let RuleCondition::IsReference(is_ref) = cond {
                        if is_ref.negated {
                            continue;
                        }
                        let ref_rule_name = is_ref.rule_name.to_string();
                        if let Some(ref_rule) = program.rule_catalog.get(&ref_rule_name) {
                            let ref_key_columns: Vec<String> = ref_rule
                                .yield_schema
                                .iter()
                                .filter(|c| c.is_key)
                                .map(|c| c.name.clone())
                                .collect();

                            let ref_facts: Vec<Row> = derived_store
                                .get(&ref_rule_name)
                                .map(|r| r.rows.clone())
                                .unwrap_or_default();

                            let matching_ref_facts: Vec<Row> = ref_facts
                                .into_iter()
                                .filter(|ref_fact| {
                                    let subjects_match =
                                        is_ref.subjects.iter().enumerate().all(|(i, subject)| {
                                            if i < ref_key_columns.len() {
                                                let subject_val = evidence_row
                                                    .get(subject)
                                                    .or_else(|| fact.get(subject));
                                                match subject_val {
                                                    Some(val) => {
                                                        ref_fact.get(&ref_key_columns[i])
                                                            == Some(val)
                                                    }
                                                    None => true,
                                                }
                                            } else {
                                                true
                                            }
                                        });
                                    let target_matches = if let Some(target) = &is_ref.target {
                                        let target_idx = is_ref.subjects.len();
                                        if target_idx < ref_key_columns.len() {
                                            let target_val = evidence_row
                                                .get(target)
                                                .or_else(|| fact.get(target));
                                            match target_val {
                                                Some(val) => {
                                                    ref_fact.get(&ref_key_columns[target_idx])
                                                        == Some(val)
                                                }
                                                None => true,
                                            }
                                        } else {
                                            true
                                        }
                                    } else {
                                        true
                                    };
                                    subjects_match && target_matches
                                })
                                .collect();

                            for ref_fact in matching_ref_facts {
                                let child = build_derivation_node(
                                    &ref_rule_name,
                                    &ref_fact,
                                    &ref_key_columns,
                                    program,
                                    fact_source,
                                    derived_store,
                                    stats,
                                    visited,
                                    max_depth - 1,
                                )
                                .await?;
                                children.push(child);
                            }
                        }
                    }
                }

                // Backtrack visited set
                visited.remove(&visit_key);

                let mut merged_bindings = evidence_row.clone();
                for (k, v) in fact {
                    merged_bindings.entry(k.clone()).or_insert(v.clone());
                }

                return Ok(DerivationNode {
                    rule: rule_name.to_string(),
                    clause_index: clause_idx,
                    priority: rule.clauses[clause_idx].priority,
                    bindings: merged_bindings,
                    along_values,
                    children,
                    graph_fact: Some(format_graph_fact(evidence_row)),
                });
            }
        }

        // No clause matched — leaf node
        visited.remove(&visit_key);
        Ok(DerivationNode {
            rule: rule_name.to_string(),
            clause_index: 0,
            priority: rule.priority,
            bindings: fact.clone(),
            along_values: extract_along_values(fact, rule),
            children: Vec::new(),
            graph_fact: Some(format_graph_fact(fact)),
        })
    })
}

fn extract_along_values(fact: &Row, rule: &CompiledRule) -> HashMap<String, Value> {
    let mut along_values = HashMap::new();
    for clause in &rule.clauses {
        for along in &clause.along {
            if let Some(v) = fact.get(&along.name) {
                along_values.insert(along.name.clone(), v.clone());
            }
        }
    }
    along_values
}

pub(crate) fn format_graph_fact(row: &Row) -> String {
    let mut entries: Vec<String> = row
        .iter()
        .map(|(k, v)| format!("{}: {}", k, format_value(v)))
        .collect();
    entries.sort();
    format!("{{{}}}", entries.join(", "))
}

fn format_value(v: &Value) -> String {
    match v {
        Value::Null => "null".to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Int(i) => i.to_string(),
        Value::Float(f) => f.to_string(),
        Value::String(s) => format!("\"{}\"", s),
        Value::List(items) => {
            let inner: Vec<String> = items.iter().map(format_value).collect();
            format!("[{}]", inner.join(", "))
        }
        Value::Map(m) => {
            let mut entries: Vec<String> = m
                .iter()
                .map(|(k, v)| format!("{}: {}", k, format_value(v)))
                .collect();
            entries.sort();
            format!("{{{}}}", entries.join(", "))
        }
        Value::Node(n) => format!("Node({})", n.vid.as_u64()),
        Value::Edge(e) => format!("Edge({})", e.eid.as_u64()),
        _ => format!("{:?}", v),
    }
}