use arcstr::ArcStr;
use grafeo_common::types::{EdgeId, EpochId, NodeId, PropertyKey, PropertyMap, Value};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone)]
pub struct Edge {
pub id: EdgeId,
pub src: NodeId,
pub dst: NodeId,
pub edge_type: ArcStr,
pub properties: PropertyMap,
}
impl Edge {
#[must_use]
pub fn new(id: EdgeId, src: NodeId, dst: NodeId, edge_type: impl Into<ArcStr>) -> Self {
Self {
id,
src,
dst,
edge_type: edge_type.into(),
properties: PropertyMap::new(),
}
}
pub fn set_property(&mut self, key: impl Into<PropertyKey>, value: impl Into<Value>) {
self.properties.insert(key.into(), value.into());
}
#[must_use]
pub fn get_property(&self, key: &str) -> Option<&Value> {
self.properties.get(&PropertyKey::new(key))
}
pub fn remove_property(&mut self, key: &str) -> Option<Value> {
self.properties.remove(&PropertyKey::new(key))
}
#[must_use]
pub fn properties_as_btree(&self) -> std::collections::BTreeMap<PropertyKey, Value> {
self.properties.to_btree_map()
}
#[must_use]
pub fn other_endpoint(&self, node: NodeId) -> Option<NodeId> {
if node == self.src {
Some(self.dst)
} else if node == self.dst {
Some(self.src)
} else {
None
}
}
}
#[repr(C)]
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct EdgeRecord {
pub id: EdgeId,
pub src: NodeId,
pub dst: NodeId,
pub type_id: u32,
pub props_offset: u32,
pub props_count: u16,
pub flags: EdgeFlags,
pub epoch: EpochId,
}
impl EdgeRecord {
pub const FLAG_DELETED: u16 = 1 << 0;
pub const FLAG_HAS_VERSION: u16 = 1 << 1;
#[must_use]
pub const fn new(id: EdgeId, src: NodeId, dst: NodeId, type_id: u32, epoch: EpochId) -> Self {
Self {
id,
src,
dst,
type_id,
props_offset: 0,
props_count: 0,
flags: EdgeFlags(0),
epoch,
}
}
#[must_use]
pub const fn is_deleted(&self) -> bool {
self.flags.contains(Self::FLAG_DELETED)
}
pub fn set_deleted(&mut self, deleted: bool) {
if deleted {
self.flags.set(Self::FLAG_DELETED);
} else {
self.flags.clear(Self::FLAG_DELETED);
}
}
}
#[repr(transparent)]
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
pub struct EdgeFlags(pub u16);
impl EdgeFlags {
#[must_use]
pub const fn contains(&self, flag: u16) -> bool {
(self.0 & flag) != 0
}
pub fn set(&mut self, flag: u16) {
self.0 |= flag;
}
pub fn clear(&mut self, flag: u16) {
self.0 &= !flag;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_edge_creation() {
let edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
assert_eq!(edge.id, EdgeId::new(1));
assert_eq!(edge.src, NodeId::new(10));
assert_eq!(edge.dst, NodeId::new(20));
assert_eq!(edge.edge_type.as_str(), "KNOWS");
}
#[test]
fn test_edge_properties() {
let mut edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
edge.set_property("since", 2020i64);
edge.set_property("weight", 1.5f64);
assert_eq!(
edge.get_property("since").and_then(|v| v.as_int64()),
Some(2020)
);
assert_eq!(
edge.get_property("weight").and_then(|v| v.as_float64()),
Some(1.5)
);
}
#[test]
fn test_edge_other_endpoint() {
let edge = Edge::new(EdgeId::new(1), NodeId::new(10), NodeId::new(20), "KNOWS");
assert_eq!(edge.other_endpoint(NodeId::new(10)), Some(NodeId::new(20)));
assert_eq!(edge.other_endpoint(NodeId::new(20)), Some(NodeId::new(10)));
assert_eq!(edge.other_endpoint(NodeId::new(30)), None);
}
#[test]
fn test_edge_record_flags() {
let mut record = EdgeRecord::new(
EdgeId::new(1),
NodeId::new(10),
NodeId::new(20),
0,
EpochId::INITIAL,
);
assert!(!record.is_deleted());
record.set_deleted(true);
assert!(record.is_deleted());
}
#[test]
fn test_edge_record_size() {
let size = std::mem::size_of::<EdgeRecord>();
assert!(size <= 64, "EdgeRecord is {} bytes", size);
}
}