codesynapse-core 0.1.3

Core graph extraction and analysis engine for codesynapse — AST parsing, semantic embeddings, BM25 index
Documentation
use crate::types::{Edge, GraphData, Node};
use std::collections::{HashMap, HashSet, VecDeque};

pub const DEFAULT_AFFECTED_RELATIONS: &[&str] = &[
    "calls",
    "references",
    "imports",
    "imports_from",
    "re_exports",
    "inherits",
    "extends",
    "implements",
    "uses",
    "mixes_in",
    "embeds",
];

#[derive(Debug, Clone, PartialEq)]
pub struct AffectedHit {
    pub node_id: String,
    pub depth: usize,
    pub via_relation: String,
}

pub fn resolve_seed(graph: &GraphData, query: &str) -> Option<String> {
    let node_ids: HashSet<&str> = graph.nodes.iter().map(|n| n.id.as_str()).collect();
    if node_ids.contains(query) {
        return Some(query.to_string());
    }
    let query_lower = query.to_lowercase();
    let exact_label: Vec<&str> = graph
        .nodes
        .iter()
        .filter(|n| n.label.to_lowercase() == query_lower)
        .map(|n| n.id.as_str())
        .collect();
    if exact_label.len() == 1 {
        return Some(exact_label[0].to_string());
    }
    let exact_source: Vec<&str> = graph
        .nodes
        .iter()
        .filter(|n| n.source_file.to_lowercase() == query_lower)
        .map(|n| n.id.as_str())
        .collect();
    if exact_source.len() == 1 {
        return Some(exact_source[0].to_string());
    }
    let contains: Vec<&str> = graph
        .nodes
        .iter()
        .filter(|n| n.label.to_lowercase().contains(&query_lower))
        .map(|n| n.id.as_str())
        .collect();
    if contains.len() == 1 {
        return Some(contains[0].to_string());
    }
    None
}

pub fn affected_nodes(
    graph: &GraphData,
    seed: &str,
    relations: &[&str],
    depth: usize,
) -> Vec<AffectedHit> {
    let relation_set: HashSet<&str> = relations.iter().copied().collect();

    let mut reverse_adj: HashMap<&str, Vec<(&str, &str)>> = HashMap::new();
    for edge in &graph.edges {
        reverse_adj
            .entry(edge.target.as_str())
            .or_default()
            .push((edge.source.as_str(), edge.relation.as_str()));
    }

    let mut seen: HashSet<&str> = HashSet::new();
    seen.insert(seed);
    let mut queue: VecDeque<(&str, usize)> = VecDeque::new();
    queue.push_back((seed, 0));
    let mut hits: Vec<AffectedHit> = Vec::new();

    while let Some((current, current_depth)) = queue.pop_front() {
        if current_depth >= depth {
            continue;
        }
        if let Some(incoming) = reverse_adj.get(current) {
            for &(source, relation) in incoming {
                if !relation_set.contains(relation) {
                    continue;
                }
                if seen.contains(source) {
                    continue;
                }
                seen.insert(source);
                hits.push(AffectedHit {
                    node_id: source.to_string(),
                    depth: current_depth + 1,
                    via_relation: relation.to_string(),
                });
                queue.push_back((source, current_depth + 1));
            }
        }
    }
    hits
}

pub fn format_affected(graph: &GraphData, query: &str, relations: &[&str], depth: usize) -> String {
    let node_map: HashMap<&str, &Node> = graph.nodes.iter().map(|n| (n.id.as_str(), n)).collect();

    let seed = match resolve_seed(graph, query) {
        Some(s) => s,
        None => return format!("No unique node match for {}", query),
    };

    let hits = affected_nodes(graph, &seed, relations, depth);
    let seed_label = node_map
        .get(seed.as_str())
        .map(|n| n.label.as_str())
        .unwrap_or(&seed);

    let mut lines = vec![
        format!("Affected nodes for {}", seed_label),
        format!("Relations: {}", relations.join(", ")),
        format!("Depth: {}", depth),
    ];

    if hits.is_empty() {
        lines.push("No affected nodes found.".to_string());
        return lines.join("\n");
    }

    for hit in &hits {
        let node = node_map.get(hit.node_id.as_str());
        let label = node.map(|n| n.label.as_str()).unwrap_or(&hit.node_id);
        let location = node
            .map(|n| {
                if let Some(ref loc) = n.source_location {
                    format!("{}:{}", n.source_file, loc)
                } else {
                    n.source_file.clone()
                }
            })
            .unwrap_or_else(|| "-".to_string());
        lines.push(format!("- {} [{}] {}", label, hit.via_relation, location));
    }
    lines.join("\n")
}

