tdb_succinct/storage/
types.rs1use async_trait::async_trait;
2use bytes::{Buf, Bytes};
3use tokio::io::{self, AsyncRead, AsyncWrite, AsyncWriteExt};
4
5use crate::{AdjacencyList, BitIndex};
6
7#[async_trait]
8pub trait SyncableFile: AsyncWrite + Unpin + Send {
9 async fn sync_all(self) -> io::Result<()>;
10}
11
12#[async_trait]
13pub trait FileStore: Clone + Send + Sync {
14 type Write: SyncableFile;
15 async fn open_write(&self) -> io::Result<Self::Write>;
16}
17
18#[async_trait]
19pub trait FileLoad: Clone + Send + Sync {
20 type Read: AsyncRead + Unpin + Send;
21
22 async fn exists(&self) -> io::Result<bool>;
23 async fn size(&self) -> io::Result<usize>;
24 async fn open_read(&self) -> io::Result<Self::Read> {
25 self.open_read_from(0).await
26 }
27 async fn open_read_from(&self, offset: usize) -> io::Result<Self::Read>;
28 async fn map(&self) -> io::Result<Bytes>;
29
30 async fn map_if_exists(&self) -> io::Result<Option<Bytes>> {
31 match self.exists().await? {
32 false => Ok(None),
33 true => {
34 let mapped = self.map().await?;
35 Ok(Some(mapped))
36 }
37 }
38 }
39}
40
41#[derive(Clone)]
42pub struct TypedDictionaryMaps {
43 pub types_present_map: Bytes,
44 pub type_offsets_map: Bytes,
45 pub blocks_map: Bytes,
46 pub offsets_map: Bytes,
47}
48
49#[derive(Clone)]
50pub struct TypedDictionaryFiles<F: 'static + FileLoad + FileStore> {
51 pub types_present_file: F,
52 pub type_offsets_file: F,
53 pub blocks_file: F,
54 pub offsets_file: F,
55}
56
57impl<F: 'static + FileLoad + FileStore> TypedDictionaryFiles<F> {
58 pub async fn map_all(&self) -> io::Result<TypedDictionaryMaps> {
59 let types_present_map = self.types_present_file.map().await?;
60 let type_offsets_map = self.type_offsets_file.map().await?;
61 let offsets_map = self.offsets_file.map().await?;
62 let blocks_map = self.blocks_file.map().await?;
63
64 Ok(TypedDictionaryMaps {
65 types_present_map,
66 type_offsets_map,
67 offsets_map,
68 blocks_map,
69 })
70 }
71
72 pub async fn write_all_from_bufs<B1: Buf, B2: Buf, B3: Buf, B4: Buf>(
73 &self,
74 types_present_buf: &mut B1,
75 type_offsets_buf: &mut B2,
76 offsets_buf: &mut B3,
77 blocks_buf: &mut B4,
78 ) -> io::Result<()> {
79 let mut types_present_writer = self.types_present_file.open_write().await?;
80 let mut type_offsets_writer = self.type_offsets_file.open_write().await?;
81 let mut offsets_writer = self.offsets_file.open_write().await?;
82 let mut blocks_writer = self.blocks_file.open_write().await?;
83
84 types_present_writer
85 .write_all_buf(types_present_buf)
86 .await?;
87 type_offsets_writer.write_all_buf(type_offsets_buf).await?;
88 offsets_writer.write_all_buf(offsets_buf).await?;
89 blocks_writer.write_all_buf(blocks_buf).await?;
90
91 types_present_writer.flush().await?;
92 types_present_writer.sync_all().await?;
93
94 type_offsets_writer.flush().await?;
95 type_offsets_writer.sync_all().await?;
96
97 offsets_writer.flush().await?;
98 offsets_writer.sync_all().await?;
99
100 blocks_writer.flush().await?;
101 blocks_writer.sync_all().await?;
102
103 Ok(())
104 }
105}
106
107#[derive(Clone)]
108pub struct DictionaryMaps {
109 pub blocks_map: Bytes,
110 pub offsets_map: Bytes,
111}
112
113#[derive(Clone)]
114pub struct DictionaryFiles<F: 'static + FileLoad + FileStore> {
115 pub blocks_file: F,
116 pub offsets_file: F,
117 }
119
120impl<F: 'static + FileLoad + FileStore> DictionaryFiles<F> {
121 pub async fn map_all(&self) -> io::Result<DictionaryMaps> {
122 let offsets_map = self.offsets_file.map().await?;
123 let blocks_map = self.blocks_file.map().await?;
124
125 Ok(DictionaryMaps {
126 offsets_map,
127 blocks_map,
128 })
129 }
130
131 pub async fn write_all_from_bufs<B1: Buf, B2: Buf>(
132 &self,
133 blocks_buf: &mut B1,
134 offsets_buf: &mut B2,
135 ) -> io::Result<()> {
136 let mut offsets_writer = self.offsets_file.open_write().await?;
137 let mut blocks_writer = self.blocks_file.open_write().await?;
138
139 offsets_writer.write_all_buf(offsets_buf).await?;
140 blocks_writer.write_all_buf(blocks_buf).await?;
141
142 offsets_writer.flush().await?;
143 offsets_writer.sync_all().await?;
144
145 blocks_writer.flush().await?;
146 blocks_writer.sync_all().await?;
147
148 Ok(())
149 }
150}
151
152#[derive(Clone)]
153pub struct BitIndexMaps {
154 pub bits_map: Bytes,
155 pub blocks_map: Bytes,
156 pub sblocks_map: Bytes,
157}
158
159impl Into<BitIndex> for BitIndexMaps {
160 fn into(self) -> BitIndex {
161 BitIndex::from_maps(self.bits_map, self.blocks_map, self.sblocks_map)
162 }
163}
164
165#[derive(Clone)]
166pub struct BitIndexFiles<F: 'static + FileLoad> {
167 pub bits_file: F,
168 pub blocks_file: F,
169 pub sblocks_file: F,
170}
171
172impl<F: 'static + FileLoad + FileStore> BitIndexFiles<F> {
173 pub async fn map_all(&self) -> io::Result<BitIndexMaps> {
174 let bits_map = self.bits_file.map().await?;
175 let blocks_map = self.blocks_file.map().await?;
176 let sblocks_map = self.sblocks_file.map().await?;
177
178 Ok(BitIndexMaps {
179 bits_map,
180 blocks_map,
181 sblocks_map,
182 })
183 }
184
185 pub async fn map_all_if_exists(&self) -> io::Result<Option<BitIndexMaps>> {
186 if self.bits_file.exists().await? {
187 Ok(Some(self.map_all().await?))
188 } else {
189 Ok(None)
190 }
191 }
192}
193
194#[derive(Clone)]
195pub struct AdjacencyListMaps {
196 pub bitindex_maps: BitIndexMaps,
197 pub nums_map: Bytes,
198}
199
200impl Into<AdjacencyList> for AdjacencyListMaps {
201 fn into(self) -> AdjacencyList {
202 AdjacencyList::parse(
203 self.nums_map,
204 self.bitindex_maps.bits_map,
205 self.bitindex_maps.blocks_map,
206 self.bitindex_maps.sblocks_map,
207 )
208 }
209}
210
211#[derive(Clone)]
212pub struct AdjacencyListFiles<F: 'static + FileLoad> {
213 pub bitindex_files: BitIndexFiles<F>,
214 pub nums_file: F,
215}
216
217impl<F: 'static + FileLoad + FileStore> AdjacencyListFiles<F> {
218 pub async fn map_all(&self) -> io::Result<AdjacencyListMaps> {
219 let bitindex_maps = self.bitindex_files.map_all().await?;
220 let nums_map = self.nums_file.map().await?;
221
222 Ok(AdjacencyListMaps {
223 bitindex_maps,
224 nums_map,
225 })
226 }
227}
228
229pub async fn copy_file<F1: FileLoad, F2: FileStore>(f1: &F1, f2: &F2) -> io::Result<()> {
230 if !f1.exists().await? {
231 return Ok(());
232 }
233 let mut input = f1.open_read().await?;
234 let mut output = f2.open_write().await?;
235
236 tokio::io::copy(&mut input, &mut output).await?;
237 output.flush().await?;
238 output.sync_all().await?;
239
240 Ok(())
241}
242
243impl<F1: 'static + FileLoad + FileStore> DictionaryFiles<F1> {
244 pub async fn copy_from<F2: 'static + FileLoad + FileStore>(
245 &self,
246 from: &DictionaryFiles<F2>,
247 ) -> io::Result<()> {
248 copy_file(&from.blocks_file, &self.blocks_file).await?;
249 copy_file(&from.offsets_file, &self.offsets_file).await?;
250
251 Ok(())
252 }
253}
254
255impl<F1: 'static + FileLoad + FileStore> TypedDictionaryFiles<F1> {
256 pub async fn copy_from<F2: 'static + FileLoad + FileStore>(
257 &self,
258 from: &TypedDictionaryFiles<F2>,
259 ) -> io::Result<()> {
260 copy_file(&from.types_present_file, &self.types_present_file).await?;
261 copy_file(&from.type_offsets_file, &self.type_offsets_file).await?;
262 copy_file(&from.blocks_file, &self.blocks_file).await?;
263 copy_file(&from.offsets_file, &self.offsets_file).await?;
264
265 Ok(())
266 }
267}
268
269impl<F1: 'static + FileLoad + FileStore> BitIndexFiles<F1> {
270 pub async fn copy_from<F2: 'static + FileLoad + FileStore>(
271 &self,
272 from: &BitIndexFiles<F2>,
273 ) -> io::Result<()> {
274 copy_file(&from.bits_file, &self.bits_file).await?;
275 copy_file(&from.blocks_file, &self.blocks_file).await?;
276 copy_file(&from.sblocks_file, &self.sblocks_file).await?;
277
278 Ok(())
279 }
280}
281impl<F1: 'static + FileLoad + FileStore> AdjacencyListFiles<F1> {
282 pub async fn copy_from<F2: 'static + FileLoad + FileStore>(
283 &self,
284 from: &AdjacencyListFiles<F2>,
285 ) -> io::Result<()> {
286 copy_file(&from.nums_file, &self.nums_file).await?;
287 self.bitindex_files.copy_from(&from.bitindex_files).await?;
288
289 Ok(())
290 }
291}