rete_core/

dictionary.rs

//! Dictionary assembly: classify RDF terms by role and assign the HDT-style
//! global IDs that triple blocks reference (SPEC.md §5, §5.1).
//!
//! Four sections are built from the §5.1 front-coded primitive:
//! `shared` (used as both subject and object), `subjects` (subject-only),
//! `objects` (object-only), and `predicates`. Subject IDs run `1..=S` (shared)
//! then `S+1..` (subject-only); object IDs reuse `1..=S` for the same shared
//! terms then `S+1..` for object-only. Predicates have their own space.

use std::collections::HashSet;

use crate::chunk_cache::ChunkCache;
use crate::dict::{ChunkedSection, DictSectionBuilder};
use crate::terms::{NodeId, ObjectId, PredicateId, SubjectId};

/// Sort a term slice ascending — in parallel when the `parallel` feature is on.
/// The front-coded dictionary sections need ascending order; doing it once here
/// (rather than per-insert in a `BTreeSet`) is the fast path for large graphs.
#[cfg(feature = "parallel")]
fn sort_terms(v: &mut [String]) {
    use rayon::slice::ParallelSliceMut;
    v.par_sort_unstable();
}
#[cfg(not(feature = "parallel"))]
fn sort_terms(v: &mut [String]) {
    v.sort_unstable();
}

/// Builds a [`Dictionary`] from observed `(subject, predicate, object)` terms.
///
/// Terms are deduped with `HashSet` during ingest (O(1) average insert, no tree
/// rebalancing) and sorted **once** in [`Self::build`]. This replaces a trio of
/// `BTreeSet<String>` whose per-insert `O(log n)` cost + a `String` allocation
/// for *every* observed term (most of them duplicates) dominated the build on
/// large graphs.
#[derive(Default)]
pub struct DictionaryBuilder {
    subjects: HashSet<String>,
    objects: HashSet<String>,
    predicates: HashSet<String>,
}

impl DictionaryBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    /// Record one triple's terms (IDs are assigned later, in [`Self::build`]).
    /// We probe with the borrowed `&str` first and only allocate a `String` for
    /// a genuinely new term — so a graph of N triples does far fewer than 3N
    /// allocations (each subject/predicate recurs across many triples).
    pub fn observe(&mut self, subject: &str, predicate: &str, object: &str) {
        if !self.subjects.contains(subject) {
            self.subjects.insert(subject.to_string());
        }
        if !self.objects.contains(object) {
            self.objects.insert(object.to_string());
        }
        if !self.predicates.contains(predicate) {
            self.predicates.insert(predicate.to_string());
        }
    }

    pub fn build(self) -> Dictionary {
        // Sort each role's terms once (ascending) — identical output order to
        // the old BTreeSet iteration, but without the per-insert tree cost.
        let mut subjects: Vec<String> = self.subjects.into_iter().collect();
        let mut predicates: Vec<String> = self.predicates.into_iter().collect();
        let object_set = self.objects; // kept for the shared (subject ∩ object) test
        sort_terms(&mut subjects);
        sort_terms(&mut predicates);

        // A term is "shared" when it appears as both a subject and an object.
        let subject_set: HashSet<&str> = subjects.iter().map(String::as_str).collect();

        let mut shared_b = DictSectionBuilder::new();
        let mut subj_b = DictSectionBuilder::new();
        let mut obj_b = DictSectionBuilder::new();
        let mut pred_b = DictSectionBuilder::new();

        // Subjects ascending: shared (also an object) vs subject-only.
        for s in &subjects {
            if object_set.contains(s) {
                shared_b.push(s.clone());
            } else {
                subj_b.push(s.clone());
            }
        }
        // Objects ascending: object-only (shared terms were emitted above).
        let mut objects: Vec<String> = object_set.into_iter().collect();
        sort_terms(&mut objects);
        for o in &objects {
            if !subject_set.contains(o.as_str()) {
                obj_b.push(o.clone());
            }
        }
        for p in &predicates {
            pred_b.push(p.clone());
        }

        // RDF-star: a quoted triple only ever appears in subject/object position,
        // so a single prefix scan of those term sets detects the whole file.
        let has_quoted_triples = subjects.iter().any(|s| s.starts_with("<<"))
            || objects.iter().any(|o| o.starts_with("<<"));

        let shared = shared_b.build();
        let mut dict =
            Dictionary::from_sections([shared, subj_b.build(), obj_b.build(), pred_b.build()]);
        dict.has_quoted_triples = has_quoted_triples;
        dict
    }
}

