impactsense_parser/store/

in_memory.rs

use std::collections::{HashMap, HashSet, VecDeque};
use std::path::Path;

use crate::compress::{
    compressor_language_from_ir_string, CompressorClient, CompressorConfig,
};
use crate::graph::{
    class_body_spans_for_file, function_body_spans_for_file, property_body_spans_for_file,
};
use crate::scanner::FileScanConfig;
use crate::scanner_incremental::scan_and_parse_incremental_vector;
use crate::ir::{
    api_endpoint_key, external_api_key, module_key, ApiEndpointIr, BehaviourIr, CallbackIr,
    ClassIr, EdgeIr, EdgeKind, ExternalApiIr, FileIr, FunctionIr, ModuleIr, PropertyIr, ProjectIr,
};
use crate::schema::NodeLabel;

/// Reference to a symbol in the graph.
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct SymbolRef {
    pub label: String,
    pub key: String,
    pub name: Option<String>,
    pub path: Option<String>,
}

/// Symbols declared in a single file.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct FileSymbols {
    pub path: String,
    pub classes: Vec<SymbolRef>,
    pub functions: Vec<SymbolRef>,
    pub modules: Vec<SymbolRef>,
    pub api_endpoints: Vec<SymbolRef>,
}

/// Bounded impact analysis result.
#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct ImpactReport {
    pub symbol: String,
    pub depth: u32,
    pub callers: Vec<String>,
    pub affected_files: Vec<String>,
    pub truncated: bool,
}

#[derive(Debug, Clone, Copy)]
pub struct QueryLimits {
    pub max_depth: u32,
    pub max_results: usize,
}

impl Default for QueryLimits {
    fn default() -> Self {
        Self {
            max_depth: 2,
            max_results: 50,
        }
    }
}

/// How [`ExplainSymbolResult::source`] was obtained.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ExplainSourceOrigin {
    Decompressed,
    FileSpan,
    Unavailable,
}

#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct ExplainOptions {
    #[serde(default)]
    pub include_callers: bool,
    #[serde(default)]
    pub include_callees: bool,
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ExplainSymbolResult {
    pub symbol: SymbolRef,
    pub source: Option<String>,
    pub source_origin: ExplainSourceOrigin,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub callers: Vec<SymbolRef>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub callees: Vec<SymbolRef>,
}

#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
pub struct RefreshReport {
    pub cleanup_targets: usize,
    pub parse_targets: usize,
    pub nodes_merged: usize,
    pub edges_merged: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum NodeKind {
    File,
    Module,
    Class,
    Property,
    Function,
    Behaviour,
    Callback,
    ApiEndpoint,
    ExternalApi,
}

#[derive(Debug, Clone)]
struct NodeRecord {
    kind: NodeKind,
    key: String,
    label: String,
    name: Option<String>,
    path: Option<String>,
    language: Option<String>,
    code_bytes: Option<Vec<u8>>,
}

type AdjKey = (usize, EdgeKind);

/// In-memory dependency graph with indexed lookups for IDE/MCP queries.
#[derive(Debug, Default)]
pub struct InMemoryGraph {
    nodes: Vec<NodeRecord>,
    nodes_by_key: HashMap<(NodeKind, String), usize>,
    symbols_by_path: HashMap<String, HashSet<usize>>,
    by_simple_name: HashMap<String, Vec<usize>>,
    outgoing: HashMap<AdjKey, Vec<usize>>,
    incoming: HashMap<AdjKey, Vec<usize>>,
}

pub trait GraphStore {
    fn callers(&self, fqn: &str) -> Vec<SymbolRef>;
    fn callees(&self, fqn: &str) -> Vec<SymbolRef>;
    fn file_dependencies(&self, path: &str) -> Vec<String>;
    fn impact(&self, fqn: &str, limits: QueryLimits) -> ImpactReport;
    fn symbols_in_file(&self, path: &str) -> FileSymbols;
    fn find_symbol(&self, query: &str) -> Vec<SymbolRef>;
    fn node_count(&self) -> usize;
    fn edge_count(&self) -> usize;
}

impl InMemoryGraph {
    pub fn from_ir(ir: ProjectIr) -> Self {
        let mut g = Self::default();
        g.merge_ir(ir);
        g
    }

