1use smallvec::SmallVec;
12use std::io;
13
14mod io_stats;
15
16pub const PREFIX_LEN: usize = 12;
18
19pub type DBValue = Vec<u8>;
21pub type DBKey = SmallVec<[u8; 32]>;
23
24pub use io_stats::{IoStats, Kind as IoStatsKind};
25
26#[derive(Default, Clone, PartialEq)]
28pub struct DBTransaction {
29 pub ops: Vec<DBOp>,
31}
32
33#[derive(Clone, PartialEq)]
35pub enum DBOp {
36 Insert { col: u32, key: DBKey, value: DBValue },
37 Delete { col: u32, key: DBKey },
38 DeletePrefix { col: u32, prefix: DBKey },
39}
40
41impl DBOp {
42 pub fn key(&self) -> &[u8] {
44 match *self {
45 DBOp::Insert { ref key, .. } => key,
46 DBOp::Delete { ref key, .. } => key,
47 DBOp::DeletePrefix { ref prefix, .. } => prefix,
48 }
49 }
50
51 pub fn col(&self) -> u32 {
53 match *self {
54 DBOp::Insert { col, .. } => col,
55 DBOp::Delete { col, .. } => col,
56 DBOp::DeletePrefix { col, .. } => col,
57 }
58 }
59}
60
61impl DBTransaction {
62 pub fn new() -> DBTransaction {
64 DBTransaction::with_capacity(256)
65 }
66
67 pub fn with_capacity(cap: usize) -> DBTransaction {
69 DBTransaction { ops: Vec::with_capacity(cap) }
70 }
71
72 pub fn put(&mut self, col: u32, key: &[u8], value: &[u8]) {
74 self.ops.push(DBOp::Insert { col, key: DBKey::from_slice(key), value: value.to_vec() })
75 }
76
77 pub fn put_vec(&mut self, col: u32, key: &[u8], value: Vec<u8>) {
79 self.ops.push(DBOp::Insert { col, key: DBKey::from_slice(key), value });
80 }
81
82 pub fn delete(&mut self, col: u32, key: &[u8]) {
84 self.ops.push(DBOp::Delete { col, key: DBKey::from_slice(key) });
85 }
86
87 pub fn delete_prefix(&mut self, col: u32, prefix: &[u8]) {
91 self.ops.push(DBOp::DeletePrefix { col, prefix: DBKey::from_slice(prefix) });
92 }
93}
94
95pub trait KeyValueDB: Sync + Send + tetsy_util_mem::MallocSizeOf {
105 fn transaction(&self) -> DBTransaction {
107 DBTransaction::new()
108 }
109
110 fn get(&self, col: u32, key: &[u8]) -> io::Result<Option<DBValue>>;
112
113 fn get_by_prefix(&self, col: u32, prefix: &[u8]) -> Option<Box<[u8]>>;
115
116 fn write(&self, transaction: DBTransaction) -> io::Result<()>;
118
119 fn iter<'a>(&'a self, col: u32) -> Box<dyn Iterator<Item = (Box<[u8]>, Box<[u8]>)> + 'a>;
121
122 fn iter_with_prefix<'a>(
125 &'a self,
126 col: u32,
127 prefix: &'a [u8],
128 ) -> Box<dyn Iterator<Item = (Box<[u8]>, Box<[u8]>)> + 'a>;
129
130 fn restore(&self, new_db: &str) -> io::Result<()>;
132
133 fn io_stats(&self, _kind: IoStatsKind) -> IoStats {
140 IoStats::empty()
141 }
142
143 fn has_key(&self, col: u32, key: &[u8]) -> io::Result<bool> {
145 self.get(col, key).map(|opt| opt.is_some())
146 }
147
148 fn has_prefix(&self, col: u32, prefix: &[u8]) -> bool {
150 self.get_by_prefix(col, prefix).is_some()
151 }
152}
153
154pub fn end_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
159 let mut end_range = prefix.to_vec();
160 while let Some(0xff) = end_range.last() {
161 end_range.pop();
162 }
163 if let Some(byte) = end_range.last_mut() {
164 *byte += 1;
165 Some(end_range)
166 } else {
167 None
168 }
169}
170
171#[cfg(test)]
172mod test {
173 use super::end_prefix;
174
175 #[test]
176 fn end_prefix_test() {
177 assert_eq!(end_prefix(&[5, 6, 7]), Some(vec![5, 6, 8]));
178 assert_eq!(end_prefix(&[5, 6, 255]), Some(vec![5, 7]));
179 assert_ne!(end_prefix(&[5, 255, 255]), Some(vec![5, 255]));
181 assert_eq!(end_prefix(&[5, 255, 255]), Some(vec![6]));
184 assert_eq!(end_prefix(&[255, 255, 255]), None);
185
186 assert_eq!(end_prefix(&[0x00, 0xff]), Some(vec![0x01]));
187 assert_eq!(end_prefix(&[0xff]), None);
188 assert_eq!(end_prefix(&[]), None);
189 assert_eq!(end_prefix(b"0"), Some(b"1".to_vec()));
190 }
191}