1#![cfg_attr(not(any(test, feature = "proptest-impl")), allow(dead_code))]
5
6use std::{
7 any::type_name,
8 collections::{BTreeMap, HashMap},
9 fmt::Debug,
10 hash::Hash,
11 marker::PhantomData,
12 ops::RangeBounds,
13};
14
15use crate::service::finalized_state::{DiskWriteBatch, FromDisk, IntoDisk, ReadDisk, WriteDisk};
16
17use super::DiskDb;
18
19#[derive(Clone)]
25pub struct TypedColumnFamily<'cf, Key, Value>
26where
27 Key: IntoDisk + FromDisk + Debug,
28 Value: IntoDisk + FromDisk,
29{
30 db: DiskDb,
32
33 cf: rocksdb::ColumnFamilyRef<'cf>,
35
36 _cf_name: String,
38
39 _marker: PhantomData<(Key, Value)>,
41}
42
43#[must_use = "batches must be written to the database"]
51#[derive(Debug, Eq, PartialEq)]
52pub struct WriteTypedBatch<'cf, Key, Value, Batch>
53where
54 Key: IntoDisk + FromDisk + Debug,
55 Value: IntoDisk + FromDisk,
56 Batch: WriteDisk,
57{
58 inner: TypedColumnFamily<'cf, Key, Value>,
59
60 batch: Batch,
61}
62
63impl<Key, Value> Debug for TypedColumnFamily<'_, Key, Value>
64where
65 Key: IntoDisk + FromDisk + Debug,
66 Value: IntoDisk + FromDisk,
67{
68 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69 f.debug_struct(&format!(
70 "TypedColumnFamily<{}, {}>",
71 type_name::<Key>(),
72 type_name::<Value>()
73 ))
74 .field("db", &self.db)
75 .field("cf", &self._cf_name)
76 .finish()
77 }
78}
79
80impl<Key, Value> PartialEq for TypedColumnFamily<'_, Key, Value>
81where
82 Key: IntoDisk + FromDisk + Debug,
83 Value: IntoDisk + FromDisk,
84{
85 fn eq(&self, other: &Self) -> bool {
86 self.db == other.db && self._cf_name == other._cf_name
87 }
88}
89
90impl<Key, Value> Eq for TypedColumnFamily<'_, Key, Value>
91where
92 Key: IntoDisk + FromDisk + Debug,
93 Value: IntoDisk + FromDisk,
94{
95}
96
97impl<'cf, Key, Value> TypedColumnFamily<'cf, Key, Value>
98where
99 Key: IntoDisk + FromDisk + Debug,
100 Value: IntoDisk + FromDisk,
101{
102 pub fn new(db: &'cf DiskDb, cf_name: &str) -> Option<Self> {
106 let cf = db.cf_handle(cf_name)?;
107
108 Some(Self {
109 db: db.clone(),
110 cf,
111 _cf_name: cf_name.to_string(),
112 _marker: PhantomData,
113 })
114 }
115
116 pub fn new_batch_for_writing(self) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
123 WriteTypedBatch {
124 inner: self,
125 batch: DiskWriteBatch::new(),
126 }
127 }
128
129 pub fn take_batch_for_writing(
134 self,
135 batch: DiskWriteBatch,
136 ) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
137 WriteTypedBatch { inner: self, batch }
138 }
139
140 pub fn with_batch_for_writing(
145 self,
146 batch: &mut DiskWriteBatch,
147 ) -> WriteTypedBatch<'cf, Key, Value, &mut DiskWriteBatch> {
148 WriteTypedBatch { inner: self, batch }
149 }
150
151 pub fn zs_is_empty(&self) -> bool {
155 self.db.zs_is_empty(&self.cf)
156 }
157
158 pub fn zs_get(&self, key: &Key) -> Option<Value> {
160 self.db.zs_get(&self.cf, key)
161 }
162
163 pub fn zs_contains(&self, key: &Key) -> bool {
165 self.db.zs_contains(&self.cf, key)
166 }
167
168 pub fn zs_first_key_value(&self) -> Option<(Key, Value)> {
172 self.db.zs_first_key_value(&self.cf)
173 }
174
175 pub fn zs_last_key_value(&self) -> Option<(Key, Value)> {
179 self.db.zs_last_key_value(&self.cf)
180 }
181
182 pub fn zs_next_key_value_from(&self, lower_bound: &Key) -> Option<(Key, Value)> {
187 self.db.zs_next_key_value_from(&self.cf, lower_bound)
188 }
189
190 pub fn zs_next_key_value_strictly_after(&self, lower_bound: &Key) -> Option<(Key, Value)> {
195 self.db
196 .zs_next_key_value_strictly_after(&self.cf, lower_bound)
197 }
198
199 pub fn zs_prev_key_value_back_from(&self, upper_bound: &Key) -> Option<(Key, Value)> {
204 self.db.zs_prev_key_value_back_from(&self.cf, upper_bound)
205 }
206
207 pub fn zs_prev_key_value_strictly_before(&self, upper_bound: &Key) -> Option<(Key, Value)> {
212 self.db
213 .zs_prev_key_value_strictly_before(&self.cf, upper_bound)
214 }
215
216 pub fn zs_forward_range_iter<Range>(
220 &self,
221 range: Range,
222 ) -> impl Iterator<Item = (Key, Value)> + '_
223 where
224 Range: RangeBounds<Key>,
225 {
226 self.db.zs_forward_range_iter(&self.cf, range)
227 }
228
229 pub fn zs_reverse_range_iter<Range>(
233 &self,
234 range: Range,
235 ) -> impl Iterator<Item = (Key, Value)> + '_
236 where
237 Range: RangeBounds<Key>,
238 {
239 self.db.zs_reverse_range_iter(&self.cf, range)
240 }
241}
242
243impl<Key, Value> TypedColumnFamily<'_, Key, Value>
244where
245 Key: IntoDisk + FromDisk + Debug + Ord,
246 Value: IntoDisk + FromDisk,
247{
248 pub fn zs_items_in_range_ordered<Range>(&self, range: Range) -> BTreeMap<Key, Value>
252 where
253 Range: RangeBounds<Key>,
254 {
255 self.db.zs_items_in_range_ordered(&self.cf, range)
256 }
257}
258
259impl<Key, Value> TypedColumnFamily<'_, Key, Value>
260where
261 Key: IntoDisk + FromDisk + Debug + Hash + Eq,
262 Value: IntoDisk + FromDisk,
263{
264 pub fn zs_items_in_range_unordered<Range>(&self, range: Range) -> HashMap<Key, Value>
268 where
269 Range: RangeBounds<Key>,
270 {
271 self.db.zs_items_in_range_unordered(&self.cf, range)
272 }
273}
274
275impl<Key, Value, Batch> WriteTypedBatch<'_, Key, Value, Batch>
276where
277 Key: IntoDisk + FromDisk + Debug,
278 Value: IntoDisk + FromDisk,
279 Batch: WriteDisk,
280{
281 pub fn zs_insert(mut self, key: &Key, value: &Value) -> Self {
286 self.batch.zs_insert(&self.inner.cf, key, value);
287
288 self
289 }
290
291 pub fn zs_delete(mut self, key: &Key) -> Self {
293 self.batch.zs_delete(&self.inner.cf, key);
294
295 self
296 }
297
298 pub fn zs_delete_range(mut self, from: &Key, until_strictly_before: &Key) -> Self {
304 self.batch
305 .zs_delete_range(&self.inner.cf, from, until_strictly_before);
306
307 self
308 }
309}
310
311impl<Key, Value> WriteTypedBatch<'_, Key, Value, DiskWriteBatch>
313where
314 Key: IntoDisk + FromDisk + Debug,
315 Value: IntoDisk + FromDisk,
316{
317 pub fn write_batch(self) -> Result<(), rocksdb::Error> {
322 self.inner.db.write(self.batch)
323 }
324}