use crate::autolink::is_auto_relation;
use crate::error::{BrainError, Result};
use crate::id::node_id_from_rel_path;
use crate::storage::Database;
use crate::symbols::{parse_symbol_path, resolve_symbol_ref};
use crate::types::{Edge, Node, NodeType};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet, VecDeque};
use std::path::Path;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum GraphDirection {
Out,
In,
#[default]
Both,
}
impl GraphDirection {
pub fn parse(s: &str) -> Option<Self> {
match s.trim().to_ascii_lowercase().as_str() {
"out" | "outgoing" | "from" => Some(Self::Out),
"in" | "incoming" | "to" => Some(Self::In),
"both" | "all" | "undirected" => Some(Self::Both),
_ => None,
}
}
}
#[derive(Debug, Clone)]
pub struct GraphOptions {
pub hops: usize,
pub include_auto: bool,
pub include_symbols: bool,
pub direction: GraphDirection,
pub max_edges: usize,
pub type_filter: Option<Vec<NodeType>>,
}
impl Default for GraphOptions {
fn default() -> Self {
Self {
hops: 1,
include_auto: true,
include_symbols: true,
direction: GraphDirection::Both,
max_edges: 200,
type_filter: None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct GraphNodeRef {
pub id: String,
pub node_type: String,
pub title: String,
pub file_path: Option<String>,
}
impl GraphNodeRef {
fn from_node(n: &Node) -> Self {
Self {
id: n.id.clone(),
node_type: n.node_type.as_str().to_string(),
title: n.title.clone(),
file_path: n.file_path.clone(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphHopEdge {
pub hop: usize,
pub direction: String,
pub relation_type: String,
pub weight: f32,
pub from_id: String,
pub to: GraphNodeRef,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphNeighborhood {
pub root: GraphNodeRef,
pub hops: usize,
pub nodes_in_subgraph: usize,
pub edges: Vec<GraphHopEdge>,
pub truncated: bool,
pub db_nodes: usize,
pub db_edges: usize,
}
impl GraphNeighborhood {
pub fn to_ascii(&self) -> String {
let mut out = String::new();
let path = self.root.file_path.as_deref().unwrap_or("-");
out.push_str(&format!(
"graph: {} ({}) \"{}\"\n",
self.root.id, self.root.node_type, self.root.title
));
out.push_str(&format!(" path: {path}\n"));
out.push_str(&format!(
" hops={} edges_shown={} nodes_in_subgraph={} db={}/{} n/e{}\n",
self.hops,
self.edges.len(),
self.nodes_in_subgraph,
self.db_nodes,
self.db_edges,
if self.truncated { " [truncated]" } else { "" }
));
if self.edges.is_empty() {
out.push_str(" (no neighbors — orphan or filters excluded all edges)\n");
out.push_str(
" tip: `rustbrain links --auto` soft-links orphans; explicit WikiLinks preferred\n",
);
return out;
}
render_tree_level(&mut out, &self.edges, &self.root.id, 1, "");
out
}
}
fn render_tree_level(
out: &mut String,
all: &[GraphHopEdge],
parent_id: &str,
hop: usize,
prefix: &str,
) {
let kids: Vec<&GraphHopEdge> = all
.iter()
.filter(|e| e.from_id == parent_id && e.hop == hop)
.collect();
let n = kids.len();
for (i, e) in kids.iter().enumerate() {
let last = i + 1 == n;
let branch = if last { "└──" } else { "├──" };
let cont = if last { " " } else { "│ " };
let arrow = if e.direction == "in" { "←" } else { "→" };
let path = e
.to
.file_path
.as_deref()
.map(|p| format!(" [{p}]"))
.unwrap_or_default();
out.push_str(&format!(
"{prefix}{branch}[{arrow} {} w={:.2}] {} ({}) \"{}\"{path}\n",
e.relation_type, e.weight, e.to.id, e.to.node_type, e.to.title
));
let child_prefix = format!("{prefix}{cont}");
render_tree_level(out, all, &e.to.id, hop + 1, &child_prefix);
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphStats {
pub nodes: usize,
pub edges: usize,
pub by_type: Vec<(String, usize)>,
pub by_relation: Vec<(String, usize)>,
pub hubs: Vec<GraphHub>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphHub {
pub id: String,
pub node_type: String,
pub title: String,
pub degree: usize,
}
impl GraphStats {
pub fn to_ascii(&self) -> String {
let mut out = String::new();
out.push_str(&format!(
"graph stats: nodes={} edges={}\n",
self.nodes, self.edges
));
if !self.by_type.is_empty() {
out.push_str(" by type:\n");
for (t, n) in &self.by_type {
out.push_str(&format!(" {t}: {n}\n"));
}
}
if !self.by_relation.is_empty() {
out.push_str(" by relation:\n");
for (r, n) in &self.by_relation {
out.push_str(&format!(" {r}: {n}\n"));
}
}
if !self.hubs.is_empty() {
out.push_str(" hubs (degree = in+out):\n");
for h in &self.hubs {
out.push_str(&format!(
" deg={:>3} {} ({}) \"{}\"\n",
h.degree, h.id, h.node_type, h.title
));
}
}
out
}
}
pub fn graph_stats(db: &Database) -> Result<GraphStats> {
let nodes = db.get_all_nodes()?;
let edges = db.get_all_edges()?;
let mut by_type: HashMap<String, usize> = HashMap::new();
for n in &nodes {
*by_type.entry(n.node_type.as_str().to_string()).or_default() += 1;
}
let mut by_relation: HashMap<String, usize> = HashMap::new();
let mut degree: HashMap<String, usize> = HashMap::new();
for e in &edges {
*by_relation.entry(e.relation_type.clone()).or_default() += 1;
*degree.entry(e.source_id.clone()).or_default() += 1;
*degree.entry(e.target_id.clone()).or_default() += 1;
}
let mut by_type: Vec<(String, usize)> = by_type.into_iter().collect();
by_type.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
let mut by_relation: Vec<(String, usize)> = by_relation.into_iter().collect();
by_relation.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
let id_to_node: HashMap<&str, &Node> = nodes.iter().map(|n| (n.id.as_str(), n)).collect();
let mut hubs: Vec<GraphHub> = degree
.into_iter()
.filter_map(|(id, deg)| {
let n = id_to_node.get(id.as_str())?;
Some(GraphHub {
id,
node_type: n.node_type.as_str().to_string(),
title: n.title.clone(),
degree: deg,
})
})
.collect();
hubs.sort_by(|a, b| {
b.degree
.cmp(&a.degree)
.then_with(|| a.id.cmp(&b.id))
});
hubs.truncate(15);
Ok(GraphStats {
nodes: nodes.len(),
edges: edges.len(),
by_type,
by_relation,
hubs,
})
}
pub fn neighborhood(db: &Database, target: &str, opts: &GraphOptions) -> Result<GraphNeighborhood> {
let root = resolve_graph_target(db, target)?;
let all_edges = db.get_all_edges()?;
let all_nodes = db.get_all_nodes()?;
let node_map: HashMap<String, Node> = all_nodes.into_iter().map(|n| (n.id.clone(), n)).collect();
let hops = opts.hops.max(1);
let mut out_adj: HashMap<String, Vec<&Edge>> = HashMap::new();
let mut in_adj: HashMap<String, Vec<&Edge>> = HashMap::new();
for e in &all_edges {
if !opts.include_auto && is_auto_relation(&e.relation_type) {
continue;
}
out_adj.entry(e.source_id.clone()).or_default().push(e);
in_adj.entry(e.target_id.clone()).or_default().push(e);
}
let mut visited: HashSet<String> = HashSet::new();
visited.insert(root.id.clone());
let mut queue: VecDeque<(String, usize)> = VecDeque::new();
queue.push_back((root.id.clone(), 0));
let mut result_edges: Vec<GraphHopEdge> = Vec::new();
let mut truncated = false;
while let Some((curr, depth)) = queue.pop_front() {
if depth >= hops {
continue;
}
if result_edges.len() >= opts.max_edges {
truncated = true;
break;
}
let mut candidates: Vec<(&Edge, &str, &str)> = Vec::new();
if matches!(opts.direction, GraphDirection::Out | GraphDirection::Both) {
if let Some(list) = out_adj.get(&curr) {
for e in list {
candidates.push((e, e.target_id.as_str(), "out"));
}
}
}
if matches!(opts.direction, GraphDirection::In | GraphDirection::Both) {
if let Some(list) = in_adj.get(&curr) {
for e in list {
candidates.push((e, e.source_id.as_str(), "in"));
}
}
}
candidates.sort_by(|a, b| {
b.0.weight
.partial_cmp(&a.0.weight)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.1.cmp(b.1))
.then_with(|| a.0.relation_type.cmp(&b.0.relation_type))
});
for (e, neighbor_id, dir) in candidates {
if result_edges.len() >= opts.max_edges {
truncated = true;
break;
}
if visited.contains(neighbor_id) {
continue;
}
let Some(neighbor) = node_map.get(neighbor_id) else {
continue;
};
if !opts.include_symbols && neighbor.node_type == NodeType::Symbol {
continue;
}
if let Some(filter) = &opts.type_filter {
if !filter.contains(&neighbor.node_type) {
continue;
}
}
visited.insert(neighbor_id.to_string());
let hop = depth + 1;
result_edges.push(GraphHopEdge {
hop,
direction: dir.to_string(),
relation_type: e.relation_type.clone(),
weight: e.weight,
from_id: curr.clone(),
to: GraphNodeRef::from_node(neighbor),
});
if hop < hops {
queue.push_back((neighbor_id.to_string(), hop));
}
}
}
Ok(GraphNeighborhood {
root: GraphNodeRef::from_node(&root),
hops,
nodes_in_subgraph: visited.len(),
edges: result_edges,
truncated,
db_nodes: node_map.len(),
db_edges: all_edges.len(),
})
}
pub fn resolve_graph_target(db: &Database, target: &str) -> Result<Node> {
let raw = target.trim().trim_start_matches("./").replace('\\', "/");
if raw.is_empty() {
return Err(BrainError::other(
"empty graph target — pass a node id, path, title, or symbol:…",
));
}
if let Some(n) = db.get_node(&raw)? {
return Ok(n);
}
let as_path = Path::new(&raw);
let slug = node_id_from_rel_path(as_path);
if let Some(n) = db.get_node(&slug)? {
return Ok(n);
}
if !raw.starts_with("docs/") {
let with_docs = format!("docs/{raw}");
if let Some(n) = db.get_node(&with_docs)? {
return Ok(n);
}
let slug2 = node_id_from_rel_path(Path::new(&with_docs));
if let Some(n) = db.get_node(&slug2)? {
return Ok(n);
}
}
let sym_raw = raw
.strip_prefix("symbol:")
.or_else(|| raw.strip_prefix("[[symbol:").and_then(|s| s.strip_suffix("]]")))
.unwrap_or(&raw);
if raw.starts_with("symbol:") || raw.contains("::") {
if let Some(sym) = parse_symbol_path(sym_raw) {
let ids: HashSet<String> = db
.get_all_node_ids()?
.into_iter()
.filter(|id| id.starts_with("symbol/"))
.collect();
if let Some(id) = resolve_symbol_ref(&sym, &ids) {
if let Some(n) = db.get_node(&id)? {
return Ok(n);
}
}
}
}
let all = db.get_all_nodes()?;
let path_hits: Vec<&Node> = all
.iter()
.filter(|n| {
n.file_path
.as_ref()
.map(|p| {
let p = p.replace('\\', "/");
p == raw
|| p.ends_with(&raw)
|| p.ends_with(&format!("/{raw}"))
|| p.trim_end_matches(".md") == raw
|| p.trim_end_matches(".md").ends_with(&format!("/{raw}"))
})
.unwrap_or(false)
})
.collect();
if path_hits.len() == 1 {
return Ok(path_hits[0].clone());
}
if path_hits.len() > 1 {
return Err(BrainError::other(format!(
"ambiguous path `{raw}` matches {} nodes — use a fuller path or exact id",
path_hits.len()
)));
}
let lower = raw.to_ascii_lowercase();
let title_hits: Vec<&Node> = all
.iter()
.filter(|n| n.title.to_ascii_lowercase() == lower)
.collect();
if title_hits.len() == 1 {
return Ok(title_hits[0].clone());
}
if title_hits.len() > 1 {
return Err(BrainError::other(format!(
"ambiguous title `{raw}` matches {} nodes — use path or id",
title_hits.len()
)));
}
let sub_hits: Vec<&Node> = all
.iter()
.filter(|n| n.title.to_ascii_lowercase().contains(&lower) && lower.len() >= 3)
.collect();
if sub_hits.len() == 1 {
return Ok(sub_hits[0].clone());
}
Err(BrainError::other(format!(
"no node found for `{raw}` — try path (docs/…md), node id, exact title, or symbol:Name after `rustbrain sync`"
)))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{Edge, Node, NodeType};
use tempfile::tempdir;
fn mk_db() -> (tempfile::TempDir, Database) {
let dir = tempdir().unwrap();
let db = Database::open(dir.path().join("db.sqlite")).unwrap();
let now = 1_700_000_000i64;
let nodes = [
Node {
id: "docs/concepts/raft".into(),
node_type: NodeType::Concept,
title: "Raft".into(),
file_path: Some("docs/concepts/raft.md".into()),
symbol_hash: None,
summary: Some("consensus".into()),
content_hash: None,
created_at: now,
updated_at: now,
},
Node {
id: "docs/concepts/logcompaction".into(),
node_type: NodeType::Concept,
title: "Log Compaction".into(),
file_path: Some("docs/concepts/logcompaction.md".into()),
symbol_hash: None,
summary: None,
content_hash: None,
created_at: now,
updated_at: now,
},
Node {
id: "docs/adr/0001-use-raft".into(),
node_type: NodeType::Adr,
title: "Use Raft".into(),
file_path: Some("docs/adr/0001-use-raft.md".into()),
symbol_hash: None,
summary: None,
content_hash: None,
created_at: now,
updated_at: now,
},
Node {
id: "symbol/demo/lib/storageengine".into(),
node_type: NodeType::Symbol,
title: "StorageEngine".into(),
file_path: Some("src/lib.rs".into()),
symbol_hash: Some(1),
summary: None,
content_hash: None,
created_at: now,
updated_at: now,
},
];
for n in &nodes {
db.insert_node(n).unwrap();
}
db.insert_edge(&Edge {
source_id: "docs/concepts/raft".into(),
target_id: "docs/concepts/logcompaction".into(),
relation_type: "relates_to".into(),
weight: 1.0,
decay_rate: 0.0,
created_at: now,
})
.unwrap();
db.insert_edge(&Edge {
source_id: "docs/concepts/raft".into(),
target_id: "symbol/demo/lib/storageengine".into(),
relation_type: "anchors".into(),
weight: 1.0,
decay_rate: 0.0,
created_at: now,
})
.unwrap();
db.insert_edge(&Edge {
source_id: "docs/adr/0001-use-raft".into(),
target_id: "docs/concepts/raft".into(),
relation_type: "relates_to".into(),
weight: 0.9,
decay_rate: 0.0,
created_at: now,
})
.unwrap();
db.insert_edge(&Edge {
source_id: "docs/concepts/raft".into(),
target_id: "docs/concepts/logcompaction".into(),
relation_type: "auto_filename".into(),
weight: 0.4,
decay_rate: 0.0,
created_at: now,
})
.unwrap();
(dir, db)
}
#[test]
fn neighborhood_both_directions() {
let (_dir, db) = mk_db();
let nb = neighborhood(&db, "docs/concepts/raft", &GraphOptions::default()).unwrap();
assert_eq!(nb.root.id, "docs/concepts/raft");
assert!(nb.edges.len() >= 3);
let ids: HashSet<_> = nb.edges.iter().map(|e| e.to.id.as_str()).collect();
assert!(ids.contains("docs/concepts/logcompaction"));
assert!(ids.contains("symbol/demo/lib/storageengine"));
assert!(ids.contains("docs/adr/0001-use-raft"));
let ascii = nb.to_ascii();
assert!(ascii.contains("relates_to"));
assert!(ascii.contains("anchors"));
}
#[test]
fn no_auto_and_no_symbols_filters() {
let (_dir, db) = mk_db();
let opts = GraphOptions {
include_auto: false,
include_symbols: false,
..GraphOptions::default()
};
let nb = neighborhood(&db, "docs/concepts/raft.md", &opts).unwrap();
for e in &nb.edges {
assert!(!e.relation_type.starts_with("auto_"));
assert_ne!(e.to.node_type, "symbol");
}
assert!(nb.edges.iter().any(|e| e.to.id == "docs/concepts/logcompaction"));
assert!(nb.edges.iter().any(|e| e.to.id == "docs/adr/0001-use-raft"));
}
#[test]
fn resolve_by_title() {
let (_dir, db) = mk_db();
let n = resolve_graph_target(&db, "Raft").unwrap();
assert_eq!(n.id, "docs/concepts/raft");
}
#[test]
fn stats_counts() {
let (_dir, db) = mk_db();
let s = graph_stats(&db).unwrap();
assert_eq!(s.nodes, 4);
assert!(s.edges >= 3);
assert!(!s.hubs.is_empty());
assert!(s.to_ascii().contains("by type"));
}
#[test]
fn direction_out_only() {
let (_dir, db) = mk_db();
let opts = GraphOptions {
direction: GraphDirection::Out,
include_auto: false,
..GraphOptions::default()
};
let nb = neighborhood(&db, "docs/concepts/raft", &opts).unwrap();
assert!(nb.edges.iter().all(|e| e.direction == "out"));
assert!(!nb.edges.iter().any(|e| e.to.id == "docs/adr/0001-use-raft"));
}
}