kglite 0.16.4

//! Thin pure-Rust graph handle for Rust embedders.
//!
//! Bridges `Arc<DirGraph>` (the engine) and the minimal set of
//! convenience methods that protocol-server binaries
//! (`kglite-mcp-server`, `kglite-bolt-server`) and other Rust
//! embedders need without taking on the wheel crate's full
//! Python-flavored state (selection / reports / mutation stats /
//! temporal context / default timeout / default max rows).
//!
//! This is the **Rust-side** `KnowledgeGraph`. The Python-side
//! `KnowledgeGraph` (the `#[pyclass]` wrapper backing
//! `pip install kglite`'s `import kglite`) lives in the
//! `kglite-py` crate at `crates/kglite-py/src/graph/mod.rs`. Two
//! types named `KnowledgeGraph` exist in distinct crates with
//! distinct audiences; mirrors the polars precedent
//! (`polars::DataFrame` vs `polars.DataFrame`).
//!
//! The heavy logic — `source_location` + `resolve_code_entity` —
//! lives as free functions in this module so the wheel's full
//! `KnowledgeGraph` can delegate to the same implementation,
//! keeping the single source of truth in `kglite` (the core).

use std::sync::Arc;

use petgraph::graph::NodeIndex;

use crate::datatypes::values::{raw_string, Value};
use crate::graph::dir_graph::DirGraph;
use crate::graph::embedder::Embedder;
use crate::graph::schema;
use crate::graph::storage::GraphRead;
use crate::graph::{SourceLocation, SourceLookup};

/// Code-entity node types used by `source_location` / `resolve_code_entity`
/// when the caller doesn't specify a `node_type`. Matches what the
/// code-graph builders (e.g. codingest) emit — language-specific subsets (Rust:
/// `Struct`/`Enum`/`Trait`; Python: `Class`/`Mixin`/`Protocol`; etc.)
/// are all listed so a single search covers every supported source
/// language.
pub const CODE_TYPES: &[&str] = &[
    "Function",
    "Struct",
    "Class",
    "Mixin",
    "Enum",
    "Trait",
    "Protocol",
    "Interface",
    "Module",
    "Constant",
];

/// Name-matching strategy for [`find_code_entities`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CodeEntityMatch {
    Exact,
    Contains,
    StartsWith,
}

/// Search code-entity type indices by `name` or `title`.
///
/// This is the binding-neutral scan behind the wheel's `find()` method. It
/// returns typed [`schema::NodeInfo`] values; dict/object marshalling remains
/// in the consuming wrapper.
pub fn find_code_entities(
    dir: &Arc<DirGraph>,
    name: &str,
    node_type: Option<&str>,
    match_type: CodeEntityMatch,
) -> Vec<schema::NodeInfo> {
    let _arena_guard = dir.graph.begin_query();
    let name_lower = name.to_lowercase();
    let name_value = Value::String(name.to_string());
    let types_to_search: Vec<&str> = match node_type {
        Some(nt) => vec![nt],
        None => CODE_TYPES.to_vec(),
    };

    let mut results = Vec::new();
    for node_type in types_to_search {
        let Some(indices) = dir.type_indices.get(node_type) else {
            continue;
        };
        for index in indices.iter() {
            let Some(node) = dir.node_view(index) else {
                continue;
            };
            // `title` is a primary NodeData field, not an ordinary property.
            // Resolve it explicitly: `field_*_ci("title")` only covers the
            // property store and silently missed titles extracted at load.
            let title = node.title();
            let title_string = match &*title {
                Value::String(value) => Some(value.as_str()),
                _ => None,
            };
            let matches = match match_type {
                CodeEntityMatch::Contains => {
                    node.field_contains_ci("name", &name_lower)
                        || title_string
                            .is_some_and(|value| value.to_lowercase().contains(&name_lower))
                }
                CodeEntityMatch::StartsWith => {
                    node.field_starts_with_ci("name", &name_lower)
                        || title_string
                            .is_some_and(|value| value.to_lowercase().starts_with(&name_lower))
                }
                CodeEntityMatch::Exact => {
                    node.get_field_ref("name")
                        .is_some_and(|value| *value == name_value)
                        || *title == name_value
                }
            };
            if matches {
                results.push(node.to_node_info(&dir.interner));
            }
        }
    }
    results
}

/// Resolved code-entity neighborhood, kept directional for neutral bindings.
#[derive(Debug)]
pub struct CodeEntityContext {
    pub node: schema::NodeInfo,
    pub defined_in: Option<String>,
    pub outgoing: std::collections::HashMap<String, Vec<schema::NodeInfo>>,
    pub incoming: std::collections::HashMap<String, Vec<schema::NodeInfo>>,
}

/// Outcome of resolving a code entity for [`code_entity_context`].
#[derive(Debug)]
pub enum CodeContextLookup {
    Found(Box<CodeEntityContext>),
    Ambiguous(Vec<schema::NodeInfo>),
    NotFound,
}

