1#![cfg_attr(
5 not(any(test, feature = "proptest-impl", feature = "shielded-scan")),
6 allow(dead_code)
7)]
8
9use std::{
10 any::type_name,
11 collections::{BTreeMap, HashMap},
12 fmt::Debug,
13 hash::Hash,
14 marker::PhantomData,
15 ops::RangeBounds,
16};
17
18use crate::service::finalized_state::{DiskWriteBatch, FromDisk, IntoDisk, ReadDisk, WriteDisk};
19
20use super::DiskDb;
21
22#[derive(Clone)]
28pub struct TypedColumnFamily<'cf, Key, Value>
29where
30 Key: IntoDisk + FromDisk + Debug,
31 Value: IntoDisk + FromDisk,
32{
33 db: DiskDb,
35
36 cf: rocksdb::ColumnFamilyRef<'cf>,
38
39 _cf_name: String,
41
42 _marker: PhantomData<(Key, Value)>,
44}
45
46#[must_use = "batches must be written to the database"]
54#[derive(Debug, Eq, PartialEq)]
55pub struct WriteTypedBatch<'cf, Key, Value, Batch>
56where
57 Key: IntoDisk + FromDisk + Debug,
58 Value: IntoDisk + FromDisk,
59 Batch: WriteDisk,
60{
61 inner: TypedColumnFamily<'cf, Key, Value>,
62
63 batch: Batch,
64}
65
66impl<Key, Value> Debug for TypedColumnFamily<'_, Key, Value>
67where
68 Key: IntoDisk + FromDisk + Debug,
69 Value: IntoDisk + FromDisk,
70{
71 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72 f.debug_struct(&format!(
73 "TypedColumnFamily<{}, {}>",
74 type_name::<Key>(),
75 type_name::<Value>()
76 ))
77 .field("db", &self.db)
78 .field("cf", &self._cf_name)
79 .finish()
80 }
81}
82
83impl<Key, Value> PartialEq for TypedColumnFamily<'_, Key, Value>
84where
85 Key: IntoDisk + FromDisk + Debug,
86 Value: IntoDisk + FromDisk,
87{
88 fn eq(&self, other: &Self) -> bool {
89 self.db == other.db && self._cf_name == other._cf_name
90 }
91}
92
93impl<Key, Value> Eq for TypedColumnFamily<'_, Key, Value>
94where
95 Key: IntoDisk + FromDisk + Debug,
96 Value: IntoDisk + FromDisk,
97{
98}
99
100impl<'cf, Key, Value> TypedColumnFamily<'cf, Key, Value>
101where
102 Key: IntoDisk + FromDisk + Debug,
103 Value: IntoDisk + FromDisk,
104{
105 pub fn new(db: &'cf DiskDb, cf_name: &str) -> Option<Self> {
109 let cf = db.cf_handle(cf_name)?;
110
111 Some(Self {
112 db: db.clone(),
113 cf,
114 _cf_name: cf_name.to_string(),
115 _marker: PhantomData,
116 })
117 }
118
119 pub fn new_batch_for_writing(self) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
126 WriteTypedBatch {
127 inner: self,
128 batch: DiskWriteBatch::new(),
129 }
130 }
131
132 pub fn take_batch_for_writing(
137 self,
138 batch: DiskWriteBatch,
139 ) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
140 WriteTypedBatch { inner: self, batch }
141 }
142
143 pub fn with_batch_for_writing(
148 self,
149 batch: &mut DiskWriteBatch,
150 ) -> WriteTypedBatch<'cf, Key, Value, &mut DiskWriteBatch> {
151 WriteTypedBatch { inner: self, batch }
152 }
153
154 pub fn zs_is_empty(&self) -> bool {
158 self.db.zs_is_empty(&self.cf)
159 }
160
161 pub fn zs_get(&self, key: &Key) -> Option<Value> {
163 self.db.zs_get(&self.cf, key)
164 }
165
166 pub fn zs_contains(&self, key: &Key) -> bool {
168 self.db.zs_contains(&self.cf, key)
169 }
170
171 pub fn zs_first_key_value(&self) -> Option<(Key, Value)> {
175 self.db.zs_first_key_value(&self.cf)
176 }
177
178 pub fn zs_last_key_value(&self) -> Option<(Key, Value)> {
182 self.db.zs_last_key_value(&self.cf)
183 }
184
185 pub fn zs_next_key_value_from(&self, lower_bound: &Key) -> Option<(Key, Value)> {
190 self.db.zs_next_key_value_from(&self.cf, lower_bound)
191 }
192
193 pub fn zs_next_key_value_strictly_after(&self, lower_bound: &Key) -> Option<(Key, Value)> {
198 self.db
199 .zs_next_key_value_strictly_after(&self.cf, lower_bound)
200 }
201
202 pub fn zs_prev_key_value_back_from(&self, upper_bound: &Key) -> Option<(Key, Value)> {
207 self.db.zs_prev_key_value_back_from(&self.cf, upper_bound)
208 }
209
210 pub fn zs_prev_key_value_strictly_before(&self, upper_bound: &Key) -> Option<(Key, Value)> {
215 self.db
216 .zs_prev_key_value_strictly_before(&self.cf, upper_bound)
217 }
218
219 pub fn zs_forward_range_iter<Range>(
223 &self,
224 range: Range,
225 ) -> impl Iterator<Item = (Key, Value)> + '_
226 where
227 Range: RangeBounds<Key>,
228 {
229 self.db.zs_forward_range_iter(&self.cf, range)
230 }
231
232 pub fn zs_reverse_range_iter<Range>(
236 &self,
237 range: Range,
238 ) -> impl Iterator<Item = (Key, Value)> + '_
239 where
240 Range: RangeBounds<Key>,
241 {
242 self.db.zs_reverse_range_iter(&self.cf, range)
243 }
244}
245
246impl<Key, Value> TypedColumnFamily<'_, Key, Value>
247where
248 Key: IntoDisk + FromDisk + Debug + Ord,
249 Value: IntoDisk + FromDisk,
250{
251 pub fn zs_items_in_range_ordered<Range>(&self, range: Range) -> BTreeMap<Key, Value>
255 where
256 Range: RangeBounds<Key>,
257 {
258 self.db.zs_items_in_range_ordered(&self.cf, range)
259 }
260}
261
262impl<Key, Value> TypedColumnFamily<'_, Key, Value>
263where
264 Key: IntoDisk + FromDisk + Debug + Hash + Eq,
265 Value: IntoDisk + FromDisk,
266{
267 pub fn zs_items_in_range_unordered<Range>(&self, range: Range) -> HashMap<Key, Value>
271 where
272 Range: RangeBounds<Key>,
273 {
274 self.db.zs_items_in_range_unordered(&self.cf, range)
275 }
276}
277
278impl<Key, Value, Batch> WriteTypedBatch<'_, Key, Value, Batch>
279where
280 Key: IntoDisk + FromDisk + Debug,
281 Value: IntoDisk + FromDisk,
282 Batch: WriteDisk,
283{
284 pub fn zs_insert(mut self, key: &Key, value: &Value) -> Self {
289 self.batch.zs_insert(&self.inner.cf, key, value);
290
291 self
292 }
293
294 pub fn zs_delete(mut self, key: &Key) -> Self {
296 self.batch.zs_delete(&self.inner.cf, key);
297
298 self
299 }
300
301 pub fn zs_delete_range(mut self, from: &Key, until_strictly_before: &Key) -> Self {
307 self.batch
308 .zs_delete_range(&self.inner.cf, from, until_strictly_before);
309
310 self
311 }
312}
313
314impl<Key, Value> WriteTypedBatch<'_, Key, Value, DiskWriteBatch>
316where
317 Key: IntoDisk + FromDisk + Debug,
318 Value: IntoDisk + FromDisk,
319{
320 pub fn write_batch(self) -> Result<(), rocksdb::Error> {
325 self.inner.db.write(self.batch)
326 }
327}