use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use petgraph::graph::{EdgeIndex, NodeIndex};
use petgraph::stable_graph::StableDiGraph;
use petgraph::visit::NodeIndexable;
use rustc_hash::FxHashMap;
use crate::datatypes::Value;
use crate::graph::core::iterators::GraphNodeIndices;
use crate::graph::schema::{EdgeData, InternedKey, NodeData};
use crate::graph::storage::column_store::ColumnStore;
use crate::graph::storage::undo::UndoJournal;
use crate::graph::storage::{GraphRead, GraphWrite, MemoryGraph};
pub struct ForkedGraph {
base: Arc<MemoryGraph>,
nodes: FxHashMap<u32, NodeData>,
edges: FxHashMap<u32, EdgeData>,
appended: u32,
column_stores: FxHashMap<InternedKey, Arc<ColumnStore>>,
peer_counts: RwLock<HashMap<u64, Arc<super::MemoryPeerCounts>>>,
undo: Option<Box<UndoJournal>>,
slot_mirror: super::slot_mirror::SlotMirror,
}
pub(crate) fn can_fork(base: &MemoryGraph) -> bool {
let node_bound = base.inner().node_bound();
let edge_bound = petgraph::visit::EdgeIndexable::edge_bound(base.inner());
base.slot_mirror.predict_next_node(node_bound) == Some(NodeIndex::new(node_bound))
&& base.slot_mirror.predict_next_edge(edge_bound) == Some(EdgeIndex::new(edge_bound))
}
impl ForkedGraph {
pub(crate) fn new(base: Arc<MemoryGraph>) -> Self {
let column_stores = base.column_stores.clone();
let slot_mirror = base.slot_mirror.clone();
Self {
base,
nodes: FxHashMap::default(),
edges: FxHashMap::default(),
appended: 0,
column_stores,
peer_counts: RwLock::new(HashMap::new()),
undo: None,
slot_mirror,
}
}
#[inline]
fn append_floor(&self) -> usize {
self.base.inner().node_bound()
}
#[cfg(test)]
#[inline]
pub(crate) fn overlay_node_count(&self) -> usize {
self.nodes.len()
}
fn apply_overlay(&mut self, target: &mut MemoryGraph) {
let floor = target.inner().node_bound();
for offset in 0..self.appended {
let idx = floor as u32 + offset;
let data = self
.nodes
.remove(&idx)
.expect("an appended overlay index must carry its node weight");
let actual = GraphWrite::add_node(target, data);
assert_eq!(
actual.index() as u32,
idx,
"fold-back allocated node {} where the overlay handed out {idx}; \
slot identity is broken and every index keyed on the overlay's \
number is now wrong (see storage/forked.rs)",
actual.index()
);
}
for (idx, data) in self.nodes.drain() {
if let Some(slot) = target
.inner_mut()
.node_weight_mut(NodeIndex::new(idx as usize))
{
*slot = data;
}
}
for (idx, data) in self.edges.drain() {
if let Some(slot) = target
.inner_mut()
.edge_weight_mut(EdgeIndex::new(idx as usize))
{
*slot = data;
}
}
target.column_stores = std::mem::take(&mut self.column_stores);
target.undo = self.undo.take();
self.appended = 0;
}
pub(crate) fn try_compact(mut self: Box<Self>) -> Result<MemoryGraph, Box<Self>> {
if Arc::get_mut(&mut self.base).is_none() {
return Err(self);
}
let mut owned = Arc::try_unwrap(std::mem::replace(
&mut self.base,
Arc::new(MemoryGraph::new()),
))
.unwrap_or_else(|_| unreachable!("get_mut proved unique ownership"));
self.apply_overlay(&mut owned);
Ok(owned)
}
pub(crate) fn materialise(&mut self) -> MemoryGraph {
#[cfg(test)]
super::backend::note_nodes_copied(self.base.inner().node_count());
let mut owned = self.base.deep_clone();
self.apply_overlay(&mut owned);
owned
}
pub(crate) fn to_memory_graph(&self) -> MemoryGraph {
let mut clone = ForkedGraph {
base: Arc::clone(&self.base),
nodes: self.nodes.clone(),
edges: self.edges.clone(),
appended: self.appended,
column_stores: self.column_stores.clone(),
peer_counts: RwLock::new(HashMap::new()),
undo: None,
slot_mirror: self.slot_mirror.clone(),
};
let mut owned = self.base.deep_clone();
clone.apply_overlay(&mut owned);
owned
}
#[inline]
pub(crate) fn begin_undo(&mut self) {
self.undo = Some(Box::new(UndoJournal::new()));
}
#[inline]
pub(crate) fn take_undo(&mut self) -> Option<Box<UndoJournal>> {
self.undo.take()
}
#[inline]
pub(crate) fn undo_journal_mut(&mut self) -> Option<&mut UndoJournal> {
self.undo.as_deref_mut()
}
#[inline]
fn invalidate_peer_counts(&mut self) {
if let Ok(mut cache) = self.peer_counts.write() {
cache.clear();
}
}
#[inline]
fn cow_node(&mut self, idx: NodeIndex) -> Option<&mut NodeData> {
let raw = idx.index() as u32;
if !self.nodes.contains_key(&raw) {
let base = self.base.inner().node_weight(idx)?.clone();
self.nodes.insert(raw, base);
}
self.nodes.get_mut(&raw)
}
#[inline]
fn cow_edge(&mut self, idx: EdgeIndex) -> Option<&mut EdgeData> {
let raw = idx.index() as u32;
if !self.edges.contains_key(&raw) {
let base = self.base.inner().edge_weight(idx)?.clone();
self.edges.insert(raw, base);
}
self.edges.get_mut(&raw)
}
#[cold]
fn capture_node_weight(&mut self, idx: NodeIndex) {
let current = GraphRead::node_weight(self, idx).cloned();
if let Some(journal) = self.undo.as_deref_mut() {
journal.note_node_weight(idx, || current);
}
}
#[cold]
fn capture_edge_weight(&mut self, idx: EdgeIndex) {
let current = GraphRead::edge_weight(self, idx).cloned();
if let Some(journal) = self.undo.as_deref_mut() {
journal.note_edge_weight(idx, || current);
}
}
#[inline]
fn capture_property_pre_image(&mut self, idx: NodeIndex) {
if self.undo.is_none() {
return;
}
let Some(nd) = GraphRead::node_weight(self, idx) else {
return;
};
match nd.properties.columnar_row_id() {
None => self.capture_node_weight(idx),
Some(_) => {
let type_key = nd.node_type;
let prior = self.column_stores.get(&type_key).map(Arc::clone);
if let Some(journal) = self.undo.as_deref_mut() {
journal.note_columnar_fork(type_key, || prior);
}
}
}
}
#[inline]
fn columnar_row_of(&self, idx: NodeIndex) -> Option<(InternedKey, u32)> {
let nd = GraphRead::node_weight(self, idx)?;
nd.properties
.columnar_row_id()
.map(|row_id| (nd.node_type, row_id))
}
#[inline]
pub(crate) fn base_stable_digraph(&self) -> &StableDiGraph<NodeData, EdgeData> {
self.base.inner()
}
}
impl Clone for ForkedGraph {
fn clone(&self) -> Self {
Self {
base: Arc::clone(&self.base),
nodes: self.nodes.clone(),
edges: self.edges.clone(),
appended: self.appended,
column_stores: self.column_stores.clone(),
peer_counts: RwLock::new(HashMap::new()),
undo: None,
slot_mirror: self.slot_mirror.clone(),
}
}
}
impl std::fmt::Debug for ForkedGraph {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"ForkedGraph {{ base: {} nodes / {} edges, overlay: {} node weights, \
{} edge weights, {} appended }}",
self.base.inner().node_count(),
self.base.inner().edge_count(),
self.nodes.len(),
self.edges.len(),
self.appended
)
}
}
impl GraphRead for ForkedGraph {
type NodeIndicesIter<'a> = GraphNodeIndices<'a>;
type EdgeIndicesIter<'a> = <MemoryGraph as GraphRead>::EdgeIndicesIter<'a>;
type EdgesIter<'a> = <MemoryGraph as GraphRead>::EdgesIter<'a>;
type EdgeReferencesIter<'a> = <MemoryGraph as GraphRead>::EdgeReferencesIter<'a>;
type EdgesConnectingIter<'a> = <MemoryGraph as GraphRead>::EdgesConnectingIter<'a>;
type NeighborsIter<'a> = <MemoryGraph as GraphRead>::NeighborsIter<'a>;
#[inline]
fn node_count(&self) -> usize {
self.base.inner().node_count() + self.appended as usize
}
#[inline]
fn edge_count(&self) -> usize {
self.base.inner().edge_count()
}
#[inline]
fn node_bound(&self) -> usize {
self.append_floor() + self.appended as usize
}
#[inline]
fn is_memory(&self) -> bool {
true
}
#[inline]
fn node_weight(&self, idx: NodeIndex) -> Option<&NodeData> {
match self.nodes.get(&(idx.index() as u32)) {
Some(data) => Some(data),
None => self.base.inner().node_weight(idx),
}
}
#[inline]
fn node_type_of(&self, idx: NodeIndex) -> Option<InternedKey> {
self.node_weight(idx).map(|n| n.node_type)
}
#[inline]
fn get_node_property(&self, idx: NodeIndex, key: InternedKey) -> Option<Value> {
self.node_view(idx)?.get_value(key)
}
#[inline]
fn get_node_id(&self, idx: NodeIndex) -> Option<Value> {
Some(self.node_view(idx)?.id().into_owned())
}
#[inline]
fn get_node_title(&self, idx: NodeIndex) -> Option<Value> {
Some(self.node_view(idx)?.title().into_owned())
}
#[inline]
fn str_prop_eq(&self, idx: NodeIndex, key: InternedKey, target: &str) -> Option<bool> {
self.node_view(idx)?.str_prop_eq(key, target)
}
#[inline]
fn edges_directed_filtered(
&self,
idx: NodeIndex,
dir: petgraph::Direction,
conn_type_filter: Option<InternedKey>,
) -> Self::EdgesIter<'_> {
GraphRead::edges_directed_filtered(&*self.base, idx, dir, conn_type_filter)
}
fn edge_endpoint_keys<'a>(
&'a self,
) -> Box<dyn Iterator<Item = (NodeIndex, NodeIndex, InternedKey)> + 'a> {
GraphRead::edge_endpoint_keys(&*self.base)
}
fn count_edges_grouped_by_peer(
&self,
conn_type: InternedKey,
dir: petgraph::Direction,
deadline: Option<std::time::Instant>,
) -> Result<HashMap<u32, i64>, String> {
GraphRead::count_edges_grouped_by_peer(&*self.base, conn_type, dir, deadline)
}
fn count_edges_filtered(
&self,
node: NodeIndex,
dir: petgraph::Direction,
conn_type: Option<InternedKey>,
other_node_type: Option<InternedKey>,
deadline: Option<std::time::Instant>,
) -> Result<usize, String> {
GraphRead::count_edges_filtered(
&*self.base,
node,
dir,
conn_type,
other_node_type,
deadline,
)
}
#[inline]
fn column_store(&self, type_key: InternedKey) -> Option<&Arc<ColumnStore>> {
self.column_stores.get(&type_key)
}
fn column_stores_iter(
&self,
) -> Box<dyn Iterator<Item = (InternedKey, &Arc<ColumnStore>)> + '_> {
Box::new(self.column_stores.iter().map(|(k, v)| (*k, v)))
}
#[inline]
fn node_indices(&self) -> Self::NodeIndicesIter<'_> {
GraphNodeIndices::Forked {
base: Box::new(self.base.inner().node_indices()),
appended: self.append_floor()..self.node_bound(),
}
}
#[inline]
fn edge_indices(&self) -> Self::EdgeIndicesIter<'_> {
GraphRead::edge_indices(&*self.base)
}
#[inline]
fn edge_references(&self) -> Self::EdgeReferencesIter<'_> {
GraphRead::edge_references(&*self.base)
}
fn edge_weights<'a>(&'a self) -> Box<dyn Iterator<Item = &'a EdgeData> + 'a> {
GraphRead::edge_weights(&*self.base)
}
#[inline]
fn edges_directed(&self, idx: NodeIndex, dir: petgraph::Direction) -> Self::EdgesIter<'_> {
GraphRead::edges_directed(&*self.base, idx, dir)
}
#[inline]
fn edges(&self, idx: NodeIndex) -> Self::EdgesIter<'_> {
GraphRead::edges(&*self.base, idx)
}
#[inline]
fn edges_connecting(&self, a: NodeIndex, b: NodeIndex) -> Self::EdgesConnectingIter<'_> {
GraphRead::edges_connecting(&*self.base, a, b)
}
#[inline]
fn edge_weight(&self, idx: EdgeIndex) -> Option<&EdgeData> {
match self.edges.get(&(idx.index() as u32)) {
Some(data) => Some(data),
None => self.base.inner().edge_weight(idx),
}
}
#[inline]
fn find_edge(&self, a: NodeIndex, b: NodeIndex) -> Option<EdgeIndex> {
GraphRead::find_edge(&*self.base, a, b)
}
#[inline]
fn edge_endpoints(&self, idx: EdgeIndex) -> Option<(NodeIndex, NodeIndex)> {
GraphRead::edge_endpoints(&*self.base, idx)
}
#[inline]
fn neighbors_directed(
&self,
idx: NodeIndex,
dir: petgraph::Direction,
) -> Self::NeighborsIter<'_> {
GraphRead::neighbors_directed(&*self.base, idx, dir)
}
#[inline]
fn neighbors_undirected(&self, idx: NodeIndex) -> Self::NeighborsIter<'_> {
GraphRead::neighbors_undirected(&*self.base, idx)
}
}
impl GraphWrite for ForkedGraph {
#[inline]
fn node_weight_mut(&mut self, idx: NodeIndex) -> Option<&mut NodeData> {
if self.undo.is_some() {
self.capture_node_weight(idx);
}
self.cow_node(idx)
}
#[inline]
fn node_weight_mut_silent(&mut self, idx: NodeIndex) -> Option<&mut NodeData> {
self.cow_node(idx)
}
#[inline]
fn edge_weight_mut(&mut self, idx: EdgeIndex) -> Option<&mut EdgeData> {
self.invalidate_peer_counts();
if self.undo.is_some() {
self.capture_edge_weight(idx);
}
self.cow_edge(idx)
}
#[inline]
fn install_column_store(&mut self, type_key: InternedKey, store: Arc<ColumnStore>) {
self.column_stores.insert(type_key, store);
}
#[inline]
fn column_store_mut(&mut self, type_key: InternedKey) -> Option<&mut Arc<ColumnStore>> {
self.column_stores.get_mut(&type_key)
}
#[inline]
fn take_column_store(&mut self, type_key: InternedKey) -> Option<Arc<ColumnStore>> {
self.column_stores.remove(&type_key)
}
#[inline]
fn clear_column_stores(&mut self) {
self.column_stores.clear();
}
fn set_node_property(&mut self, idx: NodeIndex, key: InternedKey, value: Value) {
self.capture_property_pre_image(idx);
if let Some((type_key, row_id)) = self.columnar_row_of(idx) {
if let Some(store) = self.column_stores.get_mut(&type_key) {
Arc::make_mut(store).set(row_id, key, &value, None);
}
return;
}
if let Some(nd) = self.cow_node(idx) {
nd.properties.insert(key, value);
}
}
fn set_node_property_if_absent(&mut self, idx: NodeIndex, key: InternedKey, value: Value) {
if GraphRead::node_has_property(self, idx, key) {
return;
}
GraphWrite::set_node_property(self, idx, key, value);
}
fn remove_node_property(&mut self, idx: NodeIndex, key: InternedKey) -> Option<Value> {
let previous = GraphRead::get_node_property(self, idx, key);
self.capture_property_pre_image(idx);
if let Some((type_key, row_id)) = self.columnar_row_of(idx) {
if let Some(store) = self.column_stores.get_mut(&type_key) {
Arc::make_mut(store).set(row_id, key, &Value::Null, None);
}
return previous;
}
self.cow_node(idx)?.properties.remove(key)
}
fn clear_node_property(&mut self, idx: NodeIndex, key: InternedKey) -> Option<Value> {
GraphWrite::remove_node_property(self, idx, key)
}
fn replace_node_properties(&mut self, idx: NodeIndex, pairs: Vec<(InternedKey, Value)>) {
self.capture_property_pre_image(idx);
if self.columnar_row_of(idx).is_some() {
for (key, value) in pairs {
GraphWrite::set_node_property(self, idx, key, value);
}
return;
}
if let Some(nd) = self.cow_node(idx) {
nd.properties.replace_all(pairs);
}
}
#[inline]
fn add_node(&mut self, data: NodeData) -> NodeIndex {
let node_type = data.node_type;
let bound_before = self.node_bound();
let idx = NodeIndex::new(bound_before);
self.nodes.insert(idx.index() as u32, data);
self.appended += 1;
self.slot_mirror.note_node_added(bound_before, idx);
if let Some(journal) = self.undo.as_deref_mut() {
journal.note_node_added(idx, node_type);
}
idx
}
fn remove_node(&mut self, _idx: NodeIndex) -> Option<NodeData> {
unreachable!("forked backend must be materialised before remove_node")
}
fn add_edge(&mut self, _a: NodeIndex, _b: NodeIndex, _data: EdgeData) -> EdgeIndex {
unreachable!("forked backend must be materialised before add_edge")
}
fn remove_edge(&mut self, _idx: EdgeIndex) -> Option<EdgeData> {
unreachable!("forked backend must be materialised before remove_edge")
}
}