/// Reusable per-window scratch for [`Dictionary::resolve_window`]: one job list
/// per section (`(local_id, out_slot)` pairs) and the resolved-term output.
/// Held across windows by the export loop so a window allocates only when a
/// buffer must grow.
#[derive(Default)]
pub struct WindowResolver {
    /// Jobs per dictionary section (shared, subject-only, object-only,
    /// predicates): each `(section-local id, output slot)`.
    jobs: [Vec<(u32, usize)>; 4],
    /// Resolved terms, indexed by output slot `row * 3 + {0:s, 1:p, 2:o}`.
    out: Vec<Option<String>>,
}

/// A read-only dictionary mapping terms ↔ role-specific IDs.
///
/// Holds the four sections as [`ChunkedSection`]s — local files keep each
/// section as one resident chunk; a lazily-opened remote file faults
/// individual chunks in on first touch. Lookups never re-parse a header.
pub struct Dictionary {
    /// shared, subjects, objects, predicates.
    sections: [ChunkedSection; 4],
    shared_len: u32,
    /// Whether any term is an RDF-star quoted triple (`<< … >>`). Set at build
    /// time (from the observed terms); `false` on the read path, where the
    /// header's `FLAG_HAS_QUOTED_TRIPLES` is the source of truth instead.
    has_quoted_triples: bool,
}

impl Dictionary {
    /// Rebuild from four serialized sections (shared, subjects, objects,
    /// predicates), e.g. when reading a `.rete` file. The four sections share one
    /// [`ChunkCache`] (unlimited cap here, so every section body stays resident
    /// exactly as before).
    pub fn from_sections(sections: [Vec<u8>; 4]) -> Self {
        let cache = ChunkCache::unlimited_arc();
        let mut section_index = 0u8;
        let chunked = sections.map(|bytes| {
            let s = ChunkedSection::local(bytes, cache.clone(), section_index);
            section_index += 1;
            s
        });
        Self::from_chunked_sections(chunked)
    }

    /// Rebuild from four already-chunked sections (the remote lazy-open path).
    pub fn from_chunked_sections(sections: [ChunkedSection; 4]) -> Self {
        let shared_len = sections[0].term_count();
        Dictionary {
            sections,
            shared_len,
            has_quoted_triples: false,
        }
    }

    /// Total number of distinct terms across all four sections.
    pub fn term_count(&self) -> u32 {
        self.sections.iter().map(|s| s.term_count()).sum()
    }

    /// Set the byte cap of the shared chunk cache (bounded-export phase 2).
    /// `u64::MAX` (the default from every open) means unlimited — no eviction,
    /// residency identical to the pre-budget design. A finite cap evicts
    /// least-recently-used chunk bodies down to it on every fault, so a full
    /// dump's peak dictionary residency is bounded rather than the whole
    /// decompressed dictionary. All four sections share one cache; setting it on
    /// each is idempotent (they are the same `Arc`).
    pub fn set_cache_cap(&self, cap: u64) {
        for s in &self.sections {
            s.cache().set_cap(cap);
        }
    }

    /// Heap bytes the four sections' chunk directories hold for as long as
    /// the dictionary is open (see [`ChunkedSection::directory_bytes`]).
    ///
    /// [`ChunkedSection::directory_bytes`]: crate::dict::ChunkedSection::directory_bytes
    pub fn directory_bytes(&self) -> u64 {
        self.sections.iter().map(|s| s.directory_bytes()).sum()
    }

    /// The current chunk-cache byte cap (`u64::MAX` = unlimited).
    pub fn cache_cap(&self) -> u64 {
        self.sections[0].cache().cap()
    }