    pub fn merge_ir(&mut self, delta: ProjectIr) {
        for f in delta.files {
            self.upsert_file(f);
        }
        for m in delta.modules {
            self.upsert_module(m);
        }
        for c in delta.classes {
            self.upsert_class(c);
        }
        for p in delta.properties {
            self.upsert_property(p);
        }
        for f in delta.functions {
            self.upsert_function(f);
        }
        for b in delta.behaviours {
            self.upsert_behaviour(b);
        }
        for c in delta.callbacks {
            self.upsert_callback(c);
        }
        for a in delta.api_endpoints {
            self.upsert_api_endpoint(a);
        }
        for e in delta.external_apis {
            self.upsert_external_api(e);
        }
        for edge in delta.edges {
            self.add_edge(edge);
        }
    }

    pub fn remove_file(&mut self, path: &str) {
        let normalized = path.replace('\\', "/");
        let to_remove: Vec<usize> = self
            .nodes
            .iter()
            .enumerate()
            .filter(|(_, n)| n.path.as_deref() == Some(normalized.as_str()))
            .map(|(i, _)| i)
            .collect();

        if to_remove.is_empty() {
            return;
        }

        let remove_set: HashSet<usize> = to_remove.into_iter().collect();
        self.remove_nodes(&remove_set);
    }

    fn remove_nodes(&mut self, remove_set: &HashSet<usize>) {
        let mut remap: HashMap<usize, Option<usize>> = HashMap::new();
        let mut new_nodes: Vec<NodeRecord> = Vec::new();

        for (old_id, node) in self.nodes.iter().enumerate() {
            if remove_set.contains(&old_id) {
                remap.insert(old_id, None);
                self.nodes_by_key.remove(&(node.kind, node.key.clone()));
                if let Some(p) = &node.path {
                    if let Some(set) = self.symbols_by_path.get_mut(p) {
                        set.remove(&old_id);
                    }
                }
                if let Some(name) = &node.name {
                    if let Some(ids) = self.by_simple_name.get_mut(name) {
                        ids.retain(|id| *id != old_id);
                    }
                }
            } else {
                let new_id = new_nodes.len();
                remap.insert(old_id, Some(new_id));
                new_nodes.push(node.clone());
            }
        }

        self.nodes = new_nodes;
        self.rebuild_adjacency(&remap);
        self.rebuild_indexes();
    }

    fn rebuild_adjacency(&mut self, remap: &HashMap<usize, Option<usize>>) {
        let mut new_out: HashMap<AdjKey, Vec<usize>> = HashMap::new();
        let mut new_in: HashMap<AdjKey, Vec<usize>> = HashMap::new();

        for ((from, kind), targets) in &self.outgoing {
            let Some(&Some(new_from)) = remap.get(from) else {
                continue;
            };
            for to in targets {
                if let Some(Some(new_to)) = remap.get(to) {
                    new_out.entry((new_from, *kind)).or_default().push(*new_to);
                    new_in.entry((*new_to, *kind)).or_default().push(new_from);
                }
            }
        }
        self.outgoing = new_out;
        self.incoming = new_in;
    }

    fn rebuild_indexes(&mut self) {
        self.nodes_by_key.clear();
        self.symbols_by_path.clear();
        self.by_simple_name.clear();
        for (id, node) in self.nodes.iter().enumerate() {
            self.nodes_by_key.insert((node.kind, node.key.clone()), id);
            if let Some(p) = &node.path {
                self.symbols_by_path.entry(p.clone()).or_default().insert(id);
            }
            if let Some(name) = &node.name {
                self.by_simple_name.entry(name.clone()).or_default().push(id);
            }
        }
    }

    fn upsert_file(&mut self, f: FileIr) {
        let _id = self.ensure_node(
            NodeKind::File,
            f.path.clone(),
            NodeLabel::File.to_string(),
            None,
            Some(f.path),
        );
    }

