umadb_core/
tracking_tree_nodes.rs

1use crate::common::{PageID, Position};
2use byteorder::{ByteOrder, LittleEndian};
3use umadb_dcb::{DCBError, DCBResult};
4
5/// A simple leaf-only node for the tracking tree.
6/// Keys are UTF-8 source strings; values are positions.
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct TrackingLeafNode {
9    pub keys: Vec<String>,
10    pub values: Vec<Position>,
11}
12
13impl TrackingLeafNode {
14    pub fn new() -> Self {
15        Self {
16            keys: Vec::new(),
17            values: Vec::new(),
18        }
19    }
20
21    pub fn calc_serialized_size(&self) -> usize {
22        // layout (v1 and newer on write):
23        // u8: node version (1 for sorted keys)
24        // u16: key_count
25        // repeated key_count times: u8 key_len + [u8; key_len]
26        // repeated key_count times: u64 position
27        let mut size = 1 + 2; // version + count(u16)
28        for k in &self.keys {
29            size += 1 + k.len();
30        }
31        size += 8 * self.values.len();
32        size
33    }
34
35    pub fn serialize_into(&self, buf: &mut [u8]) -> usize {
36        let needed = self.calc_serialized_size();
37        assert!(buf.len() >= needed, "buffer too small for TrackingLeafNode");
38        buf[0] = 1; // version >=1 means keys are sorted
39        LittleEndian::write_u16(&mut buf[1..3], self.keys.len() as u16);
40        let mut off = 3;
41        // Keys are expected to be maintained in sorted order by the node
42        for k in &self.keys {
43            let kb = k.as_bytes();
44            assert!(
45                kb.len() <= u8::MAX as usize,
46                "tracking key too long to serialize (len>{})",
47                u8::MAX
48            );
49            buf[off] = kb.len() as u8;
50            off += 1;
51            buf[off..off + kb.len()].copy_from_slice(kb);
52            off += kb.len();
53        }
54        for v in &self.values {
55            LittleEndian::write_u64(&mut buf[off..off + 8], v.0);
56            off += 8;
57        }
58        needed
59    }
60
61    pub fn from_slice(slice: &[u8]) -> DCBResult<Self> {
62        if slice.is_empty() {
63            return Err(DCBError::DeserializationError(
64                "tracking leaf too small".to_string(),
65            ));
66        }
67        let ver = slice[0];
68        let (mut off, count): (usize, usize) = if ver == 0 {
69            if slice.len() < 5 {
70                return Err(DCBError::DeserializationError(
71                    "tracking leaf too small (v0)".to_string(),
72                ));
73            }
74            let cnt_u32 = LittleEndian::read_u32(&slice[1..5]);
75            if cnt_u32 > u16::MAX as u32 {
76                return Err(DCBError::DeserializationError(
77                    "v0 tracking leaf count exceeds u16".to_string(),
78                ));
79            }
80            (5, cnt_u32 as usize)
81        } else {
82            if slice.len() < 3 {
83                return Err(DCBError::DeserializationError(
84                    "tracking leaf too small (v1)".to_string(),
85                ));
86            }
87            (3, LittleEndian::read_u16(&slice[1..3]) as usize)
88        };
89        let mut keys = Vec::with_capacity(count);
90        for _ in 0..count {
91            if ver == 0 {
92                if off + 4 > slice.len() {
93                    return Err(DCBError::DeserializationError(
94                        "tracking leaf truncated klen (v0)".to_string(),
95                    ));
96                }
97                let klen_u32 = LittleEndian::read_u32(&slice[off..off + 4]);
98                if klen_u32 > u8::MAX as u32 {
99                    return Err(DCBError::DeserializationError(
100                        "v0 tracking leaf key length exceeds u8".to_string(),
101                    ));
102                }
103                let klen = klen_u32 as usize;
104                off += 4;
105                if off + klen > slice.len() {
106                    return Err(DCBError::DeserializationError(
107                        "tracking leaf truncated key (v0)".to_string(),
108                    ));
109                }
110                let k = std::str::from_utf8(&slice[off..off + klen])
111                    .map_err(|e| DCBError::DeserializationError(format!("invalid utf8: {e}")))?
112                    .to_string();
113                off += klen;
114                keys.push(k);
115            } else {
116                if off + 1 > slice.len() {
117                    return Err(DCBError::DeserializationError(
118                        "tracking leaf truncated klen (v1)".to_string(),
119                    ));
120                }
121                let klen = slice[off] as usize;
122                off += 1;
123                if off + klen > slice.len() {
124                    return Err(DCBError::DeserializationError(
125                        "tracking leaf truncated key (v1)".to_string(),
126                    ));
127                }
128                let k = std::str::from_utf8(&slice[off..off + klen])
129                    .map_err(|e| DCBError::DeserializationError(format!("invalid utf8: {e}")))?
130                    .to_string();
131                off += klen;
132                keys.push(k);
133            }
134        }
135        let mut values = Vec::with_capacity(count);
136        for _ in 0..count {
137            if off + 8 > slice.len() {
138                return Err(DCBError::DeserializationError(
139                    "tracking leaf truncated value".to_string(),
140                ));
141            }
142            let v = LittleEndian::read_u64(&slice[off..off + 8]);
143            off += 8;
144            values.push(Position(v));
145        }
146        if ver == 0 {
147            // Old format: keys may be unsorted. Sort keys and keep values aligned.
148            let mut pairs: Vec<(String, Position)> = keys.into_iter().zip(values).collect();
149            pairs.sort_by(|a, b| a.0.cmp(&b.0));
150            let (new_keys, new_vals): (Vec<String>, Vec<Position>) = pairs.into_iter().unzip();
151            Ok(Self {
152                keys: new_keys,
153                values: new_vals,
154            })
155        } else {
156            // v>=1: keys are already sorted
157            Ok(Self { keys, values })
158        }
159    }
160
161    pub fn get(&self, source: &str) -> Option<Position> {
162        match self.keys.binary_search_by(|k| k.as_str().cmp(source)) {
163            Ok(i) => Some(self.values[i]),
164            Err(_) => None,
165        }
166    }
167
168    /// Insert or update a key with value while enforcing page size limit (no splitting).
169    /// Returns Err(InternalError("not implemented: tracking split")) if it would exceed page capacity.
170    pub fn upsert_no_split(
171        &mut self,
172        source: &str,
173        pos: Position,
174        page_body_capacity: usize,
175    ) -> DCBResult<()> {
176        match self.keys.binary_search_by(|k| k.as_str().cmp(source)) {
177            Ok(idx) => {
178                // Key exists; update value in place
179                self.values[idx] = pos;
180                Ok(())
181            }
182            Err(ins_idx) => {
183                // Key not found; check capacity before inserting
184                let key_len = source.len();
185                // Check free space as key bytes length + 8 for position (as requested)
186                let additional = key_len + 8;
187                let current = self.calc_serialized_size();
188                if current + additional > page_body_capacity {
189                    return Err(DCBError::InternalError(
190                        "not implemented: tracking split".to_string(),
191                    ));
192                }
193                self.keys.insert(ins_idx, source.to_string());
194                self.values.insert(ins_idx, pos);
195                Ok(())
196            }
197        }
198    }
199}
200
201#[derive(Debug, Clone, PartialEq, Eq)]
202pub struct TrackingInternalNode {
203    pub keys: Vec<String>,
204    pub child_ids: Vec<PageID>,
205}
206
207impl TrackingInternalNode {
208    pub fn child_index_for_key(&self, key: &str) -> usize {
209        match self.keys.binary_search_by(|k| k.as_str().cmp(key)) {
210            Ok(idx) => idx + 1,
211            Err(idx) => idx,
212        }
213    }
214
215    pub fn replace_child_id_at(
216        &mut self,
217        idx: usize,
218        old_id: PageID,
219        new_id: PageID,
220    ) -> DCBResult<()> {
221        if idx >= self.child_ids.len() {
222            return Err(DCBError::DatabaseCorrupted(
223                "child index out of bounds".to_string(),
224            ));
225        }
226        if self.child_ids[idx] != old_id {
227            return Err(DCBError::DatabaseCorrupted("Child ID mismatch".to_string()));
228        }
229        self.child_ids[idx] = new_id;
230        Ok(())
231    }
232
233    pub fn insert_promoted_at(&mut self, idx: usize, key: String, right_child: PageID) {
234        self.keys.insert(idx, key);
235        self.child_ids.insert(idx + 1, right_child);
236    }
237
238    /// Split this internal node around the middle key. Returns (promoted_key, right_keys, right_child_ids).
239    pub fn split_off(&mut self) -> DCBResult<(String, Vec<String>, Vec<PageID>)> {
240        if self.child_ids.len() < 4 || self.keys.len() + 1 != self.child_ids.len() {
241            return Err(DCBError::DatabaseCorrupted(
242                "Cannot split tracking internal with insufficient arity".to_string(),
243            ));
244        }
245        let mid = self.keys.len() / 2; // promote this key
246        let promoted_key = self.keys[mid].clone();
247        let right_keys: Vec<String> = self.keys[mid + 1..].to_vec();
248        let right_child_ids: Vec<PageID> = self.child_ids[mid + 1..].to_vec();
249        // Truncate left
250        self.keys.truncate(mid);
251        self.child_ids.truncate(mid + 1);
252        Ok((promoted_key, right_keys, right_child_ids))
253    }
254    pub fn new() -> Self {
255        Self {
256            keys: Vec::new(),
257            child_ids: Vec::new(),
258        }
259    }
260
261    pub fn calc_serialized_size(&self) -> usize {
262        // layout (v1 write):
263        // u8: version (1)
264        // u16: key_count
265        // repeated key_count times: u8 key_len + [u8; key_len]
266        // repeated (key_count + 1) times: u64 child_page_id
267        let mut size = 1 + 2; // version + count(u16)
268        for k in &self.keys {
269            size += 1 + k.len();
270        }
271        size += 8 * (self.keys.len() + 1);
272        size
273    }
274
275    pub fn serialize_into(&self, buf: &mut [u8]) -> usize {
276        let needed = self.calc_serialized_size();
277        assert!(
278            buf.len() >= needed,
279            "buffer too small for TrackingInternalNode"
280        );
281        buf[0] = 1; // version 1
282        LittleEndian::write_u16(&mut buf[1..3], self.keys.len() as u16);
283        let mut off = 3;
284        for k in &self.keys {
285            let kb = k.as_bytes();
286            assert!(
287                kb.len() <= u8::MAX as usize,
288                "tracking internal key too long to serialize (len>{})",
289                u8::MAX
290            );
291            buf[off] = kb.len() as u8;
292            off += 1;
293            buf[off..off + kb.len()].copy_from_slice(kb);
294            off += kb.len();
295        }
296        // children count is implied as keys.len() + 1
297        for id in &self.child_ids {
298            LittleEndian::write_u64(&mut buf[off..off + 8], id.0);
299            off += 8;
300        }
301        needed
302    }
303
304    pub fn from_slice(slice: &[u8]) -> DCBResult<Self> {
305        if slice.is_empty() {
306            return Err(DCBError::DeserializationError(
307                "tracking internal too small".to_string(),
308            ));
309        }
310        let ver = slice[0];
311        if ver == 0 {
312            return Err(DCBError::DeserializationError(
313                "unsupported tracking internal version 0".to_string(),
314            ));
315        }
316        if slice.len() < 3 {
317            return Err(DCBError::DeserializationError(
318                "tracking internal too small (v1)".to_string(),
319            ));
320        }
321        let count = LittleEndian::read_u16(&slice[1..3]) as usize;
322        let mut off = 3;
323        let mut keys = Vec::with_capacity(count);
324        for _ in 0..count {
325            if off + 1 > slice.len() {
326                return Err(DCBError::DeserializationError(
327                    "tracking internal truncated klen (v1)".to_string(),
328                ));
329            }
330            let klen = slice[off] as usize;
331            off += 1;
332            if off + klen > slice.len() {
333                return Err(DCBError::DeserializationError(
334                    "tracking internal truncated key (v1)".to_string(),
335                ));
336            }
337            let k = std::str::from_utf8(&slice[off..off + klen])
338                .map_err(|e| DCBError::DeserializationError(format!("invalid utf8: {e}")))?
339                .to_string();
340            off += klen;
341            keys.push(k);
342        }
343        // Now children: implied count = keys.len() + 1
344        let child_count = keys.len() + 1;
345        let need = off + 8 * child_count;
346        if slice.len() < need {
347            return Err(DCBError::DeserializationError(
348                "tracking internal truncated children".to_string(),
349            ));
350        }
351        let mut child_ids = Vec::with_capacity(child_count);
352        for _ in 0..child_count {
353            let v = LittleEndian::read_u64(&slice[off..off + 8]);
354            off += 8;
355            child_ids.push(PageID(v));
356        }
357        Ok(Self { keys, child_ids })
358    }
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn test_tracking_leaf_roundtrip() {
367        let mut node = TrackingLeafNode::new();
368        node.keys = vec!["a".to_string(), "b".to_string()];
369        node.values = vec![Position(1), Position(2)];
370        let mut buf = vec![0u8; node.calc_serialized_size()];
371        let n = node.serialize_into(&mut buf);
372        assert_eq!(n, buf.len());
373        let dec = TrackingLeafNode::from_slice(&buf).unwrap();
374        assert_eq!(node, dec);
375        assert_eq!(dec.get("a"), Some(Position(1)));
376        assert_eq!(dec.get("z"), None);
377    }
378
379    #[test]
380    fn test_deserialize_v0_unsorted_sorts_and_aligns() {
381        // Manually craft a v0 buffer with unsorted keys: ["b", "c", "a"] and values [2, 3, 1]
382        let keys = vec!["b", "c", "a"];
383        let values = vec![Position(2), Position(3), Position(1)];
384        // compute size for v0: 1 (ver) + 4 (count) + sum(4 + klen) + 8*count
385        let key_bytes: Vec<Vec<u8>> = keys.iter().map(|k| k.as_bytes().to_vec()).collect();
386        let mut size = 1 + 4;
387        for kb in &key_bytes {
388            size += 4 + kb.len();
389        }
390        size += 8 * values.len();
391        let mut buf = vec![0u8; size];
392        buf[0] = 0; // version 0
393        LittleEndian::write_u32(&mut buf[1..5], keys.len() as u32);
394        let mut off = 5;
395        for kb in &key_bytes {
396            LittleEndian::write_u32(&mut buf[off..off + 4], kb.len() as u32);
397            off += 4;
398            buf[off..off + kb.len()].copy_from_slice(kb);
399            off += kb.len();
400        }
401        for v in &values {
402            LittleEndian::write_u64(&mut buf[off..off + 8], v.0);
403            off += 8;
404        }
405        // Now deserialize; it should sort keys and align values accordingly
406        let dec = TrackingLeafNode::from_slice(&buf).unwrap();
407        assert_eq!(dec.keys, vec!["a", "b", "c"]);
408        assert_eq!(dec.values, vec![Position(1), Position(2), Position(3)]);
409        // Binary search based get should work
410        assert_eq!(dec.get("a"), Some(Position(1)));
411        assert_eq!(dec.get("c"), Some(Position(3)));
412        assert_eq!(dec.get("z"), None);
413    }
414
415    #[test]
416    fn test_upsert_maintains_sorted_and_capacity_check() {
417        let mut node = TrackingLeafNode::new();
418        // small capacity; account for version+count and future entries
419        let capacity = 1 + 4 + (4 + 1) + (4 + 1) + (4 + 1) + 8 * 3; // rough upper bound; we pre-check len+8 anyway
420        node.upsert_no_split("b", Position(2), capacity).unwrap();
421        node.upsert_no_split("a", Position(1), capacity).unwrap();
422        node.upsert_no_split("c", Position(3), capacity).unwrap();
423        assert_eq!(
424            node.keys,
425            vec!["a".to_string(), "b".to_string(), "c".to_string()]
426        );
427        assert_eq!(node.get("b"), Some(Position(2)));
428        // Now try to insert with too small capacity
429        let mut node2 = TrackingLeafNode::new();
430        let small_capacity = 1 + 4; // too small to accept any key/value
431        let err = node2
432            .upsert_no_split("x", Position(9), small_capacity)
433            .unwrap_err();
434        match err {
435            DCBError::InternalError(s) => assert!(s.contains("tracking split")),
436            _ => panic!("unexpected error type"),
437        }
438    }
439
440    #[test]
441    fn test_tracking_internal_roundtrip() {
442        let node = TrackingInternalNode {
443            keys: vec!["alpha".into(), "beta".into(), "gamma".into()],
444            child_ids: vec![PageID(10), PageID(20), PageID(30), PageID(40)],
445        };
446        let mut buf = vec![0u8; node.calc_serialized_size()];
447        let n = node.serialize_into(&mut buf);
448        assert_eq!(n, buf.len());
449        let dec = TrackingInternalNode::from_slice(&buf).unwrap();
450        assert_eq!(dec, node);
451    }
452
453    #[test]
454    fn test_tracking_internal_empty_keys_one_child_roundtrip() {
455        let node = TrackingInternalNode {
456            keys: vec![],
457            child_ids: vec![PageID(123)],
458        };
459        let mut buf = vec![0u8; node.calc_serialized_size()];
460        let n = node.serialize_into(&mut buf);
461        assert_eq!(n, buf.len());
462        let dec = TrackingInternalNode::from_slice(&buf).unwrap();
463        assert_eq!(dec, node);
464    }
465
466    #[test]
467    fn test_tracking_internal_from_slice_truncated_children_err() {
468        // Build a valid buffer then truncate last 8 bytes to simulate missing child
469        let node = TrackingInternalNode {
470            keys: vec!["k1".into()],
471            child_ids: vec![PageID(1), PageID(2)],
472        };
473        let mut buf = vec![0u8; node.calc_serialized_size()];
474        let _ = node.serialize_into(&mut buf);
475        buf.truncate(buf.len() - 4); // remove less than 8 to force error path
476        let err = TrackingInternalNode::from_slice(&buf).unwrap_err();
477        match err {
478            DCBError::DeserializationError(_) => {}
479            _ => panic!("expected DeserializationError"),
480        }
481    }
482}
umadb_core/tracking_tree_nodes.rs

umadb_core/
tracking_tree_nodes.rs
