use std::path::Path;
use anyhow::{Context, Result};
use crate::graph::node::Language;
use crate::graph::unified::{CodeGraph, EdgeKind, NodeKind};
use super::edge::EdgeId;
use super::node::NodeId;
pub struct GraphWriteTxn<'a> {
graph: &'a mut CodeGraph,
pending_nodes: Vec<(Language, String, String, NodeKind, Option<String>)>,
pending_edges: Vec<(NodeId, NodeId, EdgeKind)>,
committed: bool,
}
impl<'a> GraphWriteTxn<'a> {
#[must_use]
pub fn new(graph: &'a mut CodeGraph) -> Self {
Self {
graph,
pending_nodes: Vec::new(),
pending_edges: Vec::new(),
committed: false,
}
}
pub fn add_node(
&mut self,
language: Language,
file: impl Into<String>,
symbol: impl Into<String>,
kind: NodeKind,
signature: Option<String>,
) -> Result<NodeId> {
let node_index =
u32::try_from(self.pending_nodes.len()).expect("pending node index exceeds u32::MAX");
let temp_id = NodeId::new(node_index, 0);
self.pending_nodes
.push((language, file.into(), symbol.into(), kind, signature));
Ok(temp_id)
}
pub fn add_edge(&mut self, source: NodeId, target: NodeId, kind: EdgeKind) -> Result<EdgeId> {
let edge_index =
u32::try_from(self.pending_edges.len()).expect("pending edge index exceeds u32::MAX");
let temp_id = EdgeId::new(edge_index);
self.pending_edges.push((source, target, kind));
Ok(temp_id)
}
pub fn commit(mut self) -> Result<()> {
assert!(!self.committed, "Transaction already committed");
self.committed = true;
let mut node_id_map = Vec::new();
for (_language, file, symbol, kind, signature) in &self.pending_nodes {
use crate::graph::unified::storage::arena::NodeEntry;
let file_id = self
.graph
.files_mut()
.register(Path::new(file))
.with_context(|| format!("Failed to register file: {file}"))?;
let name_id = self.graph.strings_mut().intern(symbol)?;
let mut entry = NodeEntry::new(*kind, name_id, file_id);
if let Some(sig) = signature {
let signature_id = self.graph.strings_mut().intern(sig)?;
entry = entry.with_signature(signature_id);
}
let node_id = self
.graph
.nodes_mut()
.alloc(entry)
.with_context(|| format!("Failed to allocate node for symbol: {symbol}"))?;
node_id_map.push(node_id);
}
for (source_temp, target_temp, kind) in &self.pending_edges {
let source = *node_id_map
.get(source_temp.index() as usize)
.ok_or_else(|| {
anyhow::anyhow!(
"Source node {} not found in transaction",
source_temp.index()
)
})?;
let target = *node_id_map
.get(target_temp.index() as usize)
.ok_or_else(|| {
anyhow::anyhow!(
"Target node {} not found in transaction",
target_temp.index()
)
})?;
let file_id = if let Some(entry) = self.graph.nodes().get(source) {
entry.file
} else {
self.graph.files_mut().register(Path::new(""))?
};
self.graph
.edges_mut()
.add_edge(source, target, kind.clone(), file_id);
}
self.graph.bump_epoch();
Ok(())
}
#[must_use]
pub fn pending_nodes(&self) -> usize {
self.pending_nodes.len()
}
#[must_use]
pub fn pending_edges(&self) -> usize {
self.pending_edges.len()
}
}
impl Drop for GraphWriteTxn<'_> {
fn drop(&mut self) {
if !self.committed && (!self.pending_nodes.is_empty() || !self.pending_edges.is_empty()) {
log::warn!(
"GraphWriteTxn dropped without commit ({} nodes, {} edges discarded)",
self.pending_nodes.len(),
self.pending_edges.len()
);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_txn_new() {
let mut graph = CodeGraph::new();
let txn = GraphWriteTxn::new(&mut graph);
assert_eq!(txn.pending_nodes(), 0);
assert_eq!(txn.pending_edges(), 0);
}
#[test]
fn test_txn_add_node() {
let mut graph = CodeGraph::new();
let mut txn = GraphWriteTxn::new(&mut graph);
let node_id = txn
.add_node(
Language::Rust,
"main.rs",
"main",
NodeKind::Function,
Some("fn main()".to_string()),
)
.expect("add_node");
assert_eq!(txn.pending_nodes(), 1);
assert_eq!(node_id.index(), 0);
}
#[test]
fn test_txn_add_edge() {
let mut graph = CodeGraph::new();
let mut txn = GraphWriteTxn::new(&mut graph);
let source = txn
.add_node(
Language::Rust,
"main.rs",
"caller",
NodeKind::Function,
None,
)
.expect("add_node");
let target = txn
.add_node(
Language::Rust,
"main.rs",
"callee",
NodeKind::Function,
None,
)
.expect("add_node");
let _edge_id = txn
.add_edge(
source,
target,
EdgeKind::Calls {
argument_count: 0,
is_async: false,
},
)
.expect("add_edge");
assert_eq!(txn.pending_nodes(), 2);
assert_eq!(txn.pending_edges(), 1);
}
#[test]
fn test_txn_commit_empty() {
let mut graph = CodeGraph::new();
let txn = GraphWriteTxn::new(&mut graph);
txn.commit().expect("commit empty txn");
}
#[test]
fn test_txn_commit_nodes() {
let mut graph = CodeGraph::new();
let initial_epoch = graph.epoch();
let mut txn = GraphWriteTxn::new(&mut graph);
txn.add_node(
Language::Rust,
"main.rs",
"main",
NodeKind::Function,
Some("fn main()".to_string()),
)
.expect("add_node");
txn.commit().expect("commit");
assert_eq!(graph.epoch(), initial_epoch + 1);
assert_eq!(graph.nodes().len(), 1);
}
#[test]
fn test_txn_commit_edges() {
let mut graph = CodeGraph::new();
let mut txn = GraphWriteTxn::new(&mut graph);
let source = txn
.add_node(
Language::Rust,
"main.rs",
"caller",
NodeKind::Function,
None,
)
.expect("add_node");
let target = txn
.add_node(
Language::Rust,
"main.rs",
"callee",
NodeKind::Function,
None,
)
.expect("add_node");
txn.add_edge(
source,
target,
EdgeKind::Calls {
argument_count: 0,
is_async: false,
},
)
.expect("add_edge");
txn.commit().expect("commit");
assert_eq!(graph.nodes().len(), 2);
}
}