    fn upsert_module(&mut self, m: ModuleIr) {
        let key = module_key(&m.name, &m.path);
        let id = self.ensure_node(
            NodeKind::Module,
            key,
            NodeLabel::Module.to_string(),
            Some(m.name),
            Some(m.path),
        );
        self.apply_symbol_metadata(id, Some(&m.language), m.code_bytes);
    }

    fn upsert_class(&mut self, c: ClassIr) {
        let id = self.ensure_node(
            NodeKind::Class,
            c.fqn.clone(),
            "Class".to_string(),
            Some(c.name),
            Some(c.path),
        );
        self.apply_symbol_metadata(id, Some(&c.language), c.code_bytes);
    }

    fn upsert_property(&mut self, p: PropertyIr) {
        let id = self.ensure_node(
            NodeKind::Property,
            p.fqn.clone(),
            "Property".to_string(),
            Some(p.name),
            Some(p.path),
        );
        self.apply_symbol_metadata(id, Some(&p.language), p.code_bytes);
    }

    fn upsert_function(&mut self, f: FunctionIr) {
        let id = self.ensure_node(
            NodeKind::Function,
            f.fqn.clone(),
            NodeLabel::Function.to_string(),
            Some(f.name),
            Some(f.path),
        );
        self.apply_symbol_metadata(id, Some(&f.language), f.code_bytes);
    }

    fn apply_symbol_metadata(
        &mut self,
        id: usize,
        language: Option<&str>,
        code_bytes: Option<Vec<u8>>,
    ) {
        let Some(n) = self.nodes.get_mut(id) else {
            return;
        };
        if let Some(lang) = language {
            n.language = Some(lang.to_string());
        }
        if let Some(bytes) = code_bytes {
            n.code_bytes = Some(bytes);
        }
    }

    fn upsert_behaviour(&mut self, b: BehaviourIr) {
        let _id = self.ensure_node(
            NodeKind::Behaviour,
            b.name.clone(),
            NodeLabel::Behaviour.to_string(),
            Some(b.name.clone()),
            b.path,
        );
    }

    fn upsert_callback(&mut self, c: CallbackIr) {
        let _id = self.ensure_node(
            NodeKind::Callback,
            c.fqn.clone(),
            NodeLabel::Callback.to_string(),
            Some(c.name),
            None,
        );
    }

    fn upsert_api_endpoint(&mut self, a: ApiEndpointIr) {
        let key = api_endpoint_key(&a.methods, &a.path);
        let _id = self.ensure_node(
            NodeKind::ApiEndpoint,
            key,
            NodeLabel::ApiEndpoint.to_string(),
            Some(a.path.clone()),
            None,
        );
    }

    fn upsert_external_api(&mut self, e: ExternalApiIr) {
        let key = if let (Some(base), Some(norm)) = (&e.base_url, &e.norm_path) {
            external_api_key(base, norm)
        } else {
            e.name.clone()
        };
        let _id = self.ensure_node(
            NodeKind::ExternalApi,
            key,
            NodeLabel::ExternalApi.to_string(),
            Some(e.name),
            None,
        );
    }

    fn ensure_node(
        &mut self,
        kind: NodeKind,
        key: String,
        label: String,
        name: Option<String>,
        path: Option<String>,
    ) -> usize {
        if let Some(&id) = self.nodes_by_key.get(&(kind, key.clone())) {
            if let Some(n) = self.nodes.get_mut(id) {
                if name.is_some() {
                    n.name = name.clone();
                }
                if let Some(new_path) = path.clone() {
                    if n.path.as_deref() != Some(new_path.as_str()) {
                        if let Some(old_path) = n.path.take() {
                            if let Some(set) = self.symbols_by_path.get_mut(&old_path) {
                                set.remove(&id);
                            }
                        }
                        n.path = Some(new_path.clone());
                        self.symbols_by_path
                            .entry(new_path)
                            .or_default()
                            .insert(id);
                    }
                }
            }
            return id;
        }
        let id = self.nodes.len();
        let record = NodeRecord {
            kind,
            key: key.clone(),
            label,
            name: name.clone(),
            path: path.clone(),
            language: None,
            code_bytes: None,
        };
        self.nodes.push(record);
        self.nodes_by_key.insert((kind, key), id);
        if let Some(p) = path {
            self.symbols_by_path.entry(p).or_default().insert(id);
        }
        if let Some(n) = name {
            self.by_simple_name.entry(n).or_default().push(id);
        }
        id
    }

