1use memmap2::Mmap;
2use std::{
3 fs::File,
4 io::{BufWriter, Write},
5 marker::PhantomData,
6 path::Path,
7};
8use tempfile::TempPath;
9
10use super::traits::{PayloadBuilderOps, open_mmap};
11use super::{OwnedPayloadStore, RefPayloadStore};
12use crate::error::Result;
13use crate::persist::move_file;
14
15pub struct PodStoreBuilder<T> {
24 writer: BufWriter<File>,
25 temp_path: TempPath,
26 n_pushed: usize,
27 _marker: PhantomData<T>,
28}
29
30impl<T: bytemuck::Pod> PodStoreBuilder<T> {
31 pub fn new() -> Result<Self> {
32 let named = tempfile::NamedTempFile::new()?;
33 let (file, temp_path) = named.into_parts();
34 Ok(Self {
35 writer: BufWriter::new(file),
36 temp_path,
37 n_pushed: 0,
38 _marker: PhantomData,
39 })
40 }
41
42 fn push_inner(&mut self, payload: T) -> Result<()> {
43 self.writer.write_all(bytemuck::bytes_of(&payload))?;
44 self.n_pushed += 1;
45 Ok(())
46 }
47
48 fn finish_inner(self) -> Result<PodStore<T>> {
49 let Self {
50 writer,
51 temp_path,
52 n_pushed,
53 _marker,
54 } = self;
55 let total = (n_pushed * size_of::<T>()) as u64;
56 writer.into_inner().map_err(|e| e.into_error())?;
57 let mmap = open_mmap(temp_path.as_ref(), total)?;
58 Ok(PodStore {
59 mmap,
60 _marker: PhantomData,
61 })
62 }
63
64 fn finish_persisted_inner(self, payloads_path: &Path) -> Result<PodStore<T>> {
65 let Self {
66 writer,
67 temp_path,
68 n_pushed,
69 _marker,
70 } = self;
71 let total = (n_pushed * size_of::<T>()) as u64;
72 writer.into_inner().map_err(|e| e.into_error())?;
73 move_file(temp_path.as_ref(), payloads_path)?;
76 let mmap = open_mmap(payloads_path, total)?;
77 Ok(PodStore {
78 mmap,
79 _marker: PhantomData,
80 })
81 }
82}
83
84impl<T: bytemuck::Pod> PayloadBuilderOps<T> for PodStoreBuilder<T> {
85 type Store = PodStore<T>;
86
87 fn push(&mut self, payload: T) -> Result<()> {
88 self.push_inner(payload)
89 }
90
91 fn finish(self) -> Result<PodStore<T>> {
92 self.finish_inner()
93 }
94
95 fn finish_persisted(self, payloads_path: &Path, _offsets_path: &Path) -> Result<PodStore<T>> {
97 self.finish_persisted_inner(payloads_path)
98 }
99}
100
101pub struct PodStore<T> {
104 mmap: Option<Mmap>,
105 _marker: PhantomData<T>,
107}
108
109impl<T: bytemuck::Pod> PodStore<T> {
110 pub fn load(payloads_path: &Path) -> Result<Self> {
111 let total_bytes = std::fs::metadata(payloads_path)?.len();
112 let mmap = open_mmap(payloads_path, total_bytes)?;
113 Ok(PodStore {
114 mmap,
115 _marker: PhantomData,
116 })
117 }
118}
119
120impl<T: bytemuck::Pod> OwnedPayloadStore<T> for PodStore<T> {
121 fn fetch_owned(&self, id: u32) -> crate::error::Result<T> {
123 Ok(*self.fetch_ref(id))
124 }
125}
126
127impl<T: bytemuck::Pod> RefPayloadStore<T> for PodStore<T> {
128 fn fetch_ref(&self, id: u32) -> &T {
130 let size = std::mem::size_of::<T>();
131 let offset = id as usize * size;
132 bytemuck::from_bytes(
133 &self.mmap.as_ref().expect("fetch_ref on empty store")[offset..offset + size],
134 )
135 }
136}