Pass in GS as a generic

Author: fabubaker

db4-storage/src/lib.rs

@@ -46,12 +46,12 @@ pub type Extension = NoOpStrategy;
 pub type NS<P> = NodeSegmentView<P>;
 pub type ES<P> = EdgeSegmentView<P>;
 pub type GS = GraphSegmentView;
-pub type Layer<P> = GraphStore<NS<P>, ES<P>, P>;
+pub type Layer<P> = GraphStore<NS<P>, ES<P>, GS, P>;
 
 pub type WalImpl = NoWal;
 pub type GIDResolver = MappingResolver;
 
-pub type ReadLockedLayer<P> = ReadLockedGraphStore<NS<P>, ES<P>, P>;
+pub type ReadLockedLayer<P> = ReadLockedGraphStore<NS<P>, ES<P>, GS, P>;
 pub type ReadLockedNodes<P> = ReadLockedNodeStorage<NS<P>, P>;
 pub type ReadLockedEdges<P> = ReadLockedEdgeStorage<ES<P>, P>;
 

db4-storage/src/pages/mod.rs

@@ -1,12 +1,10 @@
 use crate::{
-    LocalPOS,
-    api::{edges::EdgeSegmentOps, nodes::NodeSegmentOps},
-    error::StorageError,
-    pages::{edge_store::ReadLockedEdgeStorage, node_store::ReadLockedNodeStorage},
+    api::{edges::EdgeSegmentOps, graph::GraphSegmentOps, nodes::NodeSegmentOps},
+    error::StorageError, pages::{edge_store::ReadLockedEdgeStorage, node_store::ReadLockedNodeStorage},
     persist::strategy::{Config, PersistentStrategy},
     properties::props_meta_writer::PropsMetaWriter,
     segments::{edge::segment::MemEdgeSegment, node::segment::MemNodeSegment},
-    GS,
+    LocalPOS,
 };
 use edge_page::writer::EdgeWriter;
 use edge_store::EdgeStorageInner;
@@ -47,7 +45,7 @@ pub mod test_utils;
 // graph // (node/edges) // segment // layer_ids (0, 1, 2, ...) // actual graphy bits
 
 #[derive(Debug)]
-pub struct GraphStore<NS, ES, EXT: Config> {
+pub struct GraphStore<NS, ES, GS, EXT: Config> {
     nodes: Arc<NodeStorageInner<NS, EXT>>,
     edges: Arc<EdgeStorageInner<ES, EXT>>,
     graph: Arc<GraphStorageInner<GS, EXT>>,
@@ -60,20 +58,22 @@ pub struct GraphStore<NS, ES, EXT: Config> {
 pub struct ReadLockedGraphStore<
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
     EXT: Config,
 > {
     pub nodes: Arc<ReadLockedNodeStorage<NS, EXT>>,
     pub edges: Arc<ReadLockedEdgeStorage<ES, EXT>>,
-    pub graph: Arc<GraphStore<NS, ES, EXT>>,
+    pub graph: Arc<GraphStore<NS, ES, GS, EXT>>,
 }
 
 impl<
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
     EXT: PersistentStrategy,
-> GraphStore<NS, ES, EXT>
+> GraphStore<NS, ES, GS, EXT>
 {
-    pub fn read_locked(self: &Arc<Self>) -> ReadLockedGraphStore<NS, ES, EXT> {
+    pub fn read_locked(self: &Arc<Self>) -> ReadLockedGraphStore<NS, ES, GS, EXT> {
         let nodes = self.nodes.locked().into();
         let edges = self.edges.locked().into();
 
@@ -342,7 +342,7 @@ impl<
         src: VID,
         dst: VID,
         e_id: Option<EID>,
-    ) -> WriteSession<'_, NS, ES, EXT> {
+    ) -> WriteSession<'_, NS, ES, GS, EXT> {
         let (src_chunk, _) = self.nodes.resolve_pos(src);
         let (dst_chunk, _) = self.nodes.resolve_pos(dst);
 
@@ -375,7 +375,7 @@ impl<
         src: VID,
         dst: VID,
         e_id: Option<EID>,
-    ) -> WriteSession<'_, NS, ES, EXT> {
+    ) -> WriteSession<'_, NS, ES, GS, EXT> {
         let (src_chunk, _) = self.nodes.resolve_pos(src);
         let (dst_chunk, _) = self.nodes.resolve_pos(dst);
 
@@ -422,7 +422,7 @@ impl<
     }
 }
 
