1use std::collections::BTreeMap;
13use std::ops::Bound;
14use std::sync::Arc;
15
16use crate::backend::StorageReadView;
17use crate::read_view::{ReadLimits, ScanProjection};
18use crate::{Op, StorageError, Store};
19
20pub trait DeriveRead: Sync {
22 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError>;
27
28 fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError>;
33
34 fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError>;
40}
41
42fn map_range<'a, V: 'a>(
44 map: &'a BTreeMap<Vec<u8>, V>,
45 start: &[u8],
46 end: &[u8],
47) -> impl Iterator<Item = (&'a Vec<u8>, &'a V)> {
48 let upper = if end.is_empty() {
49 Bound::Unbounded
50 } else {
51 Bound::Excluded(end.to_vec())
52 };
53 map.range::<Vec<u8>, _>((Bound::Included(start.to_vec()), upper))
54}
55
56pub trait Deriver: Send + Sync {
62 fn derive(
68 &self,
69 ci: u64,
70 base_ops: &[Op],
71 before: &dyn DeriveRead,
72 ) -> Result<Vec<Op>, StorageError>;
73
74 fn is_derived_key(&self, key: &[u8]) -> bool;
77
78 fn prefix_reaches_derived(&self, prefix: &[u8]) -> bool;
80}
81
82impl DeriveRead for crate::ReadView {
83 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
84 crate::ReadView::get(self, key)
85 }
86
87 fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
88 crate::ReadView::range_prefix(self, prefix)
89 }
90
91 fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
92 crate::ReadView::range(self, start, end)
93 }
94}
95
96pub(crate) struct Layered<'a, 'b> {
101 pub(crate) base: &'a dyn DeriveRead,
102 pub(crate) buffer: &'a BTreeMap<&'b [u8], Option<&'b [u8]>>,
103}
104
105impl Layered<'_, '_> {
106 fn overlay(
107 &self,
108 mut merged: BTreeMap<Vec<u8>, Vec<u8>>,
109 start: &[u8],
110 end: Bound<&[u8]>,
111 ) -> Vec<(Vec<u8>, Vec<u8>)> {
112 for (key, value) in self.buffer.range::<[u8], _>((Bound::Included(start), end)) {
113 match value {
114 Some(value) => {
115 merged.insert(key.to_vec(), value.to_vec());
116 }
117 None => {
118 merged.remove(*key);
119 }
120 }
121 }
122 merged.into_iter().collect()
123 }
124}
125
126impl DeriveRead for Layered<'_, '_> {
127 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
128 match self.buffer.get(key) {
129 Some(value) => Ok(value.map(<[u8]>::to_vec)),
130 None => self.base.get(key),
131 }
132 }
133
134 fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
135 let merged = self.base.range_prefix(prefix)?.into_iter().collect();
136 let upper = crate::prefix_range(prefix).1;
137 let end = if upper.is_empty() {
138 Bound::Unbounded
139 } else {
140 Bound::Excluded(upper.as_slice())
141 };
142 Ok(self.overlay(merged, prefix, end))
143 }
144
145 fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
146 let merged = self.base.range(start, end)?.into_iter().collect();
147 let end = if end.is_empty() {
148 Bound::Unbounded
149 } else {
150 Bound::Excluded(end)
151 };
152 Ok(self.overlay(merged, start, end))
153 }
154}
155
156pub(crate) fn fold<'b>(
159 buffer: &mut BTreeMap<&'b [u8], Option<&'b [u8]>>,
160 ops: impl IntoIterator<Item = &'b Op>,
161) {
162 for op in ops {
163 match op {
164 Op::Put(key, value) | Op::Set(key, value) => {
165 buffer.insert(key, Some(value));
166 }
167 Op::Del(key) | Op::Purge(key) => {
168 buffer.insert(key, None);
169 }
170 }
171 }
172}
173
174pub(crate) struct BackendRead(pub(crate) Box<dyn StorageReadView>);
176
177impl DeriveRead for BackendRead {
178 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
179 self.0.read_value(
180 key,
181 u64::MAX,
182 true,
183 ReadLimits::default(),
184 ScanProjection::KeyValue,
185 )
186 }
187
188 fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
189 let end = crate::prefix_range(prefix).1;
190 self.0
191 .current_range(
192 prefix,
193 &end,
194 ReadLimits::default(),
195 ScanProjection::KeyValue,
196 )?
197 .collect()
198 }
199
200 fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
201 self.0
202 .current_range(start, end, ReadLimits::default(), ScanProjection::KeyValue)?
203 .collect()
204 }
205}
206
207pub(crate) struct MapRead<'a>(pub(crate) &'a BTreeMap<Vec<u8>, Vec<u8>>);
209
210impl DeriveRead for MapRead<'_> {
211 fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
212 Ok(self.0.get(key).cloned())
213 }
214
215 fn range_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
216 Ok(self
217 .0
218 .range::<[u8], _>((Bound::Included(prefix), Bound::Unbounded))
219 .take_while(|(key, _)| key.starts_with(prefix))
220 .map(|(key, value)| (key.clone(), value.clone()))
221 .collect())
222 }
223
224 fn range(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
225 Ok(map_range(self.0, start, end)
226 .map(|(key, value)| (key.clone(), value.clone()))
227 .collect())
228 }
229}
230
231pub(crate) struct Empty;
233
234impl DeriveRead for Empty {
235 fn get(&self, _key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
236 Ok(None)
237 }
238
239 fn range_prefix(&self, _prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
240 Ok(Vec::new())
241 }
242
243 fn range(&self, _start: &[u8], _end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, StorageError> {
244 Ok(Vec::new())
245 }
246}
247
248impl Store {
249 pub fn set_deriver(&self, deriver: Arc<dyn Deriver>) {
254 if self
259 .underived_replay
260 .load(std::sync::atomic::Ordering::SeqCst)
261 {
262 self.health.fence();
263 }
264 *self.health.cleanup(self.deriver.write()) = Some(deriver);
265 }
266
267 #[doc(hidden)]
273 pub fn inject_without_derivation<T>(&self, f: impl FnOnce() -> T) -> T {
274 let saved = self.health.cleanup(self.deriver.write()).take();
275 let out = f();
276 *self.health.cleanup(self.deriver.write()) = saved;
277 out
278 }
279
280 pub(crate) fn deriver(&self) -> Option<Arc<dyn Deriver>> {
281 self.health.cleanup(self.deriver.read()).clone()
282 }
283
284 pub(crate) fn derive_ops(
289 deriver: Option<&Arc<dyn Deriver>>,
290 ci: u64,
291 base_ops: &[Op],
292 before: &dyn DeriveRead,
293 ) -> Result<Vec<Op>, StorageError> {
294 match deriver {
295 Some(deriver) => deriver.derive(ci, base_ops, before),
296 None => Ok(Vec::new()),
297 }
298 }
299}