    /// A snapshot of the shared chunk cache's observability counters — for
    /// `RETE_OPEN_DEBUG` accounting and the eviction tests.
    pub fn cache_stats(&self) -> crate::chunk_cache::CacheStats {
        self.sections[0].cache().stats()
    }

    /// Does this dataset contain RDF-star quoted triples? Meaningful only for a
    /// freshly built dictionary (a read-back one carries the truth in the file
    /// header's `FLAG_HAS_QUOTED_TRIPLES`, not here).
    pub fn has_quoted_triples(&self) -> bool {
        self.has_quoted_triples
    }

    /// Did any lazy chunk fetch fail since this dictionary was opened — or
    /// since the last [`reset_load_failure`](Self::reset_load_failure)?
    pub fn load_incomplete(&self) -> bool {
        self.sections.iter().any(|s| s.load_incomplete())
    }

    /// Forget recorded fetch failures across every section (see
    /// `GraphIndex::reset_load_failure` — the per-query reset for resident
    /// sessions). Failed chunks were never cached, so they simply retry.
    pub fn reset_load_failure(&self) {
        self.sections.iter().for_each(|s| s.reset_load_failure());
    }

    /// Batch-fault every unloaded chunk of every section (no-op for local
    /// dictionaries). Callers about to resolve *every* term — export, dump —
    /// call this once so the sweep coalesces into a few range reads instead
    /// of one fetch per chunk.
    pub fn prefetch_all(&self) {
        for s in &self.sections {
            s.prefetch_all();
        }
    }

    /// Batch-fault just the chunks needed to resolve a bounded result set:
    /// `node_ids` are unified-node ids (the `Val::Id(id)` with `id >= 0`),
    /// `predicate_ids` are predicate-space ids. Groups them by `(section,
    /// chunk)` and coalesces each section's faults into a few range reads —
    /// turning "one request per distinct output term" into "a handful per
    /// query". No-op for a local dictionary (every chunk is already resident).
    pub fn prefetch_terms(&self, node_ids: &[u32], predicate_ids: &[u32]) {
        self.prefetch_terms_batched(node_ids, predicate_ids, usize::MAX);
    }

    /// [`prefetch_terms`](Self::prefetch_terms), faulting at most `batch` chunks
    /// per bulk read. A bulk read holds every chunk it fetches — the coalesced
    /// span, the per-chunk slices and the decoded bodies — until all of them
    /// are cached, so an unbatched prefetch of a scattered window materialized
    /// up to thousands of chunks at once, outside every cache cap. A bounded
    /// export batches it so that transient stays a few MB.
    pub fn prefetch_terms_batched(&self, node_ids: &[u32], predicate_ids: &[u32], batch: usize) {
        let batch = batch.max(2);
        let mut want: [std::collections::BTreeSet<usize>; 4] = Default::default();
        for &n in node_ids {
            if let Some((si, ci)) = self.node_chunk(n) {
                want[si].insert(ci);
            }
        }
        for &p in predicate_ids {
            if let Some(ci) = self.sections[3].chunk_of_id(p) {
                want[3].insert(ci);
            }
        }
        for (si, set) in want.iter().enumerate() {
            if set.len() >= 2 {
                let cis: Vec<usize> = set.iter().copied().collect();
                for part in cis.chunks(batch) {
                    self.sections[si].prefetch_chunks(part);
                }
            }
        }
    }

    /// The `(section, chunk)` holding a unified node id — mirrors the section
    /// routing in [`node_term`](Self::node_term).
    fn node_chunk(&self, node: u32) -> Option<(usize, usize)> {
        let su = self.subject_only_count();
        let (si, local) = if node < self.shared_len + su {
            let id = node + 1; // subject_term(id)
            if id <= self.shared_len {
                (0, id)
            } else {
                (1, id - self.shared_len)
            }
        } else {
            let id = node + 1 - su; // object_term(id)
            if id <= self.shared_len {
                (0, id)
            } else {
                (2, id - self.shared_len)
            }
        };
        Some((si, self.sections[si].chunk_of_id(local)?))
    }

