use std::collections::{HashSet, VecDeque};
use crate::graph::node::Language;
use crate::graph::unified::concurrent::GraphSnapshot;
use crate::graph::unified::edge::EdgeKind;
use crate::graph::unified::node::NodeId;
pub const DEFAULT_MAX_DEPTH: usize = 2;
pub const MAX_DEPTH_CAP: usize = 5;
pub const DEFAULT_MAX_RESULTS: usize = 1000;
#[derive(Debug, Clone)]
pub struct SeededSubgraphConfig {
pub max_depth: usize,
pub max_results: usize,
pub languages: Vec<Language>,
}
impl Default for SeededSubgraphConfig {
fn default() -> Self {
Self {
max_depth: DEFAULT_MAX_DEPTH,
max_results: DEFAULT_MAX_RESULTS,
languages: Vec::new(),
}
}
}
impl SeededSubgraphConfig {
#[must_use]
pub fn normalized(mut self) -> Self {
self.max_depth = self.max_depth.min(MAX_DEPTH_CAP);
if self.max_results == 0 {
self.max_results = DEFAULT_MAX_RESULTS;
}
self
}
}
#[derive(Debug, Clone)]
pub struct EdgeRetention<C> {
pub class: C,
pub traverse: bool,
}
#[derive(Debug, Clone)]
pub struct RetainedEdge<C> {
pub from: NodeId,
pub to: NodeId,
pub class: C,
}
#[derive(Debug, Clone)]
pub struct SeededSubgraph<C> {
pub nodes: Vec<NodeId>,
pub edges: Vec<RetainedEdge<C>>,
pub truncated: bool,
}
pub fn build_seeded_subgraph<C, F>(
snapshot: &GraphSnapshot,
seeds: &[NodeId],
config: &SeededSubgraphConfig,
classify_edge: F,
) -> SeededSubgraph<C>
where
F: Fn(&EdgeKind) -> Option<EdgeRetention<C>>,
C: Clone,
{
let language_ok = |node: NodeId| -> bool {
if config.languages.is_empty() {
return true;
}
snapshot
.get_node(node)
.and_then(|entry| snapshot.files().language_for_file(entry.file))
.is_some_and(|lang| config.languages.contains(&lang))
};
let mut ordered_seeds: Vec<NodeId> = seeds
.iter()
.copied()
.filter(|&id| {
snapshot.get_node(id).is_some_and(|e| !e.is_unified_loser()) && language_ok(id)
})
.collect();
ordered_seeds.sort_by_key(|id| id.index());
ordered_seeds.dedup();
let mut visited: HashSet<NodeId> = HashSet::new();
let mut discovery: Vec<NodeId> = Vec::new();
let mut pending_edges: Vec<(NodeId, NodeId, C)> = Vec::new();
let mut queue: VecDeque<(NodeId, usize)> =
ordered_seeds.iter().map(|&id| (id, 0usize)).collect();
let mut truncated = false;
while let Some((current, depth)) = queue.pop_front() {
if depth > config.max_depth || visited.contains(¤t) {
continue;
}
if !language_ok(current) {
continue;
}
if discovery.len() >= config.max_results {
truncated = true;
break;
}
visited.insert(current);
discovery.push(current);
let mut outgoing: Vec<(NodeId, EdgeKind, u64, C, bool)> = Vec::new();
for edge in snapshot.edges().edges_from(current) {
let Some(retention) = classify_edge(&edge.kind) else {
continue;
};
if snapshot
.get_node(edge.target)
.is_none_or(|e| e.is_unified_loser())
|| !language_ok(edge.target)
{
continue;
}
outgoing.push((
edge.target,
edge.kind.clone(),
edge.seq,
retention.class,
retention.traverse,
));
}
outgoing.sort_by(|a, b| {
a.0.index()
.cmp(&b.0.index())
.then_with(|| edge_kind_rank(&a.1).cmp(&edge_kind_rank(&b.1)))
.then_with(|| a.2.cmp(&b.2))
});
for (target, _kind, _seq, class, traverse) in outgoing {
pending_edges.push((current, target, class));
if traverse && depth < config.max_depth && !visited.contains(&target) {
queue.push_back((target, depth + 1));
}
}
}
let edges: Vec<RetainedEdge<C>> = pending_edges
.into_iter()
.filter(|(from, to, _)| visited.contains(from) && visited.contains(to))
.map(|(from, to, class)| RetainedEdge { from, to, class })
.collect();
SeededSubgraph {
nodes: discovery,
edges,
truncated,
}
}
fn edge_kind_rank(kind: &EdgeKind) -> u16 {
match kind {
EdgeKind::HttpRequest { .. } => 0,
EdgeKind::GrpcCall { .. } => 1,
EdgeKind::DbQuery { .. } => 2,
EdgeKind::TableRead { .. } => 3,
EdgeKind::TableWrite { .. } => 4,
EdgeKind::TriggeredBy { .. } => 5,
EdgeKind::MessageQueue { .. } => 6,
EdgeKind::WebSocket { .. } => 7,
EdgeKind::GraphQLOperation { .. } => 8,
EdgeKind::FfiCall { .. } => 9,
EdgeKind::WebAssemblyCall => 10,
EdgeKind::Calls { .. } => 11,
EdgeKind::Imports { .. } => 12,
EdgeKind::Exports { .. } => 13,
EdgeKind::TypeOf { .. } => 14,
EdgeKind::References => 15,
EdgeKind::Inherits => 16,
EdgeKind::Implements => 17,
EdgeKind::Defines => 18,
EdgeKind::Contains => 19,
other => 20u16.wrapping_add(discriminant_index(other)),
}
}
fn discriminant_index(kind: &EdgeKind) -> u16 {
let label = format!("{kind:?}");
let mut acc: u16 = 0;
for b in label.bytes().take(8) {
acc = acc.wrapping_mul(31).wrapping_add(u16::from(b));
}
acc
}