    fn add_edge(&mut self, edge: EdgeIr) {
        let Some(from_id) = self.lookup_id(&edge.from_label, &edge.from_key) else {
            return;
        };
        let Some(to_id) = self.lookup_id(&edge.to_label, &edge.to_key) else {
            return;
        };
        let out_list = self.outgoing.entry((from_id, edge.kind)).or_default();
        if !out_list.contains(&to_id) {
            out_list.push(to_id);
        }
        let in_list = self.incoming.entry((to_id, edge.kind)).or_default();
        if !in_list.contains(&from_id) {
            in_list.push(from_id);
        }
    }

    fn lookup_id(&self, label: &str, key: &str) -> Option<usize> {
        let kind = node_kind_from_label(label)?;
        self.nodes_by_key.get(&(kind, key.to_string())).copied()
    }

    fn function_id(&self, fqn: &str) -> Option<usize> {
        if let Some(&id) = self
            .nodes_by_key
            .get(&(NodeKind::Function, fqn.to_string()))
        {
            return Some(id);
        }
        let suffix = format!("::{fqn}");
        let mut suffix_matches = self
            .nodes
            .iter()
            .enumerate()
            .filter(|(_, n)| n.kind == NodeKind::Function && n.key.ends_with(&suffix));
        if let Some((id, _)) = suffix_matches.next() {
            if suffix_matches.next().is_none() {
                return Some(id);
            }
            return None;
        }
        if let Some(ids) = self.by_simple_name.get(fqn) {
            let fn_ids: Vec<usize> = ids
                .iter()
                .copied()
                .filter(|id| self.nodes.get(*id).is_some_and(|n| n.kind == NodeKind::Function))
                .collect();
            if fn_ids.len() == 1 {
                return Some(fn_ids[0]);
            }
        }
        None
    }

    fn node_to_symbol(&self, id: usize) -> Option<SymbolRef> {
        self.nodes.get(id).map(|n| SymbolRef {
            label: n.label.clone(),
            key: n.key.clone(),
            name: n.name.clone(),
            path: n.path.clone(),
        })
    }

    #[cfg(test)]
    fn code_bytes_for_symbol(&self, query: &str) -> Option<Vec<u8>> {
        let id = self.resolve_symbol(query).ok()?;
        self.nodes[id].code_bytes.clone()
    }

    #[cfg(test)]
    fn language_for_symbol(&self, query: &str) -> Option<String> {
        let id = self.resolve_symbol(query).ok()?;
        self.nodes[id].language.clone()
    }

    /// Resolve a symbol query to a single graph node id, if unambiguous.
    fn resolve_symbol(&self, query: &str) -> Result<usize, String> {
        let q = query.trim();
        if q.is_empty() {
            return Err("empty symbol query".into());
        }

        for kind in [
            NodeKind::Function,
            NodeKind::Class,
            NodeKind::Property,
            NodeKind::Module,
        ] {
            if let Some(&id) = self.nodes_by_key.get(&(kind, q.to_string())) {
                return Ok(id);
            }
        }

        if let Some(id) = self.function_id(q) {
            return Ok(id);
        }

        let suffix = format!("::{q}");
        let mut class_matches = self
            .nodes
            .iter()
            .enumerate()
            .filter(|(_, n)| n.kind == NodeKind::Class && n.key.ends_with(&suffix));
        if let Some((id, _)) = class_matches.next() {
            if class_matches.next().is_none() {
                return Ok(id);
            }
            return Err(format!("ambiguous class match for `{q}`"));
        }

        if let Some(ids) = self.by_simple_name.get(q) {
            let symbol_ids: Vec<usize> = ids
                .iter()
                .copied()
                .filter(|id| {
                    self.nodes.get(*id).is_some_and(|n| {
                        matches!(
                            n.kind,
                            NodeKind::Function | NodeKind::Class | NodeKind::Property
                        )
                    })
                })
                .collect();
            if symbol_ids.len() == 1 {
                return Ok(symbol_ids[0]);
            }
            if symbol_ids.len() > 1 {
                return Err(format!("ambiguous symbol match for `{q}`"));
            }
        }

        if q.contains('@') {
            if let Some(&id) = self.nodes_by_key.get(&(NodeKind::Module, q.to_string())) {
                return Ok(id);
            }
        }

        Err(format!("symbol not found: `{q}`"))
    }