    /// Number of shared terms `S`.
    pub fn shared_count(&self) -> u32 {
        self.shared_len
    }

    /// Number of subject-only terms `Su`.
    pub fn subject_only_count(&self) -> u32 {
        self.sections[1].term_count()
    }

    /// Number of object-only terms `Oo`.
    pub fn object_only_count(&self) -> u32 {
        self.sections[2].term_count()
    }

    // --- unified node space (for graph algorithms / the pyramid) ----------
    //
    // Subjects and objects use *separate* ID spaces that overlap on shared
    // terms; the pyramid needs one node per distinct term. We lay them out as:
    //   [0 .. S)            shared
    //   [S .. S+Su)         subject-only
    //   [S+Su .. S+Su+Oo)   object-only

    /// Total distinct graph nodes `N = S + Su + Oo`.
    pub fn node_count(&self) -> u32 {
        self.shared_len + self.subject_only_count() + self.object_only_count()
    }

    /// Unified node ID for a subject-role ID. Shared and subject-only IDs are
    /// already contiguous (`1..=S+Su`), so this is just `sid - 1`.
    pub fn subject_node(&self, sid: SubjectId) -> NodeId {
        // saturating_sub: a corrupt block may carry id 0; valid ids are ≥1 so this
        // is identical for well-formed files but never underflow-panics.
        sid.saturating_sub(1)
    }

    /// Unified node ID for an object-role ID: shared stay at `oid-1`, object-only
    /// IDs are shifted past the subject-only block.
    pub fn object_node(&self, oid: ObjectId) -> NodeId {
        if oid <= self.shared_len {
            oid.saturating_sub(1)
        } else {
            oid.saturating_sub(1) + self.subject_only_count()
        }
    }

    /// Resolve a unified node ID back to its term.
    pub fn node_term(&self, node: NodeId) -> Option<String> {
        let su = self.subject_only_count();
        if node < self.shared_len + su {
            self.subject_term(node + 1)
        } else {
            self.object_term(node + 1 - su)
        }
    }

    /// Batch-fault the dictionary chunks that decoding these subject-role ids
    /// will touch — one coalesced read per section instead of a blocking round
    /// trip per id. Purely a cache warmer (resident sections no-op).
    pub fn prefetch_subject_terms(&self, ids: &[SubjectId]) {
        let mut per_section: [Vec<usize>; 2] = Default::default();
        for &id in ids {
            let (si, local) = if id <= self.shared_len {
                (0usize, id)
            } else {
                (1, id - self.shared_len)
            };
            if let Some(ci) = self.sections[si].chunk_of_term(local) {
                per_section[si].push(ci);
            }
        }
        for (si, cis) in per_section.iter_mut().enumerate() {
            cis.sort_unstable();
            cis.dedup();
            self.sections[si].prefetch_chunks(cis);
        }
    }

    /// [`prefetch_subject_terms`](Self::prefetch_subject_terms) for unified
    /// node ids — the id space solution rows carry. Warms every chunk a batch
    /// of row decodes (a FILTER's string ops, the projection) is about to
    /// fault, using [`node_term`](Self::node_term)'s exact routing.
    pub fn prefetch_node_terms(&self, nodes: &[NodeId]) {
        let su = self.subject_only_count();
        let mut per_section: [Vec<usize>; 3] = Default::default();
        for &node in nodes {
            let (si, local) = if node < self.shared_len + su {
                let sid = node + 1;
                if sid <= self.shared_len {
                    (0usize, sid)
                } else {
                    (1, sid - self.shared_len)
                }
            } else {
                let oid = node + 1 - su;
                if oid <= self.shared_len {
                    (0, oid)
                } else {
                    (2, oid - self.shared_len)
                }
            };
            if let Some(ci) = self.sections[si].chunk_of_term(local) {
                per_section[si].push(ci);
            }
        }
        for (si, cis) in per_section.iter_mut().enumerate() {
            cis.sort_unstable();
            cis.dedup();
            self.sections[si].prefetch_chunks(cis);
        }
    }