/// Resolve a code entity and collect its neighborhood up to `hops` away.
///
/// The traversal and edge-type grouping are shared engine logic. Bindings may
/// flatten or rename the directional groups to suit their native result shape.
pub fn code_entity_context(
    dir: &Arc<DirGraph>,
    name: &str,
    node_type: Option<&str>,
    hops: usize,
) -> CodeContextLookup {
    let _arena_guard = dir.graph.begin_query();
    let (resolved, matches) = resolve_code_entity(dir, name, node_type);
    let Some(target_idx) = resolved else {
        return if matches.is_empty() {
            CodeContextLookup::NotFound
        } else {
            CodeContextLookup::Ambiguous(matches.into_iter().map(|(_, info)| info).collect())
        };
    };
    let Some(target_node) = dir.node_view(target_idx) else {
        return CodeContextLookup::NotFound;
    };

    let neighbor_indices = if hops <= 1 {
        let mut neighbors = std::collections::HashSet::new();
        for edge in dir
            .graph
            .edges_directed(target_idx, petgraph::Direction::Outgoing)
        {
            neighbors.insert(edge.target());
        }
        for edge in dir
            .graph
            .edges_directed(target_idx, petgraph::Direction::Incoming)
        {
            neighbors.insert(edge.source());
        }
        neighbors
    } else {
        let mut visited = std::collections::HashSet::from([target_idx]);
        let mut frontier = std::collections::HashSet::from([target_idx]);
        for _ in 0..hops {
            let mut next_frontier = std::collections::HashSet::new();
            for &node in &frontier {
                for neighbor in dir.graph.neighbors_undirected(node) {
                    if visited.insert(neighbor) {
                        next_frontier.insert(neighbor);
                    }
                }
            }
            if next_frontier.is_empty() {
                break;
            }
            frontier = next_frontier;
        }
        visited.remove(&target_idx);
        visited
    };

    let mut outgoing_indices: std::collections::HashMap<String, Vec<NodeIndex>> =
        std::collections::HashMap::new();
    let mut incoming_indices: std::collections::HashMap<String, Vec<NodeIndex>> =
        std::collections::HashMap::new();
    for edge in dir
        .graph
        .edges_directed(target_idx, petgraph::Direction::Outgoing)
    {
        let target = edge.target();
        if hops <= 1 || neighbor_indices.contains(&target) {
            outgoing_indices
                .entry(edge.weight().connection_type_str(&dir.interner).to_string())
                .or_default()
                .push(target);
        }
    }
    for edge in dir
        .graph
        .edges_directed(target_idx, petgraph::Direction::Incoming)
    {
        let source = edge.source();
        if hops <= 1 || neighbor_indices.contains(&source) {
            incoming_indices
                .entry(edge.weight().connection_type_str(&dir.interner).to_string())
                .or_default()
                .push(source);
        }
    }
    if hops > 1 {
        for &node_idx in &neighbor_indices {
            for edge in dir
                .graph
                .edges_directed(node_idx, petgraph::Direction::Outgoing)
            {
                let target = edge.target();
                if target != target_idx && neighbor_indices.contains(&target) {
                    outgoing_indices
                        .entry(edge.weight().connection_type_str(&dir.interner).to_string())
                        .or_default()
                        .push(target);
                }
            }
        }
    }

    let materialise_groups = |groups: std::collections::HashMap<String, Vec<NodeIndex>>| {
        groups
            .into_iter()
            .map(|(edge_type, indices)| {
                let mut seen = std::collections::HashSet::new();
                let nodes = indices
                    .into_iter()
                    .filter(|index| seen.insert(*index))
                    .filter_map(|index| dir.node_view(index))
                    .map(|node| node.to_node_info(&dir.interner))
                    .collect();
                (edge_type, nodes)
            })
            .collect()
    };

    CodeContextLookup::Found(Box::new(CodeEntityContext {
        node: target_node.to_node_info(&dir.interner),
        defined_in: match target_node.get_field_ref("file_path").as_deref() {
            Some(Value::String(path)) => Some(path.clone()),
            _ => None,
        },
        outgoing: materialise_groups(outgoing_indices),
        incoming: materialise_groups(incoming_indices),
    }))
}

/// Resolve a name (or qualified-name suffix) to a single code-entity
/// `NodeIndex`. Returns `(Some(idx), Vec::new())` for an unambiguous
/// match, `(None, matches)` when 0 or >1 candidates matched.
///
/// Lookup order:
/// 1. Exact match on `node.id()` (the qualified name, e.g.
///    `crate::graph::languages::cypher::executor::CypherExecutor::execute_single_clause`)
/// 2. Suffix match on `node.id()` if `name` contains `::`
///    (e.g. `CypherExecutor::execute_single_clause` matches the above)
/// 3. Exact match on `node.get_field_ref("name")` or
///    `node.get_field_ref("title")` — bare-name fallback
///
/// When `node_type` is `None`, searches across every entry in
/// [`CODE_TYPES`]; otherwise restricted to the single type.
pub fn resolve_code_entity(
    dir: &Arc<DirGraph>,
    name: &str,
    node_type: Option<&str>,
) -> (Option<NodeIndex>, Vec<(NodeIndex, schema::NodeInfo)>) {
    // Arena guard: disk-backed node reads materialize into the query arena
    // (protocol in disk/graph.rs); no-op on memory/mapped.
    let _arena_guard = dir.graph.begin_query();
    let name_val = Value::String(name.to_string());
    let types_to_search: Vec<&str> = match node_type {
        Some(nt) => vec![nt],
        None => CODE_TYPES.to_vec(),
    };

    // Try qualified_name (stored as "id") exact match first
    for nt in &types_to_search {
        if let Some(indices) = dir.type_indices.get(nt) {
            for idx in indices.iter() {
                if let Some(node) = dir.node_view(idx) {
                    if *node.id() == name_val {
                        return (Some(idx), Vec::new());
                    }
                }
            }
        }
    }

    // Try qualified_name suffix match (e.g. "CypherExecutor::execute_single_clause"
    // matches "crate::graph::languages::cypher::executor::CypherExecutor::execute_single_clause")
    if name.contains("::") {
        let suffix = format!("::{}", name);
        let mut matches: Vec<(NodeIndex, schema::NodeInfo)> = Vec::new();
        for nt in &types_to_search {
            if let Some(indices) = dir.type_indices.get(nt) {
                for idx in indices.iter() {
                    if let Some(node) = dir.node_view(idx) {
                        if let Value::String(qn) = &*node.id() {
                            if qn.ends_with(&suffix) {
                                matches.push((idx, node.to_node_info(&dir.interner)));
                            }
                        }
                    }
                }
            }
        }
        if matches.len() == 1 {
            return (Some(matches[0].0), matches);
        } else if !matches.is_empty() {
            return (None, matches);
        }
    }

    // Fall back to name/title search
    let mut matches: Vec<(NodeIndex, schema::NodeInfo)> = Vec::new();
    for nt in &types_to_search {
        if let Some(indices) = dir.type_indices.get(nt) {
            for idx in indices.iter() {
                if let Some(node) = dir.node_view(idx) {
                    let name_match = node
                        .get_field_ref("name")
                        .map(|v| *v == name_val)
                        .unwrap_or(false)
                        || node
                            .get_field_ref("title")
                            .map(|v| *v == name_val)
                            .unwrap_or(false);
                    if name_match {
                        matches.push((idx, node.to_node_info(&dir.interner)));
                    }
                }
            }
        }
    }

    if matches.len() == 1 {
        (Some(matches[0].0), matches)
    } else {
        (None, matches)
    }
}