-impl<NS, ES, EXT: Config> Drop for GraphStore<NS, ES, EXT> {
+impl<NS, ES, GS, EXT: Config> Drop for GraphStore<NS, ES, GS, EXT> {
     fn drop(&mut self) {
         let node_types = self.nodes.prop_meta().get_all_node_types();
         self._ext.set_node_types(node_types);

db4-storage/src/pages/session.rs

@@ -2,10 +2,7 @@ use super::{
     GraphStore, edge_page::writer::EdgeWriter, node_page::writer::WriterPair, resolve_pos,
 };
 use crate::{
-    LocalPOS,
-    api::{edges::EdgeSegmentOps, nodes::NodeSegmentOps},
-    persist::strategy::{Config, PersistentStrategy},
-    segments::{edge::segment::MemEdgeSegment, node::segment::MemNodeSegment},
+    api::{edges::EdgeSegmentOps, graph::GraphSegmentOps, nodes::NodeSegmentOps}, persist::strategy::{Config, PersistentStrategy}, segments::{edge::segment::MemEdgeSegment, node::segment::MemNodeSegment}, LocalPOS
 };
 use parking_lot::RwLockWriteGuard;
 use raphtory_api::core::{entities::properties::prop::Prop, storage::dict_mapper::MaybeNew};
@@ -14,23 +11,24 @@ use raphtory_core::{
     storage::timeindex::AsTime,
 };
 
-pub struct WriteSession<'a, NS: NodeSegmentOps, ES: EdgeSegmentOps, EXT: Config> {
+pub struct WriteSession<'a, NS: NodeSegmentOps, ES: EdgeSegmentOps, GS: GraphSegmentOps, EXT: Config> {
     node_writers: WriterPair<'a, RwLockWriteGuard<'a, MemNodeSegment>, NS>,
     edge_writer: Option<EdgeWriter<'a, RwLockWriteGuard<'a, MemEdgeSegment>, ES>>,
-    graph: &'a GraphStore<NS, ES, EXT>,
+    graph: &'a GraphStore<NS, ES, GS, EXT>,
 }
 
 impl<
     'a,
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
     EXT: PersistentStrategy,
-> WriteSession<'a, NS, ES, EXT>
+> WriteSession<'a, NS, ES, GS, EXT>
 {
     pub fn new(
         node_writers: WriterPair<'a, RwLockWriteGuard<'a, MemNodeSegment>, NS>,
         edge_writer: Option<EdgeWriter<'a, RwLockWriteGuard<'a, MemEdgeSegment>, ES>>,
-        graph: &'a GraphStore<NS, ES, EXT>,
+        graph: &'a GraphStore<NS, ES, GS, EXT>,
     ) -> Self {
         Self {
             node_writers,

db4-storage/src/pages/test_utils/checkers.rs

@@ -15,6 +15,7 @@ use crate::{
     api::{
         edges::{EdgeEntryOps, EdgeRefOps, EdgeSegmentOps},
         nodes::{NodeEntryOps, NodeRefOps, NodeSegmentOps},
+        graph::GraphSegmentOps,
     },
     error::StorageError,
     pages::GraphStore,
@@ -26,12 +27,13 @@ use super::fixtures::{AddEdge, Fixture, NodeFixture};
 pub fn check_edges_support<
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
     EXT: PersistentStrategy,
 >(
     edges: Vec<(impl Into<VID>, impl Into<VID>, Option<usize>)>, // src, dst, optional layer_id
     par_load: bool,
     check_load: bool,
-    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, EXT>,
+    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, GS, EXT>,
 ) {
     let mut edges = edges
         .into_iter()
@@ -103,11 +105,12 @@ pub fn check_edges_support<
     fn check<
         NS: NodeSegmentOps<Extension = EXT>,
         ES: EdgeSegmentOps<Extension = EXT>,
+        GS: GraphSegmentOps,
         EXT: PersistentStrategy,
     >(
         stage: &str,
         expected_edges: &[(VID, VID, Option<usize>)], // (src, dst, layer_id)
-        graph: &GraphStore<NS, ES, EXT>,
+        graph: &GraphStore<NS, ES, GS, EXT>,
     ) {
         let nodes = graph.nodes();
         let edges = graph.edges();
@@ -189,7 +192,7 @@ pub fn check_edges_support<
     if check_load {
         drop(graph);
 
-        let maybe_ns = GraphStore::<NS, ES, EXT>::load(graph_dir.path());
+        let maybe_ns = GraphStore::<NS, ES, GS, EXT>::load(graph_dir.path());
 
         match maybe_ns {
             Ok(graph) => {
@@ -206,10 +209,11 @@ pub fn check_graph_with_nodes_support<
     EXT: PersistentStrategy,
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
 >(
     fixture: &NodeFixture,
     check_load: bool,
-    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, EXT>,
+    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, GS, EXT>,
 ) {
     let NodeFixture {
         temp_props,
@@ -234,7 +238,7 @@ pub fn check_graph_with_nodes_support<
 
     let check_fn = |temp_props: &[(VID, i64, Vec<(String, Prop)>)],
                     const_props: &[(VID, Vec<(String, Prop)>)],
-                    graph: &GraphStore<NS, ES, EXT>| {
+                    graph: &GraphStore<NS, ES, GS, EXT>| {
         let mut ts_for_nodes = HashMap::new();
         for (node, t, _) in temp_props {
             ts_for_nodes.entry(*node).or_insert_with(Vec::new).push(*t);
@@ -331,7 +335,7 @@ pub fn check_graph_with_nodes_support<
 
     if check_load {
         drop(graph);
-        let graph = GraphStore::<NS, ES, EXT>::load(graph_dir.path()).unwrap();
+        let graph = GraphStore::<NS, ES, GS, EXT>::load(graph_dir.path()).unwrap();
         check_fn(temp_props, const_props, &graph);
     }
 }
@@ -340,10 +344,11 @@ pub fn check_graph_with_props_support<
     EXT: PersistentStrategy,
     NS: NodeSegmentOps<Extension = EXT>,
     ES: EdgeSegmentOps<Extension = EXT>,
+    GS: GraphSegmentOps,
 >(
     fixture: &Fixture,
     check_load: bool,
-    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, EXT>,
+    make_graph: impl FnOnce(&Path) -> GraphStore<NS, ES, GS, EXT>,
 ) {
     let Fixture { edges, const_props } = fixture;
     let graph_dir = tempfile::tempdir().unwrap();
@@ -374,7 +379,7 @@ pub fn check_graph_with_props_support<
 
     black_box(assert!(graph.edges().num_edges() > 0));
 
-    let check_fn = |edges: &[AddEdge], graph: &GraphStore<NS, ES, EXT>| {
+    let check_fn = |edges: &[AddEdge], graph: &GraphStore<NS, ES, GS, EXT>| {
         let mut edge_groups = HashMap::new();
         let mut node_groups: HashMap<VID, Vec<i64>> = HashMap::new();
 
@@ -485,7 +490,7 @@ pub fn check_graph_with_props_support<
         // Load the graph from disk and check again
         drop(graph);
 
-        let graph = GraphStore::<NS, ES, EXT>::load(graph_dir.path()).unwrap();
+        let graph = GraphStore::<NS, ES, GS, EXT>::load(graph_dir.path()).unwrap();
         black_box(check_fn(edges, &graph));
     }
 }

db4-storage/src/persist/strategy.rs

@@ -3,8 +3,7 @@ use std::ops::DerefMut;
 use serde::{Deserialize, Serialize};
 
 use crate::segments::{
-    edge::segment::{EdgeSegmentView, MemEdgeSegment},
-    node::segment::{MemNodeSegment, NodeSegmentView},
+    edge::segment::{EdgeSegmentView, MemEdgeSegment}, graph::GraphSegmentView, node::segment::{MemNodeSegment, NodeSegmentView}
 };
 
 pub const DEFAULT_MAX_PAGE_LEN_NODES: u32 = 131_072; // 2^17
@@ -23,12 +22,15 @@ pub trait Config:
 pub trait PersistentStrategy: Config {
     type NS;
     type ES;
+    type GS;
+
     fn persist_node_segment<MP: DerefMut<Target = MemNodeSegment>>(
         &self,
         node_page: &Self::NS,
         writer: MP,
     ) where
         Self: Sized;
+
     fn persist_edge_page<MP: DerefMut<Target = MemEdgeSegment>>(
         &self,
         edge_page: &Self::ES,
@@ -86,6 +88,7 @@ impl Config for NoOpStrategy {
 impl PersistentStrategy for NoOpStrategy {
     type ES = EdgeSegmentView<Self>;
     type NS = NodeSegmentView<Self>;
+    type GS = GraphSegmentView;
 
     fn persist_node_segment<MP: DerefMut<Target = MemNodeSegment>>(
         &self,

raphtory-storage/src/mutation/addition_ops_ext.rs

@@ -15,28 +15,24 @@ use raphtory_core::{
     storage::timeindex::TimeIndexEntry,
 };
 use storage::{
-    pages::{node_page::writer::node_info_as_props, session::WriteSession},
-    persist::strategy::PersistentStrategy,
-    properties::props_meta_writer::PropsMetaWriter,
-    resolver::GIDResolverOps,
-    Extension, WalImpl, ES, NS,
+    pages::{node_page::writer::node_info_as_props, session::WriteSession}, persist::strategy::PersistentStrategy, properties::props_meta_writer::PropsMetaWriter, resolver::GIDResolverOps, Extension, WalImpl, ES, GS, NS
 };
 
 use crate::mutation::{
     addition_ops::{EdgeWriteLock, InternalAdditionOps, SessionAdditionOps},
     MutationError,
 };
 
-pub struct WriteS<'a, EXT: PersistentStrategy<NS = NS<EXT>, ES = ES<EXT>>> {
-    static_session: WriteSession<'a, NS<EXT>, ES<EXT>, EXT>,
+pub struct WriteS<'a, EXT: PersistentStrategy<NS = NS<EXT>, ES = ES<EXT>, GS = GS>> {
+    static_session: WriteSession<'a, NS<EXT>, ES<EXT>, GS, EXT>,
 }
 
 #[derive(Clone, Copy, Debug)]
 pub struct UnlockedSession<'a> {
     graph: &'a TemporalGraph<Extension>,
 }
 
-impl<'a, EXT: PersistentStrategy<NS = NS<EXT>, ES = ES<EXT>>> EdgeWriteLock for WriteS<'a, EXT> {
+impl<'a, EXT: PersistentStrategy<NS = NS<EXT>, ES = ES<EXT>, GS = GS>> EdgeWriteLock for WriteS<'a, EXT> {
     fn internal_add_static_edge(
         &mut self,
         src: impl Into<VID>,