    /// Unified node ID for a term, resolving via subject role then object role.
    pub fn node_of_term(&self, term: &str) -> Option<NodeId> {
        if let Some(sid) = self.subject_id(term) {
            return Some(self.subject_node(sid));
        }
        self.object_id(term).map(|oid| self.object_node(oid))
    }

    /// Subject-role ID of a node, or `None` if the node never appears as a
    /// subject (an object-only term).
    pub fn node_as_subject_id(&self, node: NodeId) -> Option<SubjectId> {
        if node < self.shared_len + self.subject_only_count() {
            Some(node + 1)
        } else {
            None
        }
    }

    /// Object-role ID of a node, or `None` if the node never appears as an
    /// object (a subject-only term).
    pub fn node_as_object_id(&self, node: NodeId) -> Option<ObjectId> {
        let (s, su) = (self.shared_len, self.subject_only_count());
        if node < s {
            Some(node + 1) // shared
        } else if node >= s + su {
            Some(node + 1 - su) // object-only
        } else {
            None // subject-only
        }
    }

    // --- term -> id -------------------------------------------------------

    /// Subject-role ID for `term`: shared `1..=S`, else subject-only `S+1..`.
    pub fn subject_id(&self, term: &str) -> Option<SubjectId> {
        if let Some(id) = self.sections[0].id(term) {
            return Some(id);
        }
        self.sections[1].id(term).map(|id| self.shared_len + id)
    }

    /// Object-role ID for `term`: shared `1..=S`, else object-only `S+1..`.
    pub fn object_id(&self, term: &str) -> Option<ObjectId> {
        if let Some(id) = self.sections[0].id(term) {
            return Some(id);
        }
        self.sections[2].id(term).map(|id| self.shared_len + id)
    }

    /// Predicate ID for `term` (independent space).
    pub fn predicate_id(&self, term: &str) -> Option<PredicateId> {
        self.sections[3].id(term)
    }

    // --- id -> term -------------------------------------------------------

    pub fn subject_term(&self, id: SubjectId) -> Option<String> {
        if id <= self.shared_len {
            self.sections[0].term(id)
        } else {
            self.sections[1].term(id - self.shared_len)
        }
    }

    pub fn object_term(&self, id: ObjectId) -> Option<String> {
        if id <= self.shared_len {
            self.sections[0].term(id)
        } else {
            self.sections[2].term(id - self.shared_len)
        }
    }

    pub fn predicate_term(&self, id: PredicateId) -> Option<String> {
        self.sections[3].term(id)
    }

    /// Encode a `(s, p, o)` term triple to its `(subject_id, predicate_id,
    /// object_id)`. `None` if any term is unknown.
    pub fn encode(&self, s: &str, p: &str, o: &str) -> Option<(SubjectId, PredicateId, ObjectId)> {
        Some((
            self.subject_id(s)?,
            self.predicate_id(p)?,
            self.object_id(o)?,
        ))
    }