/// Infer the node type of the current (latest level) selection by
/// sampling the first node. Returns `None` if the selection is empty
/// or the node disappeared.
///
/// **Not re-exported through `kglite::api`** — it takes a
/// `&CowSelection`, which is currently only used externally by the
/// Python wheel's fluent-API surface. A future binding cannot
/// meaningfully call this without first lifting the `Selection`
/// concept to be a stable api type. When that happens, both should
/// move to api together. The wheel reaches this directly via
/// `kglite_core::graph::handle::infer_selection_node_type` for now
/// (see `crates/kglite-py/src/graph/mod.rs`).
pub fn infer_selection_node_type(
    selection: &crate::graph::schema::CowSelection,
    dir: &Arc<DirGraph>,
) -> Option<String> {
    let level_idx = selection.get_level_count().saturating_sub(1);
    let level = selection.get_level(level_idx)?;
    let first_idx = level.iter_node_indices().next()?;
    // Arena guard: node_weight materializes on the disk backend (protocol
    // in disk/graph.rs); no-op on memory/mapped.
    let _arena_guard = dir.graph.begin_query();
    dir.graph
        .node_view(first_idx)
        .map(|n| n.node_type_str(&dir.interner).to_string())
}

/// Column names a row-oriented exporter emits from a node's **canonical
/// identity** rather than from its property bag.
///
/// A node's `id`, `title`, and structural `type` are virtuals: every table
/// exporter writes them as leading columns straight from the node header. A
/// node may *also* carry a stored property under one of these names — Cypher
/// `CREATE (:T {title: 'a'})` sets `title` both ways — and an exporter that
/// naively appends every discovered property key then emits that column
/// twice. See [`is_canonical_node_column`].
pub const CANONICAL_NODE_COLUMNS: [&str; 3] = ["id", "title", "type"];

/// Whether `key` names a column a row-oriented exporter already emits from
/// the node's canonical identity.
///
/// Property keys that collide with a canonical column are dropped from the
/// discovered property set: the canonical value wins. This is the rule the
/// SQL-dump, d3/JSON, and `to_text` exporters have always applied, and the
/// only rule that keeps a header unique. Emitting the column twice is not a
/// lossless alternative — a name-keyed column map silently overwrites the
/// canonical value with the property, so the duplicate *destroys* the
/// identity it appears to preserve.
pub fn is_canonical_node_column(key: &str) -> bool {
    CANONICAL_NODE_COLUMNS.contains(&key)
}