    fn slice_source_span(source: &str, span: (usize, usize)) -> Option<String> {
        let (lo, hi) = span;
        let lo = lo.min(source.len());
        let hi = hi.min(source.len());
        if lo >= hi {
            return None;
        }
        Some(source[lo..hi].to_string())
    }

    fn span_for_node_in_file(
        node: &NodeRecord,
        file: &crate::scanner::ParsedFile,
        source: &str,
    ) -> Option<(usize, usize)> {
        let path = node.path.as_deref()?;
        match node.kind {
            NodeKind::Function => {
                let spans = function_body_spans_for_file(file, path, source);
                spans.get(&node.key).copied()
            }
            NodeKind::Class => {
                let spans = class_body_spans_for_file(file, source);
                spans.get(&node.key).copied()
            }
            NodeKind::Property => {
                let spans = property_body_spans_for_file(file, source);
                spans.get(&node.key).copied()
            }
            NodeKind::Module => Some((0, source.len())),
            _ => None,
        }
    }

    async fn source_from_file_span(
        &self,
        node: &NodeRecord,
        root: &Path,
    ) -> Result<String, String> {
        let rel_path = node
            .path
            .as_deref()
            .ok_or_else(|| "symbol has no file path".to_string())?;
        let full_path = root.join(rel_path);
        let source = std::fs::read_to_string(&full_path)
            .map_err(|e| format!("failed to read {}: {e}", full_path.display()))?;

        let config = FileScanConfig::new(root);
        let paths = vec![std::path::PathBuf::from(rel_path)];
        let files = scan_and_parse_incremental_vector(&config, &paths)
            .map_err(|e| format!("failed to parse {rel_path}: {e}"))?;
        let file = files
            .first()
            .ok_or_else(|| format!("no parse result for {rel_path}"))?;

        let span = Self::span_for_node_in_file(node, file, &source).ok_or_else(|| {
            format!(
                "no AST span for {} `{}` in {rel_path}",
                node.label, node.key
            )
        })?;
        Self::slice_source_span(&source, span)
            .ok_or_else(|| format!("empty source span for `{}`", node.key))
    }

