use crate::error::{Error, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct GraphEdge {
id: u64,
source: u64,
target: u64,
label: String,
properties: HashMap<String, Value>,
}
impl GraphEdge {
pub fn new(id: u64, source: u64, target: u64, label: &str) -> Result<Self> {
let trimmed = label.trim();
if trimmed.is_empty() {
return Err(Error::InvalidEdgeLabel(
"Edge label cannot be empty or whitespace-only".to_string(),
));
}
Ok(Self {
id,
source,
target,
label: trimmed.to_string(),
properties: HashMap::new(),
})
}
#[must_use]
pub fn with_properties(mut self, properties: HashMap<String, Value>) -> Self {
self.properties = properties;
self
}
#[must_use]
pub fn id(&self) -> u64 {
self.id
}
#[must_use]
pub fn source(&self) -> u64 {
self.source
}
#[must_use]
pub fn target(&self) -> u64 {
self.target
}
#[must_use]
pub fn label(&self) -> &str {
&self.label
}
#[must_use]
pub fn properties(&self) -> &HashMap<String, Value> {
&self.properties
}
#[must_use]
pub fn property(&self, name: &str) -> Option<&Value> {
self.properties.get(name)
}
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct EdgeStore {
edges: HashMap<u64, GraphEdge>,
outgoing: HashMap<u64, Vec<u64>>,
incoming: HashMap<u64, Vec<u64>>,
by_label: HashMap<String, Vec<u64>>,
outgoing_by_label: HashMap<(u64, String), Vec<u64>>,
}
impl EdgeStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_capacity(expected_edges: usize, expected_nodes: usize) -> Self {
let expected_labels = 10usize;
let outgoing_by_label_cap = expected_nodes
.saturating_mul(expected_labels)
.saturating_div(10);
Self {
edges: HashMap::with_capacity(expected_edges),
outgoing: HashMap::with_capacity(expected_nodes),
incoming: HashMap::with_capacity(expected_nodes),
by_label: HashMap::with_capacity(expected_labels),
outgoing_by_label: HashMap::with_capacity(outgoing_by_label_cap),
}
}
pub fn add_edge(&mut self, edge: GraphEdge) -> Result<()> {
let id = edge.id();
let source = edge.source();
let target = edge.target();
let label = edge.label().to_string();
if self.edges.contains_key(&id) {
return Err(Error::EdgeExists(id));
}
self.outgoing.entry(source).or_default().push(id);
self.incoming.entry(target).or_default().push(id);
self.by_label.entry(label.clone()).or_default().push(id);
self.outgoing_by_label
.entry((source, label))
.or_default()
.push(id);
self.edges.insert(id, edge);
Ok(())
}
pub fn add_edge_outgoing_only(&mut self, edge: GraphEdge) -> Result<()> {
let id = edge.id();
let source = edge.source();
let label = edge.label().to_string();
if self.edges.contains_key(&id) {
return Err(Error::EdgeExists(id));
}
self.outgoing.entry(source).or_default().push(id);
self.by_label.entry(label.clone()).or_default().push(id);
self.outgoing_by_label
.entry((source, label))
.or_default()
.push(id);
self.edges.insert(id, edge);
Ok(())
}
pub fn add_edge_incoming_only(&mut self, edge: GraphEdge) -> Result<()> {
let id = edge.id();
let target = edge.target();
if self.edges.contains_key(&id) {
return Err(Error::EdgeExists(id));
}
self.incoming.entry(target).or_default().push(id);
self.edges.insert(id, edge);
Ok(())
}
#[must_use]
pub fn edge_count(&self) -> usize {
self.edges.len()
}
#[must_use]
pub fn outgoing_edge_count(&self) -> usize {
self.outgoing.values().map(Vec::len).sum()
}
#[must_use]
pub fn get_edge(&self, id: u64) -> Option<&GraphEdge> {
self.edges.get(&id)
}
#[must_use]
pub fn get_outgoing(&self, node_id: u64) -> Vec<&GraphEdge> {
self.outgoing
.get(&node_id)
.map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
.unwrap_or_default()
}
#[must_use]
pub fn get_incoming(&self, node_id: u64) -> Vec<&GraphEdge> {
self.incoming
.get(&node_id)
.map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
.unwrap_or_default()
}
#[must_use]
pub fn get_outgoing_by_label(&self, node_id: u64, label: &str) -> Vec<&GraphEdge> {
self.outgoing_by_label
.get(&(node_id, label.to_string()))
.map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
.unwrap_or_default()
}
#[must_use]
pub fn get_edges_by_label(&self, label: &str) -> Vec<&GraphEdge> {
self.by_label
.get(label)
.map(|ids| ids.iter().filter_map(|id| self.edges.get(id)).collect())
.unwrap_or_default()
}
#[must_use]
pub fn get_incoming_by_label(&self, node_id: u64, label: &str) -> Vec<&GraphEdge> {
self.get_incoming(node_id)
.into_iter()
.filter(|e| e.label() == label)
.collect()
}
#[must_use]
pub fn contains_edge(&self, edge_id: u64) -> bool {
self.edges.contains_key(&edge_id)
}
#[must_use]
pub fn len(&self) -> usize {
self.edges.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.edges.is_empty()
}
#[must_use]
pub fn all_edges(&self) -> Vec<&GraphEdge> {
self.edges.values().collect()
}
pub fn remove_edge(&mut self, edge_id: u64) {
if let Some(edge) = self.edges.remove(&edge_id) {
let source = edge.source();
let label = edge.label().to_string();
if let Some(ids) = self.outgoing.get_mut(&source) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self.incoming.get_mut(&edge.target()) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self.by_label.get_mut(&label) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self.outgoing_by_label.get_mut(&(source, label)) {
ids.retain(|&id| id != edge_id);
}
}
}
pub fn remove_edge_outgoing_only(&mut self, edge_id: u64) {
if let Some(edge) = self.edges.remove(&edge_id) {
let source = edge.source();
let label = edge.label().to_string();
if let Some(ids) = self.outgoing.get_mut(&source) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self.by_label.get_mut(&label) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self.outgoing_by_label.get_mut(&(source, label)) {
ids.retain(|&id| id != edge_id);
}
}
}
pub fn remove_edge_incoming_only(&mut self, edge_id: u64) {
if let Some(edge) = self.edges.remove(&edge_id) {
if let Some(ids) = self.incoming.get_mut(&edge.target()) {
ids.retain(|&id| id != edge_id);
}
}
}
pub fn to_bytes(&self) -> std::result::Result<Vec<u8>, postcard::Error> {
postcard::to_allocvec(self)
}
pub fn from_bytes(bytes: &[u8]) -> std::result::Result<Self, postcard::Error> {
postcard::from_bytes(bytes)
}
pub fn save_to_file(&self, path: &std::path::Path) -> std::io::Result<()> {
let bytes = self
.to_bytes()
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
std::fs::write(path, bytes)
}
pub fn load_from_file(path: &std::path::Path) -> std::io::Result<Self> {
let bytes = std::fs::read(path)?;
Self::from_bytes(&bytes)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))
}
pub fn remove_node_edges(&mut self, node_id: u64) {
let outgoing_ids: Vec<u64> = self.outgoing.remove(&node_id).unwrap_or_default();
let incoming_ids: Vec<u64> = self.incoming.remove(&node_id).unwrap_or_default();
for edge_id in outgoing_ids {
if let Some(edge) = self.edges.remove(&edge_id) {
self.purge_incoming_index(edge_id, edge.target());
self.purge_label_indices(edge_id, node_id, edge.label());
}
}
for edge_id in incoming_ids {
if let Some(edge) = self.edges.remove(&edge_id) {
let source = edge.source();
self.purge_outgoing_index(edge_id, source);
self.purge_label_indices(edge_id, source, edge.label());
}
}
}
#[inline]
fn purge_incoming_index(&mut self, edge_id: u64, target_node: u64) {
if let Some(ids) = self.incoming.get_mut(&target_node) {
ids.retain(|&id| id != edge_id);
}
}
#[inline]
fn purge_outgoing_index(&mut self, edge_id: u64, source_node: u64) {
if let Some(ids) = self.outgoing.get_mut(&source_node) {
ids.retain(|&id| id != edge_id);
}
}
#[inline]
fn purge_label_indices(&mut self, edge_id: u64, source_node: u64, label: &str) {
if let Some(ids) = self.by_label.get_mut(label) {
ids.retain(|&id| id != edge_id);
}
if let Some(ids) = self
.outgoing_by_label
.get_mut(&(source_node, label.to_string()))
{
ids.retain(|&id| id != edge_id);
}
}
}