/// Discover all unique property keys across a slice of typed nodes.
/// Returns sorted, de-duplicated key names — useful for any
/// row-oriented exporter (CSV, Parquet, DataFrame, JSON-lines) that
/// needs a stable column-name set without scanning the entire graph
/// schema. The function takes only core types (`NodeData`,
/// `StringInterner`) so every binding's table-export path can call
/// it directly.
///
/// Keys naming a canonical identity column ([`CANONICAL_NODE_COLUMNS`]) are
/// excluded, so appending the result to the exporter's leading identity
/// columns always yields a header with unique names.
pub fn discover_property_keys_from_data(
    nodes: &[(&str, crate::graph::storage::NodeView<'_>)],
    interner: &crate::graph::schema::StringInterner,
) -> Vec<String> {
    discover_property_keys_excluding(nodes, interner, &CANONICAL_NODE_COLUMNS)
}

/// [`discover_property_keys_from_data`] with an explicit exclusion set.
///
/// For an exporter that emits only *some* canonical columns — the fluent
/// `to_df(include_type=False)` drops the structural `type` column — pass the
/// names actually emitted. A canonical name that is *not* emitted carries no
/// collision, so a stored property under that name is real user data and must
/// survive.
pub fn discover_property_keys_excluding(
    nodes: &[(&str, crate::graph::storage::NodeView<'_>)],
    interner: &crate::graph::schema::StringInterner,
    excluded: &[&str],
) -> Vec<String> {
    let mut seen = std::collections::HashSet::new();
    let mut keys = Vec::new();
    for (_, node) in nodes {
        for key in node.property_keys(interner) {
            if excluded.contains(&key) {
                continue;
            }
            if seen.insert(key.to_string()) {
                keys.push(key.to_string());
            }
        }
    }
    keys.sort();
    keys
}

/// Look up the source-file location for a code-entity node.
///
/// Drives the `read_code_source` MCP tool's "qualified_name →
/// (file_path, line_number, end_line, signature)" mapping. The
/// returned [`SourceLookup`] enum distinguishes a unique match
/// ([`SourceLookup::Found`]) from ambiguous candidates
/// ([`SourceLookup::Ambiguous`] with qualified-name suggestions)
/// from a miss ([`SourceLookup::NotFound`]).
///
/// All optional fields on [`SourceLocation`] mirror the
/// corresponding node fields. Graphs built from non-code-tree
/// sources (e.g. a codingest-built code graph, or a
/// manually-constructed graph) may have fewer populated.
pub fn source_location(dir: &Arc<DirGraph>, name: &str, node_type: Option<&str>) -> SourceLookup {
    // Arena guard: disk-backed node reads materialize into the query arena
    // (protocol in disk/graph.rs); no-op on memory/mapped.
    let _arena_guard = dir.graph.begin_query();
    let (resolved, matches) = resolve_code_entity(dir, name, node_type);

    if let Some(target_idx) = resolved {
        let node = match dir.node_view(target_idx) {
            Some(n) => n,
            None => return SourceLookup::NotFound,
        };
        let type_name = node.get_node_type_ref(&dir.interner).to_string();
        let entity_name = raw_string(&node.title());
        let qname = raw_string(&node.id());
        let file_path = node.get_field_ref("file_path").as_deref().map(raw_string);
        let line_number = node
            .get_field_ref("line_number")
            .as_deref()
            .and_then(|v| match v {
                Value::Int64(n) => Some(*n),
                _ => None,
            });
        let end_line = node
            .get_field_ref("end_line")
            .as_deref()
            .and_then(|v| match v {
                Value::Int64(n) => Some(*n),
                _ => None,
            });
        let signature = node.get_field_ref("signature").as_deref().map(raw_string);
        SourceLookup::Found(SourceLocation {
            type_name,
            name: entity_name,
            qualified_name: qname,
            file_path,
            line_number,
            end_line,
            signature,
        })
    } else if matches.is_empty() {
        SourceLookup::NotFound
    } else {
        let qnames: Vec<String> = matches
            .iter()
            .map(|(_, info)| raw_string(&info.id))
            .collect();
        SourceLookup::Ambiguous(qnames)
    }
}

/// Thin pure-Rust graph handle. Holds an `Arc<DirGraph>` plus an
/// optional [`Embedder`] for `text_score()` queries. For Rust
/// embedders (mcp-server, bolt-server, third-party binaries) that
/// don't need the Python wheel's full state.
///
/// The Python wheel's `KnowledgeGraph` (in `kglite-py`) has the
/// same name but adds wheel-API state (selection, reports,
/// mutation stats, temporal context, default timeout / max-rows).
/// The two types don't share a definition; pick whichever fits
/// your audience.
pub struct KnowledgeGraph {
    inner: Arc<DirGraph>,
    embedder: Option<Arc<dyn Embedder>>,
}

impl KnowledgeGraph {
    /// Wrap an existing `Arc<DirGraph>` (e.g. one returned by
    /// [`crate::graph::io::file::load_file`] or
    /// an external code-graph builder such as codingest) into a
    /// `KnowledgeGraph` handle with no embedder set.
    pub fn from_arc(inner: Arc<DirGraph>) -> Self {
        KnowledgeGraph {
            inner,
            embedder: None,
        }
    }

    /// Borrow the underlying `Arc<DirGraph>`. Use this to reach
    /// the engine surface (`compute_schema`, `execute_read`,
    /// `compute_description`, ...) which all take `&DirGraph`.
    pub fn dir(&self) -> &Arc<DirGraph> {
        &self.inner
    }

    /// Mutable borrow of the underlying `Arc<DirGraph>` — the write
    /// counterpart of [`dir`](Self::dir). Pair with
    /// [`make_dir_graph_mut`] to obtain a `&mut DirGraph` for the
    /// mutation surface (`execute_mut`, …). Used by bindings that hold a
    /// long-lived handle and mutate it in place (e.g. the write-enabled
    /// MCP server), so the mutation lands on *this* handle's graph rather
    /// than a detached clone.
    pub fn dir_mut(&mut self) -> &mut Arc<DirGraph> {
        &mut self.inner
    }

    /// Bind an embedder implementing the [`Embedder`] trait — used
    /// by `text_score()` Cypher to map text queries onto stored
    /// vectors. Replaces any previously-bound embedder. Callers
    /// that wrap a Python embedder object should construct an
    /// adapter in the wheel crate; pure-Rust callers can pass
    /// e.g. `Arc::new(FastEmbedAdapter::new("bge-small")?)`.
    pub fn set_embedder_native(&mut self, embedder: Arc<dyn Embedder>) {
        self.embedder = Some(embedder);
    }

    /// Access the active embedder, if any. Returns `None` until
    /// [`set_embedder_native`](Self::set_embedder_native) has been
    /// called.
    pub fn embedder(&self) -> Option<&Arc<dyn Embedder>> {
        self.embedder.as_ref()
    }

    /// Look up the source-file location for a code-entity node by
    /// name (or qualified-name suffix). Delegates to the
    /// [`source_location`] free function so the wheel crate's
    /// `KnowledgeGraph` can share the same implementation.
    pub fn source_location(&self, name: &str, node_type: Option<&str>) -> SourceLookup {
        source_location(&self.inner, name, node_type)
    }
}

/// Get a `&mut DirGraph` from an `Arc<DirGraph>` and bump the version
/// counter. Wraps [`Arc::make_mut`] (which clones the inner `DirGraph`
/// if other strong refs exist) plus the canonical post-mutation version
/// increment that downstream OCC commit-checks + the plan cache rely on.
///
/// Lifted from the wheel crate in 0.10.1 so bindings + embedders that
/// hold an `Arc<DirGraph>` and want to mutate it have a single,
/// consistent entry point. Re-exported as `kglite::api::make_dir_graph_mut`.
/// (Homed here rather than in `dir_graph.rs` to keep that file under the
/// god-file ceiling.)
///
/// **Cost when other `Arc<DirGraph>` references exist** (a snapshot held by an
/// open transaction, a clone held by a still-alive `ResultView`, a `freeze()`):
///
/// * **Memory mode — a copy-on-write fork, not a copy.** D2 replaced the
///   whole-graph clone this warning used to describe: the backend forks to an
///   overlay over the shared data and the indexes layer over shared levels, so
///   the write is O(write) and the overlay folds back on the first write after
///   the last reader drops. See `docs/rust/structural-sharing.md`, and
///   `held_reference_clone_tests` below for the executable form.
/// * **Mapped and disk modes still deep-copy**, so a lingering reference there
///   does cost a full copy — every node, edge and index — on the first write.
/// * An adjacency edit (adding or removing an edge, deleting a node) is not
///   overlay-expressible and **flattens** the fork: one copy, paid once per
///   fork rather than once per statement.
///
/// Copy-on-write access that preserves disk writer authority when a shared
/// snapshot forces a clone. Does not change the graph version; callers that
/// perform semantic mutations should use [`make_dir_graph_mut`].
pub(crate) fn make_dir_graph_mut_preserving_lineage(arc: &mut Arc<DirGraph>) -> &mut DirGraph {
    let parent = if Arc::get_mut(arc).is_none() {
        Some(Arc::clone(arc))
    } else {
        None
    };
    let graph = Arc::make_mut(arc);
    if let Some(parent) = parent {
        graph.graph.adopt_shared_writer_lineage(&parent.graph);
    }
    // D2 Phase 2 compaction point. If this graph is a copy-on-write overlay and
    // the reader that forced the fork has since dropped, fold the delta back
    // into the base here and return to the flat representation. This is the
    // earliest moment the writer can observe the reader's departure —
    // `Arc::get_mut` succeeding *is* that observation — so "hold a view, write,
    // drop the view, write again" self-heals on the very next write.
    graph.graph.try_compact();
    // ...and the mirror image: this graph may itself be somebody else's base
    // (`g.copy()` forks *from* `g`), in which case writing in place would edit a
    // backend the fork is reading. `ensure_writable` turns `g` into an overlay
    // over the shared base too — one `Arc::get_mut` probe when nothing is
    // shared, which is the steady state.
    graph.graph.ensure_writable();
    // The same fold for `id_indices`, whose entries layer over a shared base of
    // their own (D2 Phase 3). Per entry this is `Arc::get_mut` + an O(delta)
    // merge, so it is a probe when nothing is shared.
    graph.id_indices.try_compact();
    // ...and for `type_indices`, whose buckets are stacks of shared levels
    // (D2). Per bucket this is an `Arc::get_mut` probe plus an O(delta) merge.
    graph.type_indices.try_compact();
    // ...and the three user index families, same mechanism over their
    // `value -> members` maps (`dir_graph/index_layer.rs`, and
    // `dir_graph/range_index_layer.rs` for the ordered one). One probe per
    // declared index, and the loops do not run at all on the overwhelmingly
    // common graph that has no user index.
    for index in graph.property_indices.values_mut() {
        index.try_compact();
    }
    for index in graph.composite_indices.values_mut() {
        index.try_compact();
    }
    for index in graph.range_indices.values_mut() {
        index.try_compact();
    }
    graph
}

pub fn make_dir_graph_mut(arc: &mut Arc<DirGraph>) -> &mut DirGraph {
    let graph = make_dir_graph_mut_preserving_lineage(arc);
    graph.bump_version();
    graph
}

#[cfg(test)]
mod boundary_lift_tests {
    use super::*;
    use crate::graph::session::{execute_mut, ExecuteOptions};
    use std::collections::HashMap;

    fn code_graph() -> Arc<DirGraph> {
        let mut graph = DirGraph::new();
        let params = HashMap::new();
        execute_mut(
            &mut graph,
            "CREATE (a:Function {id:'mod::alpha', title:'alpha', name:'alpha', file_path:'src/a.rs'}), \
             (b:Function {id:'mod::beta', title:'BetaWorker', name:'beta', file_path:'src/b.rs'}), \
             (c:Function {id:'mod::gamma', title:'gamma', name:'gamma', file_path:'src/c.rs'}), \
             (f:File {id:'src/a.rs', title:'src/a.rs'})",
            &ExecuteOptions::eager(&params),
        )
        .expect("fixture nodes");
        execute_mut(
            &mut graph,
            "MATCH (a:Function {id:'mod::alpha'}), (b:Function {id:'mod::beta'}), \
             (c:Function {id:'mod::gamma'}), (f:File {id:'src/a.rs'}) \
             CREATE (a)-[:CALLS]->(b), (b)-[:CALLS]->(c), (f)-[:DEFINES]->(a)",
            &ExecuteOptions::eager(&params),
        )
        .expect("fixture edges");
        Arc::new(graph)
    }

    #[test]
    fn find_code_entities_supports_match_modes_and_type_filter() {
        let graph = code_graph();
        let exact = find_code_entities(&graph, "alpha", Some("Function"), CodeEntityMatch::Exact);
        assert_eq!(exact.len(), 1);
        assert_eq!(exact[0].id, Value::String("mod::alpha".into()));

        let contains = find_code_entities(&graph, "et", None, CodeEntityMatch::Contains);
        assert_eq!(contains.len(), 1);
        assert_eq!(contains[0].id, Value::String("mod::beta".into()));

        let starts_with = find_code_entities(&graph, "bet", None, CodeEntityMatch::StartsWith);
        assert_eq!(starts_with.len(), 1);
        assert_eq!(starts_with[0].id, Value::String("mod::beta".into()));
    }

    #[test]
    fn code_entity_context_groups_directional_multi_hop_neighbors() {
        let graph = code_graph();
        let CodeContextLookup::Found(context) =
            code_entity_context(&graph, "alpha", Some("Function"), 2)
        else {
            panic!("expected resolved context");
        };
        assert_eq!(context.defined_in.as_deref(), Some("src/a.rs"));
        let calls = &context.outgoing["CALLS"];
        assert_eq!(calls.len(), 2);
        assert!(calls
            .iter()
            .any(|node| node.id == Value::String("mod::beta".into())));
        assert!(calls
            .iter()
            .any(|node| node.id == Value::String("mod::gamma".into())));
        assert_eq!(context.incoming["DEFINES"].len(), 1);
    }

    #[test]
    fn code_entity_context_distinguishes_miss_from_ambiguity() {
        let mut graph = match Arc::try_unwrap(code_graph()) {
            Ok(graph) => graph,
            Err(_) => panic!("expected sole graph owner"),
        };
        let params = HashMap::new();
        execute_mut(
            &mut graph,
            "CREATE (:Function {id:'other::alpha', title:'alpha', name:'alpha', file_path:'other.rs'})",
            &ExecuteOptions::eager(&params),
        )
        .expect("add ambiguous entity");
        let graph = Arc::new(graph);
        assert!(matches!(
            code_entity_context(&graph, "alpha", Some("Function"), 1),
            CodeContextLookup::Ambiguous(matches) if matches.len() == 2
        ));
        assert!(matches!(
            code_entity_context(&graph, "missing", None, 1),
            CodeContextLookup::NotFound
        ));
    }

    /// A graph whose nodes store `title` (and one storing `type`) in the
    /// property bag as well as in the canonical header — what Cypher `CREATE`
    /// produces for any ordinary graph.
    fn collision_graph() -> Arc<DirGraph> {
        let mut graph = DirGraph::new();
        let params = HashMap::new();
        execute_mut(
            &mut graph,
            "CREATE (:T {id:1, title:'a', v:2}), (:T {id:2, title:'b', type:'USER', w:3})",
            &ExecuteOptions::eager(&params),
        )
        .expect("fixture nodes");
        Arc::new(graph)
    }

    fn nodes_of(graph: &DirGraph) -> Vec<(&str, crate::graph::storage::NodeView<'_>)> {
        graph
            .graph
            .node_indices()
            .filter_map(|idx| {
                graph
                    .node_view(idx)
                    .map(|n| (n.node_type_str(&graph.interner), n))
            })
            .collect()
    }

    #[test]
    fn discovered_property_keys_exclude_canonical_columns() {
        // Every row exporter emits id/title/type from the node header, so a
        // stored property of the same name must not become a second column.
        let graph = collision_graph();
        let keys = discover_property_keys_from_data(&nodes_of(&graph), &graph.interner);
        assert_eq!(keys, vec!["v".to_string(), "w".to_string()]);
        for canonical in CANONICAL_NODE_COLUMNS {
            assert!(
                !keys.contains(&canonical.to_string()),
                "canonical column {canonical} leaked into the property key set"
            );
        }
    }

    #[test]
    fn an_unemitted_canonical_column_keeps_its_stored_property() {
        // `to_df(include_type=False)` emits no `type` column, so there is no
        // collision and the stored `type` property is real user data.
        let graph = collision_graph();
        let keys =
            discover_property_keys_excluding(&nodes_of(&graph), &graph.interner, &["id", "title"]);
        assert_eq!(
            keys,
            vec!["type".to_string(), "v".to_string(), "w".to_string()]
        );
    }

    #[test]
    fn is_canonical_node_column_covers_exactly_the_identity_names() {
        assert!(is_canonical_node_column("id"));
        assert!(is_canonical_node_column("title"));
        assert!(is_canonical_node_column("type"));
        assert!(!is_canonical_node_column("titles"));
        assert!(!is_canonical_node_column("node_type"));
        assert!(!is_canonical_node_column("Title"));
    }
}

/// The cost of holding a reader across a write — pinned as a *count*, not a
/// timing.
///
/// [`make_dir_graph_mut`] deep-clones the whole graph whenever a second
/// `Arc<DirGraph>` is alive, which the doc on
/// [`make_dir_graph_mut_preserving_lineage`] warns about. The reachable-from-
/// ordinary-code shape is mundane: a request handler keeps a query result in a
/// local, then writes. No snapshot API, no threading, no `freeze()`.
///
/// This is measured elsewhere as wall time — and that measurement is
/// structurally hard to gate on. `min`-of-N and p95 both *hide* it, because
/// only the first write after acquiring the reference pays; the Python
/// benchmark that does see it (`test_bench_first_write_after_reference`) has to
/// re-acquire the reference in an untimed `pedantic` setup before every round,
/// and even then it asserts no threshold, and lives behind a marker the default
/// suite deselects. Meanwhile `bench-check` and `bench-anchor` both run
/// `--metric min`, so neither can ever see this.
///
/// `BACKEND_CLONE_NODES` sidesteps all of it. It counts nodes copied, so the
/// intentional O(1) clone of an emptied backend registers zero while a real
/// fork registers the node count. That turns a ~28 ms cliff at 1M nodes into an
/// exact integer at ten — no idle machine, no statistics, no marker.
///
/// **Inverted 2026-08-10 by D2 Phase 2**, exactly as the previous version of
/// this paragraph instructed. `held_reader_forces_a_whole_graph_copy` is now
/// [`held_reader_copies_no_nodes`]: a held reader makes the writer fork to a
/// copy-on-write overlay (`storage/forked.rs`) and copies **zero** nodes. The
/// two arms around it — `unique_handle_copies_nothing` and
/// `dropping_the_reader_restores_in_place_mutation` — are unchanged and still
/// green, which is what keeps this a statement about *sharing* rather than
/// about `make_dir_graph_mut` having become unconditionally cheap.
///
/// The oracle itself had to be re-pointed in the same change, or it would have
/// gone on passing for the wrong reason: `impl Clone for GraphBackend` used to
/// bump `BACKEND_CLONE_NODES` by `node_count()` on *every* clone, including the
/// new shallow one. It now bumps only where node storage is genuinely
/// duplicated (`backend::note_nodes_copied`). A test that cannot tell the fix
/// from the defect is worse than no test.
///
/// What must **not** be relaxed here: these still pin behaviour in both
/// directions. If a later phase makes the fork copy nodes again, this file is
/// what says so.
#[cfg(test)]
mod held_reference_clone_tests {
    use super::*;
    use crate::graph::session::{execute_mut, ExecuteOptions};
    use crate::graph::storage::backend::{backend_clone_nodes, reset_backend_clone_count};
    use crate::graph::storage::GraphRead;
    use std::collections::HashMap;

    /// Small on purpose. The defect is O(V), so ten nodes prove the shape
    /// exactly as well as a million and keep the test in the millisecond range.
    const FIXTURE_NODES: usize = 10;

    /// `(slot, id, title)` for every live node, in scan order.
    ///
    /// Keyed by petgraph slot on purpose: content equality alone would pass a
    /// compaction that put every node back with the right values on the wrong
    /// index, which is precisely the failure `rollback.rs`'s slot-identity
    /// contract forbids and which every `DirGraph` index would then mis-key.
    fn snapshot_ids_and_titles(graph: &DirGraph) -> Vec<(usize, String, String)> {
        graph
            .graph
            .node_indices()
            .map(|idx| {
                let view = graph.graph.node_view(idx).expect("live node");
                (
                    idx.index(),
                    format!("{:?}", view.id()),
                    format!("{:?}", view.title()),
                )
            })
            .collect()
    }

    fn seeded_arc() -> Arc<DirGraph> {
        let mut graph = DirGraph::new();
        let params = HashMap::new();
        for i in 0..FIXTURE_NODES {
            execute_mut(
                &mut graph,
                &format!("CREATE (:Item {{id: {i}, name: 'item-{i}'}})"),
                &ExecuteOptions::eager(&params),
            )
            .expect("fixture node");
        }
        Arc::new(graph)
    }

    /// The baseline: a uniquely-owned handle mutates in place.
    ///
    /// Without this arm the guard below would pass on a build where *every*
    /// write copies the graph, and would still look like it was measuring
    /// something specific to holding a reader.
    #[test]
    fn unique_handle_copies_nothing() {
        let mut arc = seeded_arc();
        reset_backend_clone_count();
        let _ = make_dir_graph_mut(&mut arc);
        assert_eq!(
            backend_clone_nodes(),
            0,
            "a uniquely-owned Arc<DirGraph> must mutate in place"
        );
    }

    /// The fix, as an exact integer — the inversion of
    /// `held_reader_forces_a_whole_graph_copy`.
    ///
    /// Holding one extra `Arc` — what a live `ResultView`, a `freeze()`, a
    /// `Session` or an open transaction snapshot does — used to copy every
    /// node. It now forks to an overlay and copies none.
    #[test]
    fn held_reader_copies_no_nodes() {
        let mut arc = seeded_arc();
        let reader = Arc::clone(&arc);

        reset_backend_clone_count();
        let _ = make_dir_graph_mut(&mut arc);
        let copied = backend_clone_nodes();

        // Keep the reader alive across the write. Dropping it earlier would
        // make `Arc::get_mut` succeed and the test would measure
        // `unique_handle_copies_nothing` while appearing to measure this.
        assert_eq!(reader.graph.node_count(), FIXTURE_NODES);

        assert_eq!(
            copied, 0,
            "a live second Arc<DirGraph> must fork to a copy-on-write overlay, \
             not copy the graph; getting {FIXTURE_NODES} here means the fork \
             regressed to a deep clone (storage/forked.rs)"
        );
        assert!(
            arc.graph.is_forked(),
            "the writer's backend must be the overlay variant while the reader lives"
        );
    }

    /// The reader must still see its own pre-write graph, and the writer its
    /// post-write one, with the writer's edits landing nowhere the reader can
    /// see them.
    ///
    /// This is the semantic contract the whole programme exists to preserve.
    /// Before D2 it held for the *expensive* reason (the reader owned a private
    /// deep copy); it must now hold for the cheap one.
    #[test]
    fn a_held_reader_never_observes_the_writers_edits() {
        let mut arc = seeded_arc();
        let reader = Arc::clone(&arc);
        let before = snapshot_ids_and_titles(&reader);

        {
            let graph = make_dir_graph_mut(&mut arc);
            let params = HashMap::new();
            execute_mut(
                graph,
                "MATCH (n:Item {id: 3}) SET n.name = 'rewritten'",
                &ExecuteOptions::eager(&params),
            )
            .expect("write");
            execute_mut(
                graph,
                "CREATE (:Item {id: 999, name: 'appended'})",
                &ExecuteOptions::eager(&params),
            )
            .expect("append");
        }

        assert_eq!(
            snapshot_ids_and_titles(&reader),
            before,
            "the reader's graph must be byte-for-byte what it was before the \
             write; a difference here means the writer mutated the shared base, \
             which silently corrupts every holder of that snapshot"
        );
        assert_eq!(reader.graph.node_count(), FIXTURE_NODES);
        assert_eq!(arc.graph.node_count(), FIXTURE_NODES + 1);
    }

    /// The **id index** must answer for the graph you asked, not the one that
    /// wrote last — the D2 Phase 3 counterpart of
    /// `a_held_reader_never_observes_the_writers_edits`.
    ///
    /// `id_indices` is layered over a base the two graphs share
    /// (`storage/disk/id_index_layer.rs`), so this is where a delta leaking into
    /// the base would show up: the reader would resolve an id its snapshot has
    /// never seen, and `MATCH (n:Item {id: …})` would return a node that does
    /// not exist in that view.
    #[test]
    fn a_held_reader_resolves_ids_against_its_own_snapshot() {
        let mut arc = seeded_arc();
        let params = HashMap::new();

        // Warm the id index *before* the fork, so the lookups below take the
        // layered path rather than build-on-miss — a cold index would rebuild
        // from each graph's own state and pass this test for the wrong reason.
        {
            let graph = make_dir_graph_mut(&mut arc);
            execute_mut(
                graph,
                "MATCH (n:Item {id: 0}) RETURN n.id",
                &ExecuteOptions::eager(&params),
            )
            .expect("warm the id index");
        }
        let reader = Arc::clone(&arc);
        assert!(
            reader
                .id_indices
                .lookup("Item", &crate::datatypes::Value::Int64(0))
                .is_some(),
            "fixture must be id-indexed, or this test proves nothing"
        );

        {
            let graph = make_dir_graph_mut(&mut arc);
            execute_mut(
                graph,
                "CREATE (:Item {id: 4242, name: 'appended'})",
                &ExecuteOptions::eager(&params),
            )
            .expect("append");
        }

        let new_id = crate::datatypes::Value::Int64(4242);
        assert!(
            arc.id_indices.lookup("Item", &new_id).is_some(),
            "the writer must resolve the id it just created"
        );
        assert!(
            reader.id_indices.lookup("Item", &new_id).is_none(),
            "the reader's snapshot never saw this id; resolving it means the \
             writer's delta leaked into the shared base"
        );
        // ...and every pre-existing id still resolves identically in both.
        for id in 0..FIXTURE_NODES {
            let value = crate::datatypes::Value::Int64(id as i64);
            assert_eq!(
                reader.id_indices.lookup("Item", &value),
                arc.id_indices.lookup("Item", &value),
                "id {id} must resolve the same in both graphs"
            );
        }
    }

    /// The **type index** must scan the graph you asked, not the one that wrote
    /// last.
    ///
    /// `type_indices` buckets are stacks of shared levels
    /// (`storage/disk/type_index_layer.rs`), so this is where a `CREATE`'s
    /// append leaking into a shared level would show up: the reader's
    /// `MATCH (n:Item)` would return a node its snapshot has never seen, and —
    /// because the bucket drives the label scan directly — it would return a
    /// `NodeIndex` the reader's backend cannot even resolve.
    #[test]
    fn a_held_reader_scans_types_against_its_own_snapshot() {
        let mut arc = seeded_arc();
        let reader = Arc::clone(&arc);
        let before: Vec<usize> = reader
            .type_indices
            .get("Item")
            .expect("fixture is type-indexed")
            .iter()
            .map(|idx| idx.index())
            .collect();
        assert_eq!(before.len(), FIXTURE_NODES);

        {
            let graph = make_dir_graph_mut(&mut arc);
            let params = HashMap::new();
            execute_mut(
                graph,
                "CREATE (:Item {id: 7007, name: 'appended'})",
                &ExecuteOptions::eager(&params),
            )
            .expect("append");
        }

        let after: Vec<usize> = reader
            .type_indices
            .get("Item")
            .expect("the reader keeps its bucket")
            .iter()
            .map(|idx| idx.index())
            .collect();
        assert_eq!(
            after, before,
            "the reader's type bucket must be what it was before a write it \
             never asked for; a difference means the writer appended into a \
             level the reader shares"
        );
        assert_eq!(
            arc.type_indices.get("Item").map(|members| members.len()),
            Some(FIXTURE_NODES + 1),
            "the writer must see its own append"
        );
        // Every member the reader reports must still resolve in the reader's
        // own backend — the failure a leaked append produces is a dangling
        // NodeIndex, not merely an extra row.
        for idx in &after {
            assert!(
                reader
                    .graph
                    .node_weight(petgraph::graph::NodeIndex::new(*idx))
                    .is_some(),
                "the reader's bucket points at a node its backend does not have"
            );
        }
    }

    /// The **user index families** must answer for the graph you asked, not the
    /// one that wrote last.
    ///
    /// `property_indices` and `composite_indices` buckets live in shared,
    /// immutable levels (`dir_graph/index_layer.rs`), and both are consulted by
    /// `try_index_lookup` **unconditionally** — the matcher trusts the index
    /// rather than verifying against the graph. So a delta leaking into a shared
    /// level does not merely add a row: the reader's indexed `MATCH` returns a
    /// `NodeIndex` its own backend cannot resolve.
    ///
    /// Both directions are asserted, because they fail differently: the writer
    /// not seeing its own write is a *stale* index, the reader seeing the
    /// writer's is a *corrupt* one.
    #[test]
    fn a_held_reader_resolves_user_indexes_against_its_own_snapshot() {
        let mut arc = seeded_arc();
        let params = HashMap::new();
        {
            let graph = make_dir_graph_mut(&mut arc);
            // Give every fixture node the second composite component before the
            // index is built, so the index is over two live properties.
            execute_mut(
                graph,
                "MATCH (n:Item) SET n.qty = n.id",
                &ExecuteOptions::eager(&params),
            )
            .expect("seed qty");
            graph.create_index("Item", "name");
            graph.create_composite_index("Item", &["name", "qty"]);
        }
        let reader = Arc::clone(&arc);

        let existing = crate::datatypes::Value::String("item-3".to_string());
        let created = crate::datatypes::Value::String("appended".to_string());
        let composite_of = |value: &crate::datatypes::Value, qty: i64| {
            vec![value.clone(), crate::datatypes::Value::Int64(qty)]
        };
        assert!(
            reader
                .lookup_by_index("Item", "name", &existing)
                .is_some_and(|members| !members.is_empty()),
            "the fixture must be indexed, or this test proves nothing"
        );
        let before_existing = reader.lookup_by_index("Item", "name", &existing);

        {
            let graph = make_dir_graph_mut(&mut arc);
            execute_mut(
                graph,
                "CREATE (:Item {id: 4242, name: 'appended', qty: 7})",
                &ExecuteOptions::eager(&params),
            )
            .expect("append");
        }

        assert!(
            arc.lookup_by_index("Item", "name", &created).is_some(),
            "the writer must find the value it just indexed"
        );
        assert!(
            arc.lookup_by_composite_index(
                "Item",
                &["name".to_string(), "qty".to_string()],
                &composite_of(&created, 7)
            )
            .is_some(),
            "the writer's composite index must carry its own write"
        );
        assert_eq!(
            reader.lookup_by_index("Item", "name", &created),
            None,
            "the reader's snapshot never saw this value; finding it means the \
             writer's delta leaked into a shared level"
        );
        assert_eq!(
            reader.lookup_by_composite_index(
                "Item",
                &["name".to_string(), "qty".to_string()],
                &composite_of(&created, 7)
            ),
            None,
            "same for the composite index"
        );
        assert_eq!(
            reader.lookup_by_index("Item", "name", &existing),
            before_existing,
            "a value the reader already had must be unchanged, in bucket order"
        );
        // The other direction: the writer's own delta must not *hide* the
        // buckets it inherited. A read that stopped at the top level would
        // silently turn every pre-fork value into a miss, and `try_index_lookup`
        // trusts the index rather than falling back to a scan — so the rows
        // would simply disappear from an indexed `MATCH`.
        assert_eq!(
            arc.lookup_by_index("Item", "name", &existing),
            before_existing,
            "the writer must still resolve the values it inherited"
        );
        assert!(
            arc.lookup_by_composite_index(
                "Item",
                &["name".to_string(), "qty".to_string()],
                &composite_of(&existing, 3)
            )
            .is_some_and(|members| !members.is_empty()),
            "the writer's composite index must still carry the inherited tuples"
        );
        // Every member the reader's index reports must still resolve in the
        // reader's own backend — a leak shows up as a dangling NodeIndex, not
        // merely as an extra row.
        for idx in reader
            .lookup_by_index("Item", "name", &existing)
            .unwrap_or_default()
        {
            assert!(
                reader.graph.node_weight(idx).is_some(),
                "the reader's property index points at a node it does not have"
            );
        }
    }

    /// Compaction: once the reader drops, the next write folds the overlay back
    /// into the base and the backend returns to the flat representation — with
    /// every node still on the slot it had.
    #[test]
    fn dropping_the_reader_compacts_the_overlay_back_into_the_base() {
        let mut arc = seeded_arc();
        let reader = Arc::clone(&arc);
        let params = HashMap::new();

        {
            let graph = make_dir_graph_mut(&mut arc);
            execute_mut(
                graph,
                "CREATE (:Item {id: 999, name: 'appended'})",
                &ExecuteOptions::eager(&params),
            )
            .expect("append");
        }
        assert!(arc.graph.is_forked(), "precondition: the write forked");
        let forked_view = snapshot_ids_and_titles(&arc);

        drop(reader);
        reset_backend_clone_count();
        let _ = make_dir_graph_mut(&mut arc);

        assert!(
            !arc.graph.is_forked(),
            "the next write after the reader drops must collapse the overlay"
        );
        assert_eq!(
            backend_clone_nodes(),
            0,
            "compaction folds in place; it must not copy the graph"
        );
        assert_eq!(
            snapshot_ids_and_titles(&arc),
            forked_view,
            "compaction must preserve content AND petgraph slot identity — the \
             snapshot is keyed by slot, so a re-indexed node shows up here"
        );
    }

    /// Dropping the reader restores in-place mutation.
    ///
    /// This is what makes the guard above a statement about *sharing* rather
    /// than about `make_dir_graph_mut` being unconditionally expensive.
    #[test]
    fn dropping_the_reader_restores_in_place_mutation() {
        let mut arc = seeded_arc();
        let reader = Arc::clone(&arc);
        drop(reader);

        reset_backend_clone_count();
        let _ = make_dir_graph_mut(&mut arc);
        assert_eq!(
            backend_clone_nodes(),
            0,
            "once the extra Arc is gone the write must mutate in place again"
        );
    }
}