1use std::collections::{BTreeMap, HashMap};
10
11use nodedb_array::schema::ArraySchema;
12use nodedb_array::tile::cell_payload::{
13 CELL_GDPR_ERASURE_SENTINEL, CELL_TOMBSTONE_SENTINEL, CellPayload, is_cell_gdpr_erasure,
14 is_cell_sentinel, is_cell_tombstone,
15};
16use nodedb_array::tile::sparse_tile::{RowKind, SparseRow, SparseTile, SparseTileBuilder};
17use nodedb_array::tile::tile_id_for_cell;
18use nodedb_array::types::TileId;
19use nodedb_array::types::cell_value::value::CellValue;
20use nodedb_array::types::coord::value::CoordValue;
21use nodedb_types::{OPEN_UPPER, Surrogate};
22
23use crate::error::LiteError;
24
25pub(crate) fn encode_coord(coord: &[CoordValue]) -> Result<Vec<u8>, LiteError> {
27 zerompk::to_msgpack_vec(&coord.to_vec()).map_err(|e| LiteError::Serialization {
28 detail: format!("encode coord: {e}"),
29 })
30}
31
32fn decode_coord(key: &[u8]) -> Result<Vec<CoordValue>, LiteError> {
33 zerompk::from_msgpack(key).map_err(|e| LiteError::Serialization {
34 detail: format!("decode coord: {e}"),
35 })
36}
37
38#[derive(Debug, Default, Clone)]
40struct TileBuffer {
41 entries: HashMap<Vec<u8>, Vec<u8>>,
42}
43
44impl TileBuffer {
45 fn entry_count(&self) -> usize {
46 self.entries.len()
47 }
48
49 fn live_cell_count(&self) -> usize {
50 self.entries
51 .values()
52 .filter(|v| !is_cell_sentinel(v))
53 .count()
54 }
55
56 fn get(&self, coord: &[CoordValue]) -> Option<&[u8]> {
57 let key = encode_coord(coord).ok()?;
58 self.entries.get(&key).map(|v| v.as_slice())
59 }
60
61 fn materialise(&self, schema: &ArraySchema) -> Result<SparseTile, LiteError> {
62 let mut b = SparseTileBuilder::new(schema);
63 for (coord_key, bytes) in &self.entries {
64 let coord = decode_coord(coord_key)?;
65 if is_cell_tombstone(bytes) {
66 b.push_row(SparseRow {
67 coord: &coord,
68 attrs: &[],
69 surrogate: None,
70 valid_from_ms: 0,
71 valid_until_ms: OPEN_UPPER,
72 kind: RowKind::Tombstone,
73 })
74 .map_err(|e| LiteError::Storage {
75 detail: format!("push tombstone row: {e}"),
76 })?;
77 } else if is_cell_gdpr_erasure(bytes) {
78 b.push_row(SparseRow {
79 coord: &coord,
80 attrs: &[],
81 surrogate: None,
82 valid_from_ms: 0,
83 valid_until_ms: OPEN_UPPER,
84 kind: RowKind::GdprErased,
85 })
86 .map_err(|e| LiteError::Storage {
87 detail: format!("push erasure row: {e}"),
88 })?;
89 } else {
90 let payload = CellPayload::decode(bytes).map_err(|e| LiteError::Storage {
91 detail: format!("decode CellPayload: {e}"),
92 })?;
93 b.push_row(SparseRow {
94 coord: &coord,
95 attrs: &payload.attrs,
96 surrogate: Some(payload.surrogate),
97 valid_from_ms: payload.valid_from_ms,
98 valid_until_ms: payload.valid_until_ms,
99 kind: RowKind::Live,
100 })
101 .map_err(|e| LiteError::Storage {
102 detail: format!("push live row: {e}"),
103 })?;
104 }
105 }
106 Ok(b.build())
107 }
108}
109
110#[derive(Debug, Default)]
112pub struct MemtableStats {
113 pub tile_count: usize,
114 pub cell_count: usize,
115}
116
117pub struct PutCellArgs {
119 pub coord: Vec<CoordValue>,
120 pub attrs: Vec<CellValue>,
121 pub surrogate: Surrogate,
122 pub system_from_ms: i64,
123 pub valid_from_ms: i64,
124 pub valid_until_ms: i64,
125}
126
127#[derive(Debug, Default)]
129pub struct ArrayMemtable {
130 tiles: BTreeMap<TileId, TileBuffer>,
131}
132
133impl ArrayMemtable {
134 pub fn new() -> Self {
135 Self::default()
136 }
137
138 pub fn put_cell(
140 &mut self,
141 schema: &ArraySchema,
142 args: PutCellArgs,
143 ) -> Result<TileId, LiteError> {
144 let PutCellArgs {
145 coord,
146 attrs,
147 surrogate,
148 system_from_ms,
149 valid_from_ms,
150 valid_until_ms,
151 } = args;
152 let tile_id =
153 tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
154 detail: format!("tile_id_for_cell: {e}"),
155 })?;
156 let payload = CellPayload {
157 valid_from_ms,
158 valid_until_ms,
159 attrs,
160 surrogate,
161 };
162 let bytes = payload.encode().map_err(|e| LiteError::Storage {
163 detail: format!("encode CellPayload: {e}"),
164 })?;
165 let key = encode_coord(&coord)?;
166 self.tiles
167 .entry(tile_id)
168 .or_default()
169 .entries
170 .insert(key, bytes);
171 Ok(tile_id)
172 }
173
174 pub fn delete_cell(
176 &mut self,
177 schema: &ArraySchema,
178 coord: Vec<CoordValue>,
179 system_from_ms: i64,
180 ) -> Result<TileId, LiteError> {
181 let tile_id =
182 tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
183 detail: format!("tile_id_for_cell: {e}"),
184 })?;
185 let key = encode_coord(&coord)?;
186 self.tiles
187 .entry(tile_id)
188 .or_default()
189 .entries
190 .insert(key, CELL_TOMBSTONE_SENTINEL.to_vec());
191 Ok(tile_id)
192 }
193
194 pub fn erase_cell(
196 &mut self,
197 schema: &ArraySchema,
198 coord: Vec<CoordValue>,
199 system_from_ms: i64,
200 ) -> Result<TileId, LiteError> {
201 let tile_id =
202 tile_id_for_cell(schema, &coord, system_from_ms).map_err(|e| LiteError::Storage {
203 detail: format!("tile_id_for_cell: {e}"),
204 })?;
205 let key = encode_coord(&coord)?;
206 self.tiles
207 .entry(tile_id)
208 .or_default()
209 .entries
210 .insert(key, CELL_GDPR_ERASURE_SENTINEL.to_vec());
211 Ok(tile_id)
212 }
213
214 pub fn stats(&self) -> MemtableStats {
216 MemtableStats {
217 tile_count: self.tiles.len(),
218 cell_count: self.tiles.values().map(|b| b.live_cell_count()).sum(),
219 }
220 }
221
222 pub fn is_empty(&self) -> bool {
223 self.tiles.values().all(|b| b.entry_count() == 0)
224 }
225
226 pub fn tiles_iter(&self) -> impl Iterator<Item = TileId> + '_ {
228 self.tiles.keys().copied()
229 }
230
231 pub fn drain_and_materialise(
233 &mut self,
234 schema: &ArraySchema,
235 ) -> Result<Vec<(TileId, SparseTile)>, LiteError> {
236 let tiles: Vec<_> = std::mem::take(&mut self.tiles).into_iter().collect();
237 tiles
238 .into_iter()
239 .map(|(id, buf)| buf.materialise(schema).map(|tile| (id, tile)))
240 .collect()
241 }
242
243 pub fn iter_cell_versions(
247 &self,
248 hilbert_prefix: u64,
249 system_as_of: i64,
250 coord: &[CoordValue],
251 ) -> Vec<(TileId, Vec<u8>)> {
252 let lo = TileId::new(hilbert_prefix, i64::MIN);
253 let hi = TileId::new(hilbert_prefix, system_as_of);
254 self.tiles
255 .range(lo..=hi)
256 .rev()
257 .filter_map(|(id, buf)| buf.get(coord).map(|b| (*id, b.to_vec())))
258 .collect()
259 }
260
261 pub fn iter_tiles_for_prefix(
264 &self,
265 hilbert_prefix: u64,
266 system_as_of: i64,
267 schema: &ArraySchema,
268 ) -> Result<Vec<(TileId, SparseTile)>, LiteError> {
269 let lo = TileId::new(hilbert_prefix, i64::MIN);
270 let hi = TileId::new(hilbert_prefix, system_as_of);
271 self.tiles
272 .range(lo..=hi)
273 .map(|(id, buf)| buf.materialise(schema).map(|tile| (*id, tile)))
274 .collect()
275 }
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281 use nodedb_array::schema::ArraySchemaBuilder;
282 use nodedb_array::schema::attr_spec::{AttrSpec, AttrType};
283 use nodedb_array::schema::dim_spec::{DimSpec, DimType};
284 use nodedb_array::types::domain::{Domain, DomainBound};
285
286 fn schema() -> ArraySchema {
287 ArraySchemaBuilder::new("a")
288 .dim(DimSpec::new(
289 "x",
290 DimType::Int64,
291 Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
292 ))
293 .attr(AttrSpec::new("v", AttrType::Int64, true))
294 .tile_extents(vec![4])
295 .build()
296 .unwrap()
297 }
298
299 #[test]
300 fn put_and_stats() {
301 let s = schema();
302 let mut m = ArrayMemtable::new();
303 m.put_cell(
304 &s,
305 PutCellArgs {
306 coord: vec![CoordValue::Int64(1)],
307 attrs: vec![CellValue::Int64(10)],
308 surrogate: Surrogate::new(1),
309 system_from_ms: 100,
310 valid_from_ms: 0,
311 valid_until_ms: OPEN_UPPER,
312 },
313 )
314 .unwrap();
315 assert_eq!(m.stats().cell_count, 1);
316 }
317
318 #[test]
319 fn delete_adds_tombstone_version() {
320 let s = schema();
321 let mut m = ArrayMemtable::new();
322 m.put_cell(
323 &s,
324 PutCellArgs {
325 coord: vec![CoordValue::Int64(1)],
326 attrs: vec![CellValue::Int64(1)],
327 surrogate: Surrogate::new(1),
328 system_from_ms: 100,
329 valid_from_ms: 0,
330 valid_until_ms: OPEN_UPPER,
331 },
332 )
333 .unwrap();
334 m.delete_cell(&s, vec![CoordValue::Int64(1)], 200).unwrap();
335 assert_eq!(m.tiles.len(), 2);
336 assert_eq!(m.stats().cell_count, 1);
337 }
338
339 #[test]
340 fn erase_adds_gdpr_sentinel() {
341 let s = schema();
342 let mut m = ArrayMemtable::new();
343 m.put_cell(
344 &s,
345 PutCellArgs {
346 coord: vec![CoordValue::Int64(2)],
347 attrs: vec![CellValue::Int64(99)],
348 surrogate: Surrogate::new(1),
349 system_from_ms: 10,
350 valid_from_ms: 0,
351 valid_until_ms: OPEN_UPPER,
352 },
353 )
354 .unwrap();
355 m.erase_cell(&s, vec![CoordValue::Int64(2)], 20).unwrap();
356 assert_eq!(m.tiles.len(), 2);
357 }
358
359 #[test]
360 fn drain_and_materialise_empties() {
361 let s = schema();
362 let mut m = ArrayMemtable::new();
363 m.put_cell(
364 &s,
365 PutCellArgs {
366 coord: vec![CoordValue::Int64(3)],
367 attrs: vec![CellValue::Int64(7)],
368 surrogate: Surrogate::new(1),
369 system_from_ms: 1,
370 valid_from_ms: 0,
371 valid_until_ms: OPEN_UPPER,
372 },
373 )
374 .unwrap();
375 let tiles = m.drain_and_materialise(&s).unwrap();
376 assert_eq!(tiles.len(), 1);
377 assert!(m.is_empty());
378 }
379
380 #[test]
381 fn iter_cell_versions_newest_first() {
382 let s = schema();
383 let mut m = ArrayMemtable::new();
384 let coord = vec![CoordValue::Int64(0)];
385 for sys in [100i64, 200, 300] {
386 m.put_cell(
387 &s,
388 PutCellArgs {
389 coord: coord.clone(),
390 attrs: vec![CellValue::Int64(sys)],
391 surrogate: Surrogate::new(1),
392 system_from_ms: sys,
393 valid_from_ms: 0,
394 valid_until_ms: OPEN_UPPER,
395 },
396 )
397 .unwrap();
398 }
399 let tile_id = nodedb_array::tile::tile_id_for_cell(&s, &coord, 100).unwrap();
400 let versions = m.iter_cell_versions(tile_id.hilbert_prefix, 300, &coord);
401 assert_eq!(versions.len(), 3);
402 assert!(versions[0].0.system_from_ms >= versions[1].0.system_from_ms);
403 }
404}