    /// Return decompressed or file-sliced implementation source for a symbol.
    pub async fn explain_symbol_logic(
        &self,
        fqn: &str,
        root: &Path,
        compressor: &CompressorConfig,
        opts: ExplainOptions,
    ) -> ExplainSymbolResult {
        let empty_symbol = SymbolRef {
            label: String::new(),
            key: fqn.to_string(),
            name: None,
            path: None,
        };

        let node_id = match self.resolve_symbol(fqn) {
            Ok(id) => id,
            Err(e) => {
                return ExplainSymbolResult {
                    symbol: empty_symbol,
                    source: None,
                    source_origin: ExplainSourceOrigin::Unavailable,
                    error: Some(e),
                    callers: Vec::new(),
                    callees: Vec::new(),
                };
            }
        };

        let node = &self.nodes[node_id];
        let symbol = self.node_to_symbol(node_id).unwrap_or(empty_symbol);

        let mut result = ExplainSymbolResult {
            symbol,
            source: None,
            source_origin: ExplainSourceOrigin::Unavailable,
            error: None,
            callers: Vec::new(),
            callees: Vec::new(),
        };

        if node.kind == NodeKind::Function {
            if opts.include_callers {
                result.callers = self.callers(&node.key);
            }
            if opts.include_callees {
                result.callees = self.callees(&node.key);
            }
        }

        if let (Some(blob), Some(lang_str)) = (&node.code_bytes, &node.language) {
            if compressor.enabled {
                if let Some(api_lang) = compressor_language_from_ir_string(lang_str) {
                    if let Ok(client) = CompressorClient::from_config(compressor) {
                        match client.decompress_code(blob, api_lang).await {
                            Ok(code) => {
                                result.source = Some(code);
                                result.source_origin = ExplainSourceOrigin::Decompressed;
                                return result;
                            }
                            Err(e) => {
                                result.error = Some(format!("decompress failed: {e}"));
                            }
                        }
                    }
                }
            }
        }

        match self.source_from_file_span(node, root).await {
            Ok(source) => {
                result.source = Some(source);
                result.source_origin = ExplainSourceOrigin::FileSpan;
                result.error = None;
            }
            Err(e) => {
                if result.error.is_none() {
                    result.error = Some(e);
                } else {
                    result.error = Some(format!("{}; {}", result.error.unwrap(), e));
                }
            }
        }

        result
    }
}

impl GraphStore for InMemoryGraph {
    fn callers(&self, fqn: &str) -> Vec<SymbolRef> {
        let Some(fn_id) = self.function_id(fqn) else {
            return Vec::new();
        };
        self.incoming
            .get(&(fn_id, EdgeKind::CallsFunction))
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| self.node_to_symbol(*id))
                    .collect()
            })
            .unwrap_or_default()
    }

    fn callees(&self, fqn: &str) -> Vec<SymbolRef> {
        let Some(fn_id) = self.function_id(fqn) else {
            return Vec::new();
        };
        self.outgoing
            .get(&(fn_id, EdgeKind::CallsFunction))
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| self.node_to_symbol(*id))
                    .collect()
            })
            .unwrap_or_default()
    }

    fn file_dependencies(&self, path: &str) -> Vec<String> {
        let normalized = path.replace('\\', "/");
        let Some(&file_id) = self.nodes_by_key.get(&(NodeKind::File, normalized.clone())) else {
            return Vec::new();
        };
        self.outgoing
            .get(&(file_id, EdgeKind::DependsOnFile))
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| self.nodes.get(*id).and_then(|n| n.path.clone()))
                    .collect()
            })
            .unwrap_or_default()
    }

    fn impact(&self, fqn: &str, limits: QueryLimits) -> ImpactReport {
        let mut report = ImpactReport {
            symbol: fqn.to_string(),
            depth: limits.max_depth,
            ..Default::default()
        };
        let Some(start) = self.function_id(fqn) else {
            return report;
        };

        let mut visited: HashSet<usize> = HashSet::new();
        let mut queue: VecDeque<(usize, u32)> = VecDeque::new();
        queue.push_back((start, 0));
        visited.insert(start);

        while let Some((node_id, depth)) = queue.pop_front() {
            if depth >= limits.max_depth {
                continue;
            }
            if let Some(callers) = self.incoming.get(&(node_id, EdgeKind::CallsFunction)) {
                for caller_id in callers {
                    if visited.contains(caller_id) {
                        continue;
                    }
                    if report.callers.len() >= limits.max_results {
                        report.truncated = true;
                        return report;
                    }
                    visited.insert(*caller_id);
                    if let Some(sym) = self.node_to_symbol(*caller_id) {
                        if sym.label == NodeLabel::Function.to_string() {
                            report.callers.push(sym.key.clone());
                        }
                        if let Some(p) = sym.path {
                            if !report.affected_files.contains(&p) {
                                report.affected_files.push(p);
                            }
                        }
                    }
                    queue.push_back((*caller_id, depth + 1));
                }
            }
        }
        report
    }

    fn symbols_in_file(&self, path: &str) -> FileSymbols {
        let normalized = path.replace('\\', "/");
        let mut out = FileSymbols {
            path: normalized.clone(),
            ..Default::default()
        };
        let Some(ids) = self.symbols_by_path.get(&normalized) else {
            return out;
        };
        for id in ids {
            let Some(n) = self.nodes.get(*id) else {
                continue;
            };
            let sym = SymbolRef {
                label: n.label.clone(),
                key: n.key.clone(),
                name: n.name.clone(),
                path: n.path.clone(),
            };
            match n.kind {
                NodeKind::Class => out.classes.push(sym),
                NodeKind::Function => out.functions.push(sym),
                NodeKind::Module => out.modules.push(sym),
                NodeKind::ApiEndpoint => out.api_endpoints.push(sym),
                _ => {}
            }
        }
        out
    }

    fn find_symbol(&self, query: &str) -> Vec<SymbolRef> {
        let q = query.trim();
        if q.is_empty() {
            return Vec::new();
        }

        if let Some(&id) = self.nodes_by_key.get(&(NodeKind::Function, q.to_string())) {
            return vec![self.node_to_symbol(id).unwrap()];
        }
        if let Some(&id) = self.nodes_by_key.get(&(NodeKind::Class, q.to_string())) {
            return vec![self.node_to_symbol(id).unwrap()];
        }

        let q_lower = q.to_lowercase();
        let mut results = Vec::new();
        for (id, node) in self.nodes.iter().enumerate() {
            if node.key.to_lowercase().contains(&q_lower)
                || node
                    .name
                    .as_ref()
                    .is_some_and(|n| n.to_lowercase().contains(&q_lower))
            {
                if let Some(sym) = self.node_to_symbol(id) {
                    results.push(sym);
                }
            }
        }
        results.sort_by(|a, b| a.key.cmp(&b.key));
        results.truncate(50);
        results
    }

    fn node_count(&self) -> usize {
        self.nodes.len()
    }

    fn edge_count(&self) -> usize {
        self.outgoing.values().map(|v| v.len()).sum()
    }
}

