hypercore-kb / opencog-kb / basal-atomspace / rust-atomspace / src / query.rs

//! Query engine for pattern matching
//!
//! Supports streaming queries that never load everything into RAM.

use crate::atom::{Atom, Hash};
use crate::storage::AtomStore;

/// A query pattern - can have wildcards (None) or specific constraints
#[derive(Debug, Clone)]
pub struct Pattern {
    /// Required type (or None for any)
    pub typ: Option<String>,

    /// For nodes: required name (or None for any)
    pub name: Option<String>,

    /// For links: outgoing constraints
    /// Each element is either:
    /// - Some(hash) = exact atom required
    /// - None = any atom (wildcard)
    pub outgoing: Option<Vec<Option<Hash>>>,
}

impl Pattern {
    /// Match any atom of a type
    pub fn by_type(typ: impl Into<String>) -> Self {
        Self {
            typ: Some(typ.into()),
            name: None,
            outgoing: None,
        }
    }

    /// Match a specific node
    pub fn node(typ: impl Into<String>, name: impl Into<String>) -> Self {
        Self {
            typ: Some(typ.into()),
            name: Some(name.into()),
            outgoing: None,
        }
    }

    /// Match links with specific outgoing pattern
    pub fn link(typ: impl Into<String>, outgoing: Vec<Option<Hash>>) -> Self {
        Self {
            typ: Some(typ.into()),
            name: None,
            outgoing: Some(outgoing),
        }
    }

    /// Check if an atom matches this pattern
    pub fn matches(&self, atom: &Atom) -> bool {
        // Check type
        if let Some(ref t) = self.typ {
            if atom.typ() != t {
                return false;
            }
        }

        // Check name (for nodes)
        if let Some(ref n) = self.name {
            match atom.name() {
                Some(name) if name == n => {}
                _ => return false,
            }
        }

        // Check outgoing (for links)
        if let Some(ref pattern_out) = self.outgoing {
            match atom.outgoing() {
                Some(atom_out) => {
                    if atom_out.len() != pattern_out.len() {
                        return false;
                    }
                    for (p, a) in pattern_out.iter().zip(atom_out.iter()) {
                        if let Some(required) = p {
                            if required != a {
                                return false;
                            }
                        }
                    }
                }
                None => return false,
            }
        }

        true
    }
}

/// Query result - streaming iterator
pub struct QueryResult<'a> {
    inner: Box<dyn Iterator<Item = (Hash, Atom)> + 'a>,
}

impl<'a> Iterator for QueryResult<'a> {
    type Item = (Hash, Atom);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next()
    }
}

/// Execute a query against a store
pub fn query<'a, S: AtomStore + 'a>(store: &'a S, pattern: &Pattern) -> QueryResult<'a> {
    // Optimization: if we know the type, use the type index
    if let Some(ref typ) = pattern.typ {
        let pattern = pattern.clone();
        let iter = store
            .iter_by_type(typ)
            .filter(move |(_, atom)| pattern.matches(atom));

        return QueryResult {
            inner: Box::new(iter),
        };
    }

    // Fallback: scan everything (expensive!)
    // In practice, we'd have more indexes
    unimplemented!("Full scan not implemented - always specify a type")
}

/// Find atoms connected to a given atom
pub fn get_neighbors<S: AtomStore>(store: &S, atom_hash: &Hash, depth: usize) -> Vec<(Hash, Atom)> {
    let mut visited = std::collections::HashSet::new();
    let mut result = Vec::new();
    let mut frontier = vec![*atom_hash];

    for _ in 0..depth {
        let mut next_frontier = Vec::new();

        for hash in frontier {
            if visited.contains(&hash) {
                continue;
            }
            visited.insert(hash);

            if let Some(atom) = store.get(&hash) {
                result.push((hash, atom.clone()));

                // Follow outgoing links
                if let Some(outgoing) = atom.outgoing() {
                    for out_hash in outgoing {
                        if !visited.contains(out_hash) {
                            next_frontier.push(*out_hash);
                        }
                    }
                }

                // Follow incoming links
                for inc_hash in store.get_incoming(&hash) {
                    if !visited.contains(&inc_hash) {
                        next_frontier.push(inc_hash);
                    }
                }
            }
        }

        frontier = next_frontier;
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::storage::MemoryStore;

    #[test]
    fn test_pattern_matching() {
        let store = MemoryStore::new();

        let dog = Atom::node("ConceptNode", "dog");
        let cat = Atom::node("ConceptNode", "cat");
        let animal = Atom::node("ConceptNode", "animal");

        let dog_hash = store.add(&dog);
        let cat_hash = store.add(&cat);
        let animal_hash = store.add(&animal);

        store.add(&Atom::link("InheritanceLink", vec![dog_hash, animal_hash]));
        store.add(&Atom::link("InheritanceLink", vec![cat_hash, animal_hash]));

        // Query: what inherits from animal?
        let pattern = Pattern::link("InheritanceLink", vec![None, Some(animal_hash)]);
        let results: Vec<_> = query(&store, &pattern).collect();

        assert_eq!(results.len(), 2);
    }
}