pub fn load_graph_json(path: &std::path::Path) -> crate::error::Result<GraphData> {
    let text = std::fs::read_to_string(path)?;
    let v: serde_json::Value =
        serde_json::from_str(&text).map_err(crate::error::CodeSynapseError::Serialization)?;

    let nodes: Vec<Node> = v
        .get("nodes")
        .and_then(|n| n.as_array())
        .map(|arr| {
            arr.iter()
                .filter_map(|item| serde_json::from_value(item.clone()).ok())
                .collect()
        })
        .unwrap_or_default();

    let edge_arr = v
        .get("edges")
        .or_else(|| v.get("links"))
        .and_then(|e| e.as_array());

    let edges: Vec<Edge> = edge_arr
        .map(|arr| {
            arr.iter()
                .filter_map(|item| {
                    let obj = item.as_object()?;
                    let source = obj.get("source")?.as_str()?.to_string();
                    let target = obj.get("target")?.as_str()?.to_string();
                    let relation = obj
                        .get("relation")
                        .and_then(|v| v.as_str())
                        .unwrap_or("")
                        .to_string();
                    let confidence = obj
                        .get("confidence")
                        .and_then(|v| v.as_str())
                        .unwrap_or("EXTRACTED")
                        .to_string();
                    let source_file = obj
                        .get("source_file")
                        .and_then(|v| v.as_str())
                        .map(|s| s.to_string());
                    let weight = obj.get("weight").and_then(|v| v.as_f64()).unwrap_or(1.0);
                    Some(Edge {
                        source,
                        target,
                        relation,
                        confidence,
                        source_file,
                        weight,
                        context: None,
                    })
                })
                .collect()
        })
        .unwrap_or_default();

    Ok(GraphData {
        nodes,
        edges,
        hyperedges: None,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Edge, Node};
    use std::collections::HashMap;

    fn make_node(id: &str, label: &str, source_file: &str) -> Node {
        Node {
            id: id.to_string(),
            label: label.to_string(),
            file_type: "code".to_string(),
            source_file: source_file.to_string(),
            source_location: None,
            community: None,
            rationale: None,
            docstring: None,
            metadata: HashMap::new(),
        }
    }

    fn make_edge(source: &str, target: &str, relation: &str) -> Edge {
        Edge {
            source: source.to_string(),
            target: target.to_string(),
            relation: relation.to_string(),
            confidence: "EXTRACTED".to_string(),
            source_file: None,
            weight: 1.0,
            context: None,
        }
    }

    fn test_graph() -> GraphData {
        GraphData {
            nodes: vec![
                make_node("target", "Foo", "pkg/foo.py"),
                make_node("caller", "X()", "app.py"),
                make_node("barrel", "__init__.py", "pkg/__init__.py"),
                make_node("consumer", "app.py", "app.py"),
            ],
            edges: vec![
                make_edge("caller", "target", "calls"),
                make_edge("barrel", "target", "re_exports"),
                make_edge("consumer", "target", "imports"),
            ],
            hyperedges: None,
        }
    }

    #[test]
    fn test_resolve_seed_by_id() {
        let g = test_graph();
        assert_eq!(resolve_seed(&g, "target"), Some("target".to_string()));
    }

    #[test]
    fn test_resolve_seed_by_label() {
        let g = test_graph();
        assert_eq!(resolve_seed(&g, "Foo"), Some("target".to_string()));
    }

    #[test]
    fn test_resolve_seed_case_insensitive() {
        let g = test_graph();
        assert_eq!(resolve_seed(&g, "foo"), Some("target".to_string()));
    }

    #[test]
    fn test_resolve_seed_not_found() {
        let g = test_graph();
        assert_eq!(resolve_seed(&g, "nonexistent"), None);
    }

    #[test]
    fn test_resolve_seed_ambiguous_returns_none() {
        let g = GraphData {
            nodes: vec![
                make_node("n1", "Foo", "a.py"),
                make_node("n2", "Foo", "b.py"),
            ],
            edges: vec![],
            hyperedges: None,
        };
        assert_eq!(resolve_seed(&g, "Foo"), None);
    }

    #[test]
    fn test_affected_nodes_all_relations() {
        let g = test_graph();
        let hits = affected_nodes(&g, "target", DEFAULT_AFFECTED_RELATIONS, 2);
        let ids: HashSet<&str> = hits.iter().map(|h| h.node_id.as_str()).collect();
        assert!(ids.contains("caller"));
        assert!(ids.contains("barrel"));
        assert!(ids.contains("consumer"));
    }

    #[test]
    fn test_affected_nodes_relation_filter() {
        let g = test_graph();
        let hits = affected_nodes(&g, "target", &["calls"], 2);
        let ids: HashSet<&str> = hits.iter().map(|h| h.node_id.as_str()).collect();
        assert!(ids.contains("caller"));
        assert!(!ids.contains("barrel"));
        assert!(!ids.contains("consumer"));
    }

    #[test]
    fn test_affected_nodes_depth_zero_returns_empty() {
        let g = test_graph();
        let hits = affected_nodes(&g, "target", DEFAULT_AFFECTED_RELATIONS, 0);
        assert!(hits.is_empty());
    }

    #[test]
    fn test_affected_nodes_depth_propagation() {
        let g = GraphData {
            nodes: vec![
                make_node("a", "A", "a.py"),
                make_node("b", "B", "b.py"),
                make_node("c", "C", "c.py"),
            ],
            edges: vec![make_edge("b", "a", "calls"), make_edge("c", "b", "calls")],
            hyperedges: None,
        };
        let hits_d1 = affected_nodes(&g, "a", &["calls"], 1);
        let hits_d2 = affected_nodes(&g, "a", &["calls"], 2);
        assert_eq!(hits_d1.len(), 1);
        assert_eq!(hits_d2.len(), 2);
    }

    #[test]
    fn test_format_affected_contains_expected_output() {
        let g = test_graph();
        let out = format_affected(&g, "Foo", DEFAULT_AFFECTED_RELATIONS, 2);
        assert!(
            out.contains("Affected nodes for Foo"),
            "missing header: {out}"
        );
        assert!(out.contains("X()"), "missing X(): {out}");
        assert!(out.contains("calls"), "missing calls: {out}");
        assert!(out.contains("__init__.py"), "missing __init__: {out}");
        assert!(out.contains("re_exports"), "missing re_exports: {out}");
        assert!(out.contains("imports"), "missing imports: {out}");
    }

    #[test]
    fn test_format_affected_relation_filter() {
        let g = test_graph();
        let out = format_affected(&g, "Foo", &["calls"], 2);
        assert!(
            out.contains("Relations: calls"),
            "missing relations header: {out}"
        );
        assert!(out.contains("X()"), "missing X(): {out}");
        assert!(!out.contains("__init__.py"), "unexpected __init__: {out}");
    }

    #[test]
    fn test_format_affected_not_found() {
        let g = test_graph();
        let out = format_affected(&g, "nonexistent", DEFAULT_AFFECTED_RELATIONS, 2);
        assert!(out.contains("No unique node match"));
    }

    #[test]
    fn test_format_affected_no_hits() {
        let g = GraphData {
            nodes: vec![make_node("isolated", "Isolated", "x.py")],
            edges: vec![],
            hyperedges: None,
        };
        let out = format_affected(&g, "isolated", DEFAULT_AFFECTED_RELATIONS, 2);
        assert!(out.contains("No affected nodes found."));
    }
}