tdb_succinct/storage/
types.rs

1use 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    //    pub map_files: Option<BitIndexFiles<F>>
118}
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}