fn node_kind_from_label(label: &str) -> Option<NodeKind> {
    match label {
        "File" => Some(NodeKind::File),
        "Module" => Some(NodeKind::Module),
        "Class" => Some(NodeKind::Class),
        "Property" => Some(NodeKind::Property),
        "Function" => Some(NodeKind::Function),
        "Behaviour" => Some(NodeKind::Behaviour),
        "Callback" => Some(NodeKind::Callback),
        "ApiEndpoint" => Some(NodeKind::ApiEndpoint),
        "ExternalApi" => Some(NodeKind::ExternalApi),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::{EdgeKind, ProjectIr};

    fn sample_ir() -> ProjectIr {
        let mut ir = ProjectIr::empty();
        ir.files.push(FileIr {
            path: "src/a.rs".into(),
            language: "rust".into(),
            framework: None,
            project_name: None,
        });
        ir.functions.push(FunctionIr {
            name: "main".into(),
            fqn: "src/a.rs::main".into(),
            path: "src/a.rs".into(),
            language: "rust".into(),
            framework: None,
            project_name: None,
            arity: None,
            return_type: None,
            param_count: None,
            param_types: vec![],
            code_bytes: None,
        });
        ir.functions.push(FunctionIr {
            name: "helper".into(),
            fqn: "src/a.rs::helper".into(),
            path: "src/a.rs".into(),
            language: "rust".into(),
            framework: None,
            project_name: None,
            arity: None,
            return_type: None,
            param_count: None,
            param_types: vec![],
            code_bytes: None,
        });
        ir.edges.push(EdgeIr {
            kind: EdgeKind::DeclaresFunction,
            from_label: "File".into(),
            from_key: "src/a.rs".into(),
            to_label: "Function".into(),
            to_key: "src/a.rs::main".into(),
        });
        ir.edges.push(EdgeIr {
            kind: EdgeKind::CallsFunction,
            from_label: "Function".into(),
            from_key: "src/a.rs::main".into(),
            to_label: "Function".into(),
            to_key: "src/a.rs::helper".into(),
        });
        ir
    }

    #[test]
    fn in_memory_graph_queries_callers_and_callees() {
        let graph = InMemoryGraph::from_ir(sample_ir());
        let callees = graph.callees("src/a.rs::main");
        assert_eq!(callees.len(), 1);
        assert_eq!(callees[0].key, "src/a.rs::helper");
        let callers = graph.callers("src/a.rs::helper");
        assert_eq!(callers.len(), 1);
        assert_eq!(callers[0].key, "src/a.rs::main");
    }

    #[test]
    fn in_memory_graph_callers_by_short_name_when_unique() {
        let graph = InMemoryGraph::from_ir(sample_ir());
        let callers = graph.callers("helper");
        assert_eq!(callers.len(), 1);
        assert_eq!(callers[0].key, "src/a.rs::main");
    }

    #[test]
    fn in_memory_graph_callees_by_short_name_when_unique() {
        let graph = InMemoryGraph::from_ir(sample_ir());
        let callees = graph.callees("main");
        assert_eq!(callees.len(), 1);
        assert_eq!(callees[0].key, "src/a.rs::helper");
    }

    #[test]
    fn ensure_node_reindexes_symbols_by_path_on_path_change() {
        fn fn_ir(path: &str) -> FunctionIr {
            FunctionIr {
                name: "dup".into(),
                fqn: "dup".into(),
                path: path.into(),
                language: "rust".into(),
                framework: None,
                project_name: None,
                arity: None,
                return_type: None,
                param_count: None,
                param_types: vec![],
                code_bytes: None,
            }
        }
        let mut graph = InMemoryGraph::default();
        let mut ir = ProjectIr::empty();
        ir.functions.push(fn_ir("src/a.rs"));
        graph.merge_ir(ir);
        assert_eq!(graph.symbols_in_file("src/a.rs").functions.len(), 1);
        let mut ir2 = ProjectIr::empty();
        ir2.functions.push(fn_ir("src/b.rs"));
        graph.merge_ir(ir2);
        assert!(graph.symbols_in_file("src/a.rs").functions.is_empty());
        assert_eq!(graph.symbols_in_file("src/b.rs").functions.len(), 1);
    }

    #[test]
    fn remove_file_strips_symbols() {
        let mut graph = InMemoryGraph::from_ir(sample_ir());
        graph.remove_file("src/a.rs");
        assert!(graph.find_symbol("main").is_empty());
    }

    #[test]
    fn merge_ir_preserves_code_bytes_on_function() {
        let mut ir = sample_ir();
        ir.functions[0].code_bytes = Some(vec![1, 2, 3]);
        let graph = InMemoryGraph::from_ir(ir);
        assert_eq!(
            graph.code_bytes_for_symbol("src/a.rs::main").as_deref(),
            Some(&[1, 2, 3][..])
        );
        assert_eq!(graph.language_for_symbol("src/a.rs::main").as_deref(), Some("rust"));
    }

    #[test]
    fn merge_ir_coalesces_code_bytes_when_delta_omits_blob() {
        let mut ir = sample_ir();
        ir.functions[0].code_bytes = Some(vec![9, 8, 7]);
        let mut graph = InMemoryGraph::from_ir(ir);
        let mut ir2 = sample_ir();
        ir2.functions[0].code_bytes = None;
        graph.merge_ir(ir2);
        assert_eq!(
            graph.code_bytes_for_symbol("src/a.rs::main").as_deref(),
            Some(&[9, 8, 7][..])
        );
    }

    #[tokio::test]
    async fn explain_symbol_logic_file_span_fallback() {
        let dir = tempfile::tempdir().unwrap();
        let root = dir.path();
        std::fs::create_dir_all(root.join("src")).unwrap();
        std::fs::write(
            root.join("src/a.rs"),
            "fn main() {\n    helper();\n}\nfn helper() {}\n",
        )
        .unwrap();

        let mut ir = ProjectIr::empty();
        ir.files.push(FileIr {
            path: "src/a.rs".into(),
            language: "rust".into(),
            framework: None,
            project_name: None,
        });
        ir.functions.push(FunctionIr {
            name: "main".into(),
            fqn: "src/a.rs::main".into(),
            path: "src/a.rs".into(),
            language: "rust".into(),
            framework: None,
            project_name: None,
            arity: None,
            return_type: None,
            param_count: None,
            param_types: vec![],
            code_bytes: None,
        });

        let graph = InMemoryGraph::from_ir(ir);
        let compressor = CompressorConfig {
            enabled: false,
            ..Default::default()
        };
        let result = graph
            .explain_symbol_logic(
                "src/a.rs::main",
                root,
                &compressor,
                ExplainOptions::default(),
            )
            .await;
        assert_eq!(result.source_origin, ExplainSourceOrigin::FileSpan);
        let source = result.source.expect("source");
        assert!(
            source.contains("helper()"),
            "expected function body span, got: {source:?}"
        );
    }

}