    /// Resolve a window of `(subject_id, predicate_id, object_id)` triples in
    /// **chunk order** and emit each fully-resolved row through `emit` in the
    /// window's **original order**.
    ///
    /// This is the batch twin of resolving each row with
    /// [`subject_term`](Self::subject_term) /
    /// [`predicate_term`](Self::predicate_term) /
    /// [`object_term`](Self::object_term) and emitting on `(Some, Some, Some)`
    /// — the exact thing it replaces inside `Rete::dump_filtered_each`. A row is
    /// emitted only when all three of its terms resolve, and the bytes of each
    /// term are identical to the per-row path, so the dump stays
    /// byte-for-byte the same.
    ///
    /// What changes is *how* the terms are decoded. Each row contributes three
    /// jobs to the section that owns its id (shared/subject-only for the
    /// subject, predicates for the predicate, shared/object-only for the
    /// object); each section then resolves its jobs in one chunk-ordered,
    /// single-pass walk ([`ChunkedSection::resolve_into`]). A window whose
    /// object ids scatter across the object-only section therefore touches each
    /// chunk once — one faulted body, one forward walk — instead of the
    /// scattered re-faults and per-row re-walks the row-order path performed.
    ///
    /// `w` carries the per-window scratch buffers so back-to-back windows cost
    /// no allocation beyond growth.
    ///
    /// [`ChunkedSection::resolve_into`]: crate::dict::ChunkedSection::resolve_into
    pub fn resolve_window(
        &self,
        ids: &[(SubjectId, PredicateId, ObjectId)],
        w: &mut WindowResolver,
        mut emit: impl FnMut(&str, &str, &str),
    ) {
        for j in &mut w.jobs {
            j.clear();
        }
        w.out.clear();
        w.out.resize(ids.len() * 3, None);
        for (i, &(s, p, o)) in ids.iter().enumerate() {
            // Subject role: shared (section 0) if `id <= S`, else subject-only
            // (section 1) with the shared block subtracted — the exact routing
            // of `subject_term`.
            if s <= self.shared_len {
                w.jobs[0].push((s, i * 3));
            } else {
                w.jobs[1].push((s - self.shared_len, i * 3));
            }
            // Predicate role: its own section (3).
            w.jobs[3].push((p, i * 3 + 1));
            // Object role: shared (section 0) if `id <= S`, else object-only
            // (section 2) — the exact routing of `object_term`.
            if o <= self.shared_len {
                w.jobs[0].push((o, i * 3 + 2));
            } else {
                w.jobs[2].push((o - self.shared_len, i * 3 + 2));
            }
        }
        for si in 0..4 {
            self.sections[si].resolve_into(&mut w.jobs[si], &mut w.out);
        }
        for i in 0..ids.len() {
            if let (Some(s), Some(p), Some(o)) =
                (&w.out[i * 3], &w.out[i * 3 + 1], &w.out[i * 3 + 2])
            {
                emit(s, p, o);
            }
        }
    }

    /// The four serialized sections (header + body each), for the file writer.
    pub fn sections(&self) -> [Vec<u8>; 4] {
        [
            self.sections[0].raw_section_bytes(),
            self.sections[1].raw_section_bytes(),
            self.sections[2].raw_section_bytes(),
            self.sections[3].raw_section_bytes(),
        ]
    }
}

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

    fn dict() -> Dictionary {
        let mut b = DictionaryBuilder::new();
        // Alice knows Bob; Bob knows Carol. So Bob is both subject and object
        // (shared); Alice is subject-only; Carol is object-only.
        b.observe("Alice", "knows", "Bob");
        b.observe("Bob", "knows", "Carol");
        b.build()
    }

    #[test]
    fn shared_term_has_same_id_in_both_roles() {
        let d = dict();
        assert_eq!(d.shared_count(), 1); // just "Bob"
        let sid = d.subject_id("Bob").unwrap();
        let oid = d.object_id("Bob").unwrap();
        assert_eq!(sid, oid, "shared term must share its ID across roles");
        assert_eq!(sid, 1, "shared terms get the lowest IDs");
    }

    #[test]
    fn role_specific_terms_round_trip() {
        let d = dict();
        // subject-only Alice
        let a = d.subject_id("Alice").unwrap();
        assert!(a > d.shared_count());
        assert_eq!(d.subject_term(a).as_deref(), Some("Alice"));
        // object-only Carol
        let c = d.object_id("Carol").unwrap();
        assert!(c > d.shared_count());
        assert_eq!(d.object_term(c).as_deref(), Some("Carol"));
        // predicate
        let k = d.predicate_id("knows").unwrap();
        assert_eq!(d.predicate_term(k).as_deref(), Some("knows"));
        // Alice never appears as object, Carol never as subject.
        assert_eq!(d.object_id("Alice"), None);
        assert_eq!(d.subject_id("Carol"), None);
    }

    #[test]
    fn encode_full_triples() {
        let d = dict();
        let t1 = d.encode("Alice", "knows", "Bob").unwrap();
        let t2 = d.encode("Bob", "knows", "Carol").unwrap();
        // Bob is object in t1 and subject in t2 — same ID both times.
        assert_eq!(t1.2, t2.0);
        assert!(d.encode("Nobody", "knows", "Bob").is_none());
    }
}