Add PrimaryKeyGenerator to DataBucket

src/page/mod.rs

@@ -1,6 +1,7 @@
 mod data;
 mod header;
 mod index;
+mod primary_key;
 mod space_info;
 mod ty;
 mod util;
@@ -11,6 +12,7 @@ use rkyv::{Archive, Deserialize, Serialize};
 pub use data::Data;
 pub use header::{GeneralHeader, DATA_VERSION};
 pub use index::{map_tree_index, map_unique_tree_index, IndexPage};
+pub use primary_key::{TablePrimaryKey, PrimaryKeyGenerator, PrimaryKeyGeneratorState};
 pub use space_info::{Interval, SpaceInfo};
 pub use ty::PageType;
 pub use util::{

src/page/primary_key.rs

@@ -0,0 +1,112 @@
+use std::sync::atomic::{AtomicI64, AtomicU32, AtomicU64, Ordering};
+
+pub trait TablePrimaryKey {
+    type Generator;
+}
+
+pub trait PrimaryKeyGenerator<T> {
+    fn next(&self) -> T;
+}
+
+pub trait PrimaryKeyGeneratorState {
+    type State;
+
+    fn get_state(&self) -> Self::State;
+
+    fn from_state(state: Self::State) -> Self;
+}
+
+impl<T> PrimaryKeyGenerator<T> for AtomicU32
+where
+    T: From<u32>,
+{
+    fn next(&self) -> T {
+        self.fetch_add(1, Ordering::Relaxed).into()
+    }
+}
+
+impl PrimaryKeyGeneratorState for AtomicU32 {
+    type State = u32;
+
+    fn get_state(&self) -> Self::State {
+        self.load(Ordering::Relaxed)
+    }
+
+    fn from_state(state: Self::State) -> Self {
+        AtomicU32::from(state)
+    }
+}
+
+impl<T> PrimaryKeyGenerator<T> for AtomicU64
+where
+    T: From<u64>,
+{
+    fn next(&self) -> T {
+        self.fetch_add(1, Ordering::Relaxed).into()
+    }
+}
+
+impl PrimaryKeyGeneratorState for AtomicU64 {
+    type State = u64;
+
+    fn get_state(&self) -> Self::State {
+        self.load(Ordering::Relaxed)
+    }
+
+    fn from_state(state: Self::State) -> Self {
+        AtomicU64::from(state)
+    }
+}
+
+impl<T> PrimaryKeyGenerator<T> for AtomicI64
+where
+    T: From<i64>,
+{
+    fn next(&self) -> T {
+        self.fetch_add(1, Ordering::Relaxed).into()
+    }
+}
+
+impl PrimaryKeyGeneratorState for AtomicI64 {
+    type State = i64;
+
+    fn get_state(&self) -> Self::State {
+        self.load(Ordering::Relaxed)
+    }
+
+    fn from_state(state: Self::State) -> Self {
+        AtomicI64::from(state)
+    }
+}
+
+impl PrimaryKeyGeneratorState for () {
+    type State = ();
+
+    fn get_state(&self) -> Self::State {
+        ()
+    }
+
+    fn from_state((): Self::State) -> Self {
+        ()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::atomic::{AtomicI64, AtomicU64};
+    use super::PrimaryKeyGenerator;
+
+    #[test]
+    fn test_pk_gen_state_atomic_u64() {
+        let mut state = AtomicU64::new(0);
+        assert_eq!(<AtomicU64 as PrimaryKeyGenerator<u64>>::next(&mut state), 0);
+        assert_eq!(<AtomicU64 as PrimaryKeyGenerator<u64>>::next(&mut state), 1);
+    }
+
+    #[test]
+    fn test_pk_gen_state_atomic_i64() {
+        let mut state = AtomicI64::new(0);
+        assert_eq!(<AtomicI64 as PrimaryKeyGenerator<i64>>::next(&mut state), 0);
+        assert_eq!(<AtomicI64 as PrimaryKeyGenerator<i64>>::next(&mut state), 1);
+    }
+}

src/page/space_info.rs

@@ -8,14 +8,18 @@ use crate::page::INNER_PAGE_SIZE;
 use crate::util::Persistable;
 use crate::{space, Link};
 
+use super::PrimaryKeyGeneratorState;
+
 pub type SpaceName = String;
 
-// TODO: Minor. Add some schema description in `SpaceIndo`
+// TODO: Minor. Add some schema description in `SpaceInfo`
 
 /// Internal information about a `Space`. Always appears first before all other
 /// pages in a `Space`.
 #[derive(Archive, Clone, Deserialize, Debug, PartialEq, Serialize)]
-pub struct SpaceInfo<Pk = ()> {
+pub struct SpaceInfo<Pk = ()>
+where Pk: PrimaryKeyGeneratorState
+{
     pub id: space::Id,
     pub page_count: u32,
     pub name: SpaceName,
@@ -32,7 +36,7 @@ pub struct Interval(pub usize, pub usize);
 
 impl<Pk> Persistable for SpaceInfo<Pk>
 where
-    Pk: Archive + Serialize<AllocSerializer<{ INNER_PAGE_SIZE }>>,
+    Pk: Archive + Serialize<AllocSerializer<{ INNER_PAGE_SIZE }>> + PrimaryKeyGeneratorState,
 {
     fn as_bytes(&self) -> impl AsRef<[u8]> {
         rkyv::to_bytes::<_, { INNER_PAGE_SIZE }>(self).unwrap()
@@ -42,8 +46,9 @@ where
 #[cfg(test)]
 mod test {
     use std::collections::HashMap;
+    use std::sync::atomic::AtomicU64;
 
-    use crate::page::{SpaceInfo, INNER_PAGE_SIZE};
+    use crate::page::{PrimaryKeyGenerator, SpaceInfo, INNER_PAGE_SIZE};
     use crate::util::Persistable;
 
     #[test]
@@ -61,4 +66,19 @@ mod test {
         let bytes = info.as_bytes();
         assert!(bytes.as_ref().len() < INNER_PAGE_SIZE)
     }
+
+    #[test]
+    fn test_pk_gen_state() {
+        let info = SpaceInfo {
+            id: 0.into(),
+            page_count: 0,
+            name: "Test".to_string(),
+            primary_key_intervals: vec![],
+            secondary_index_intervals: HashMap::new(),
+            data_intervals: vec![],
+            pk_gen_state: AtomicU64::new(0),
+            empty_links_list: vec![],
+        };
+        assert_eq!(<AtomicU64 as PrimaryKeyGenerator<u64>>::next(&info.pk_gen_state), 0);
+    }
 }