1use std::collections::{BTreeMap, BTreeSet};
39use std::fs::File;
40
41use mentedb_core::MemoryNode;
42use mentedb_core::error::{MenteError, MenteResult};
43use mentedb_core::types::MemoryId;
44
45pub const BLOCK_SIZE: usize = 512;
48
49pub const EXTENT_HEADER_SIZE: usize = 48;
51
52pub const DATA_FILE: &str = "records.db";
54
55pub const WAL_FILE: &str = "records.wal";
57
58const FILE_MAGIC: &[u8; 8] = b"MDBREC01";
59const FILE_VERSION: u32 = 1;
60const EXTENT_MAGIC: &[u8; 4] = b"MREC";
61
62const PAYLOAD_V1: u8 = 1;
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
73pub struct PageId(pub u64);
74
75#[derive(Debug, Clone, Copy, PartialEq, Eq)]
77pub(crate) struct Extent {
78 pub(crate) start: u64,
79 pub(crate) blocks: u64,
80}
81
82impl Extent {
83 pub(crate) fn end(&self) -> u64 {
84 self.start + self.blocks
85 }
86
87 pub(crate) fn offset(&self) -> u64 {
88 self.start * BLOCK_SIZE as u64
89 }
90
91 pub(crate) fn byte_len(&self) -> usize {
92 self.blocks as usize * BLOCK_SIZE
93 }
94}
95
96pub(crate) fn blocks_for(payload_len: usize) -> u64 {
98 (EXTENT_HEADER_SIZE + payload_len).div_ceil(BLOCK_SIZE) as u64
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102#[repr(u8)]
103pub(crate) enum ExtentKind {
104 Live = 1,
105 Free = 2,
106}
107
108#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub(crate) struct ExtentHeader {
111 pub(crate) kind: ExtentKind,
112 pub(crate) len: u32,
113 pub(crate) blocks: u32,
114 pub(crate) locator: u64,
115 pub(crate) id: MemoryId,
116 pub(crate) payload_crc: u32,
117}
118
119impl ExtentHeader {
120 pub(crate) fn encode(&self) -> [u8; EXTENT_HEADER_SIZE] {
121 let mut b = [0u8; EXTENT_HEADER_SIZE];
122 b[0..4].copy_from_slice(EXTENT_MAGIC);
123 b[4] = self.kind as u8;
124 b[8..12].copy_from_slice(&self.len.to_le_bytes());
125 b[12..16].copy_from_slice(&self.blocks.to_le_bytes());
126 b[16..24].copy_from_slice(&self.locator.to_le_bytes());
127 b[24..40].copy_from_slice(self.id.0.as_bytes());
128 b[40..44].copy_from_slice(&self.payload_crc.to_le_bytes());
129 let crc = crc32fast::hash(&b[..44]);
130 b[44..48].copy_from_slice(&crc.to_le_bytes());
131 b
132 }
133
134 pub(crate) fn decode(b: &[u8]) -> Option<Self> {
136 if b.len() < EXTENT_HEADER_SIZE || &b[0..4] != EXTENT_MAGIC {
137 return None;
138 }
139 if crc32fast::hash(&b[..44]) != le_u32(&b[44..48]) {
140 return None;
141 }
142 let kind = match b[4] {
143 1 => ExtentKind::Live,
144 2 => ExtentKind::Free,
145 _ => return None,
146 };
147 let blocks = le_u32(&b[12..16]);
148 if blocks == 0 {
149 return None;
150 }
151 let id = uuid::Uuid::from_slice(&b[24..40]).ok()?;
152 Some(Self {
153 kind,
154 len: le_u32(&b[8..12]),
155 blocks,
156 locator: le_u64(&b[16..24]),
157 id: MemoryId(id),
158 payload_crc: le_u32(&b[40..44]),
159 })
160 }
161
162 pub(crate) fn free(blocks: u64) -> Self {
163 Self {
164 kind: ExtentKind::Free,
165 len: 0,
166 blocks: blocks as u32,
167 locator: 0,
168 id: MemoryId(uuid::Uuid::nil()),
169 payload_crc: 0,
170 }
171 }
172}
173
174pub(crate) fn le_u32(b: &[u8]) -> u32 {
175 let mut a = [0u8; 4];
176 a.copy_from_slice(&b[..4]);
177 u32::from_le_bytes(a)
178}
179
180pub(crate) fn le_u64(b: &[u8]) -> u64 {
181 let mut a = [0u8; 8];
182 a.copy_from_slice(&b[..8]);
183 u64::from_le_bytes(a)
184}
185
186pub(crate) fn live_image(locator: u64, id: MemoryId, payload: &[u8], blocks: u64) -> Vec<u8> {
188 let header = ExtentHeader {
189 kind: ExtentKind::Live,
190 len: payload.len() as u32,
191 blocks: blocks as u32,
192 locator,
193 id,
194 payload_crc: crc32fast::hash(payload),
195 };
196 let mut buf = vec![0u8; blocks as usize * BLOCK_SIZE];
197 buf[..EXTENT_HEADER_SIZE].copy_from_slice(&header.encode());
198 buf[EXTENT_HEADER_SIZE..EXTENT_HEADER_SIZE + payload.len()].copy_from_slice(payload);
199 buf
200}
201
202pub(crate) fn live_payload(buf: &[u8], locator: u64) -> MenteResult<&[u8]> {
204 let header = ExtentHeader::decode(buf)
205 .ok_or_else(|| MenteError::Storage(format!("record {locator}: invalid extent header")))?;
206 if header.kind != ExtentKind::Live || header.locator != locator {
207 return Err(MenteError::Storage(format!(
208 "record {locator}: extent holds {:?} record {}",
209 header.kind, header.locator
210 )));
211 }
212 let end = EXTENT_HEADER_SIZE + header.len as usize;
213 if end > buf.len() {
214 return Err(MenteError::Storage(format!(
215 "record {locator}: payload runs past its extent"
216 )));
217 }
218 let payload = &buf[EXTENT_HEADER_SIZE..end];
219 if crc32fast::hash(payload) != header.payload_crc {
220 return Err(MenteError::Storage(format!(
221 "record {locator}: payload checksum mismatch"
222 )));
223 }
224 Ok(payload)
225}
226
227pub(crate) fn encode_node(node: &MemoryNode) -> MenteResult<Vec<u8>> {
235 let mut rest = node.clone();
236 let embedding = std::mem::take(&mut rest.embedding);
237 let json = serde_json::to_vec(&rest).map_err(|e| MenteError::Serialization(e.to_string()))?;
238 let mut buf = Vec::with_capacity(1 + 4 + embedding.len() * 4 + json.len());
239 buf.push(PAYLOAD_V1);
240 buf.extend_from_slice(&(embedding.len() as u32).to_le_bytes());
241 for v in &embedding {
242 buf.extend_from_slice(&v.to_le_bytes());
243 }
244 buf.extend_from_slice(&json);
245 Ok(buf)
246}
247
248pub(crate) fn decode_node(payload: &[u8]) -> MenteResult<MemoryNode> {
250 let bad = |what: &str| MenteError::Serialization(format!("record payload {what}"));
251 if payload.len() < 5 {
252 return Err(bad("too short"));
253 }
254 if payload[0] != PAYLOAD_V1 {
255 return Err(bad("has an unknown encoding"));
256 }
257 let dim = le_u32(&payload[1..5]) as usize;
258 let emb_end = dim
259 .checked_mul(4)
260 .and_then(|n| n.checked_add(5))
261 .filter(|&e| e <= payload.len())
262 .ok_or_else(|| bad("embedding runs past the payload"))?;
263 let mut node: MemoryNode = serde_json::from_slice(&payload[emb_end..])
264 .map_err(|e| MenteError::Serialization(e.to_string()))?;
265 node.embedding = payload[5..emb_end]
266 .as_chunks::<4>()
267 .0
268 .iter()
269 .map(|c| f32::from_le_bytes(*c))
270 .collect();
271 Ok(node)
272}
273
274pub(crate) fn file_header() -> Vec<u8> {
279 let mut b = vec![0u8; BLOCK_SIZE];
280 b[0..8].copy_from_slice(FILE_MAGIC);
281 b[8..12].copy_from_slice(&FILE_VERSION.to_le_bytes());
282 b[12..16].copy_from_slice(&(BLOCK_SIZE as u32).to_le_bytes());
283 let crc = crc32fast::hash(&b[..16]);
284 b[16..20].copy_from_slice(&crc.to_le_bytes());
285 b
286}
287
288pub(crate) fn check_file_header(b: &[u8]) -> MenteResult<()> {
290 if b.len() < 20 || &b[0..8] != FILE_MAGIC {
291 return Err(MenteError::Storage("invalid record file magic".into()));
292 }
293 if crc32fast::hash(&b[..16]) != le_u32(&b[16..20]) {
294 return Err(MenteError::Storage(
295 "record file header checksum mismatch".into(),
296 ));
297 }
298 let version = le_u32(&b[8..12]);
299 if version != FILE_VERSION {
300 return Err(MenteError::Storage(format!(
301 "unsupported record file version {version}"
302 )));
303 }
304 let block = le_u32(&b[12..16]) as usize;
305 if block != BLOCK_SIZE {
306 return Err(MenteError::Storage(format!(
307 "record file block size {block} does not match {BLOCK_SIZE}"
308 )));
309 }
310 Ok(())
311}
312
313pub(crate) fn read_at(file: &File, buf: &mut [u8], offset: u64) -> std::io::Result<()> {
318 #[cfg(unix)]
319 {
320 std::os::unix::fs::FileExt::read_exact_at(file, buf, offset)
321 }
322 #[cfg(windows)]
323 {
324 let mut done = 0usize;
325 while done < buf.len() {
326 let n = std::os::windows::fs::FileExt::seek_read(
327 file,
328 &mut buf[done..],
329 offset + done as u64,
330 )?;
331 if n == 0 {
332 return Err(std::io::ErrorKind::UnexpectedEof.into());
333 }
334 done += n;
335 }
336 Ok(())
337 }
338}
339
340pub(crate) fn write_at(file: &File, buf: &[u8], offset: u64) -> std::io::Result<()> {
342 #[cfg(unix)]
343 {
344 std::os::unix::fs::FileExt::write_all_at(file, buf, offset)
345 }
346 #[cfg(windows)]
347 {
348 let mut done = 0usize;
349 while done < buf.len() {
350 let n = std::os::windows::fs::FileExt::seek_write(
351 file,
352 &buf[done..],
353 offset + done as u64,
354 )?;
355 if n == 0 {
356 return Err(std::io::ErrorKind::WriteZero.into());
357 }
358 done += n;
359 }
360 Ok(())
361 }
362}
363
364pub(crate) fn sync_dir(dir: &std::path::Path) -> MenteResult<()> {
366 #[cfg(unix)]
367 {
368 File::open(dir)?.sync_all()?;
369 }
370 #[cfg(not(unix))]
371 {
372 let _ = dir;
373 }
374 Ok(())
375}
376
377#[derive(Debug, Clone, PartialEq)]
382pub(crate) enum RecordOp {
383 Write {
387 locator: u64,
388 id: MemoryId,
389 extent: Extent,
390 payload: Vec<u8>,
391 remainder: Option<Extent>,
392 freed: Option<Extent>,
393 },
394 Free { locator: u64, extent: Extent },
396}
397
398const OP_VERSION: u8 = 1;
399
400fn put_extent(buf: &mut Vec<u8>, e: Option<Extent>) {
401 match e {
402 Some(e) => {
403 buf.push(1);
404 buf.extend_from_slice(&e.start.to_le_bytes());
405 buf.extend_from_slice(&e.blocks.to_le_bytes());
406 }
407 None => buf.push(0),
408 }
409}
410
411struct OpReader<'a> {
412 b: &'a [u8],
413 pos: usize,
414}
415
416impl OpReader<'_> {
417 fn take(&mut self, n: usize) -> MenteResult<&[u8]> {
418 let end = self
419 .pos
420 .checked_add(n)
421 .filter(|&e| e <= self.b.len())
422 .ok_or_else(|| MenteError::Storage("record WAL entry truncated".into()))?;
423 let s = &self.b[self.pos..end];
424 self.pos = end;
425 Ok(s)
426 }
427
428 fn u64(&mut self) -> MenteResult<u64> {
429 Ok(le_u64(self.take(8)?))
430 }
431
432 fn extent(&mut self) -> MenteResult<Option<Extent>> {
433 match self.take(1)?[0] {
434 0 => Ok(None),
435 _ => {
436 let start = self.u64()?;
437 let blocks = self.u64()?;
438 Ok(Some(Extent { start, blocks }))
439 }
440 }
441 }
442}
443
444impl RecordOp {
445 pub(crate) fn locator(&self) -> u64 {
446 match self {
447 RecordOp::Write { locator, .. } | RecordOp::Free { locator, .. } => *locator,
448 }
449 }
450
451 pub(crate) fn encode(&self) -> Vec<u8> {
454 let mut buf = vec![OP_VERSION];
455 match self {
456 RecordOp::Write {
457 id,
458 extent,
459 payload,
460 remainder,
461 freed,
462 ..
463 } => {
464 buf.extend_from_slice(id.0.as_bytes());
465 put_extent(&mut buf, Some(*extent));
466 put_extent(&mut buf, *remainder);
467 put_extent(&mut buf, *freed);
468 buf.extend_from_slice(&(payload.len() as u64).to_le_bytes());
469 buf.extend_from_slice(payload);
470 }
471 RecordOp::Free { extent, .. } => {
472 put_extent(&mut buf, Some(*extent));
473 }
474 }
475 buf
476 }
477
478 pub(crate) fn decode(write: bool, locator: u64, data: &[u8]) -> MenteResult<Self> {
480 let mut r = OpReader { b: data, pos: 0 };
481 if r.take(1)?[0] != OP_VERSION {
482 return Err(MenteError::Storage(
483 "record WAL entry has an unknown version".into(),
484 ));
485 }
486 let missing = || MenteError::Storage("record WAL entry without an extent".into());
487 if write {
488 let id = uuid::Uuid::from_slice(r.take(16)?)
489 .map_err(|e| MenteError::Storage(e.to_string()))?;
490 let extent = r.extent()?.ok_or_else(missing)?;
491 let remainder = r.extent()?;
492 let freed = r.extent()?;
493 let len = r.u64()? as usize;
494 let payload = r.take(len)?.to_vec();
495 Ok(RecordOp::Write {
496 locator,
497 id: MemoryId(id),
498 extent,
499 payload,
500 remainder,
501 freed,
502 })
503 } else {
504 let extent = r.extent()?.ok_or_else(missing)?;
505 Ok(RecordOp::Free { locator, extent })
506 }
507 }
508
509 pub(crate) fn redo(&self, file: &File) -> MenteResult<()> {
512 match self {
513 RecordOp::Write {
514 locator,
515 id,
516 extent,
517 payload,
518 remainder,
519 freed,
520 } => {
521 let image = live_image(*locator, *id, payload, extent.blocks);
522 write_at(file, &image, extent.offset())?;
523 if let Some(r) = remainder {
524 write_at(file, &ExtentHeader::free(r.blocks).encode(), r.offset())?;
525 }
526 if let Some(f) = freed {
527 write_at(file, &ExtentHeader::free(f.blocks).encode(), f.offset())?;
528 }
529 }
530 RecordOp::Free { extent, .. } => {
531 write_at(
532 file,
533 &ExtentHeader::free(extent.blocks).encode(),
534 extent.offset(),
535 )?;
536 }
537 }
538 Ok(())
539 }
540}
541
542#[derive(Debug, Default, Clone)]
547pub(crate) struct FreeSpace {
548 by_start: BTreeMap<u64, u64>,
549 by_size: BTreeSet<(u64, u64)>,
550}
551
552#[derive(Debug, Clone, Copy)]
555pub(crate) struct Placement {
556 pub(crate) extent: Extent,
557 pub(crate) remainder: Option<Extent>,
559 source: Option<Extent>,
561}
562
563impl FreeSpace {
564 pub(crate) fn free_blocks(&self) -> u64 {
566 self.by_start.values().sum()
567 }
568
569 pub(crate) fn insert(&mut self, extent: Extent) {
571 let mut start = extent.start;
572 let mut blocks = extent.blocks;
573 if let Some((&ps, &pb)) = self.by_start.range(..start).next_back()
574 && ps + pb == start
575 {
576 self.remove(ps, pb);
577 start = ps;
578 blocks += pb;
579 }
580 if let Some(&nb) = self.by_start.get(&(start + blocks)) {
581 self.remove(start + blocks, nb);
582 blocks += nb;
583 }
584 self.by_start.insert(start, blocks);
585 self.by_size.insert((blocks, start));
586 }
587
588 fn remove(&mut self, start: u64, blocks: u64) {
589 self.by_start.remove(&start);
590 self.by_size.remove(&(blocks, start));
591 }
592
593 pub(crate) fn place(&self, blocks: u64, end: u64) -> Placement {
596 match self.by_size.range((blocks, 0)..).next() {
597 Some(&(size, start)) => Placement {
598 extent: Extent { start, blocks },
599 remainder: (size > blocks).then_some(Extent {
600 start: start + blocks,
601 blocks: size - blocks,
602 }),
603 source: Some(Extent {
604 start,
605 blocks: size,
606 }),
607 },
608 None => Placement {
609 extent: Extent { start: end, blocks },
610 remainder: None,
611 source: None,
612 },
613 }
614 }
615
616 pub(crate) fn commit(&mut self, p: &Placement, end: u64) -> u64 {
618 match p.source {
619 Some(src) => {
620 self.remove(src.start, src.blocks);
621 if let Some(r) = p.remainder {
622 self.by_start.insert(r.start, r.blocks);
623 self.by_size.insert((r.blocks, r.start));
624 }
625 end
626 }
627 None => end.max(p.extent.end()),
628 }
629 }
630
631 pub(crate) fn from_live(mut live: Vec<Extent>, end: u64) -> MenteResult<Self> {
635 live.sort_by_key(|e| e.start);
636 let mut free = FreeSpace::default();
637 let mut pos = 1u64;
638 for e in live {
639 if e.start < pos || e.blocks == 0 {
640 return Err(MenteError::Storage(format!(
641 "record extents overlap at block {}",
642 e.start
643 )));
644 }
645 if e.start > pos {
646 free.insert(Extent {
647 start: pos,
648 blocks: e.start - pos,
649 });
650 }
651 pos = e.end();
652 }
653 if pos > end {
654 return Err(MenteError::Storage(format!(
655 "record extent ends at block {pos}, past the file end {end}"
656 )));
657 }
658 if end > pos {
659 free.insert(Extent {
660 start: pos,
661 blocks: end - pos,
662 });
663 }
664 Ok(free)
665 }
666}
667
668#[cfg(test)]
669mod tests {
670 use super::*;
671 use mentedb_core::memory::MemoryType;
672 use mentedb_core::types::AgentId;
673
674 #[test]
675 fn header_roundtrip_and_corruption() {
676 let h = ExtentHeader {
677 kind: ExtentKind::Live,
678 len: 1234,
679 blocks: 4,
680 locator: 99,
681 id: MemoryId::new(),
682 payload_crc: 7,
683 };
684 let mut b = h.encode();
685 assert_eq!(ExtentHeader::decode(&b), Some(h));
686 b[10] ^= 1;
687 assert_eq!(ExtentHeader::decode(&b), None);
688 }
689
690 #[test]
691 fn payload_roundtrip_is_exact_and_compact() {
692 let mut node = MemoryNode::new(
693 AgentId::new(),
694 MemoryType::Semantic,
695 "x".repeat(300),
696 (0..1024).map(|i| (i as f32).sin() / 7.0).collect(),
697 );
698 node.embedding[3] = f32::NAN;
699 let p = encode_node(&node).unwrap();
700 let json = serde_json::to_vec(&{
701 let mut n = node.clone();
702 n.embedding[3] = 0.0;
703 n
704 })
705 .unwrap();
706 assert!(
707 p.len() * 2 < json.len(),
708 "{} vs json {}",
709 p.len(),
710 json.len()
711 );
712 let back = decode_node(&p).unwrap();
713 assert_eq!(back.id, node.id);
714 assert_eq!(back.content, node.content);
715 assert_eq!(back.embedding.len(), 1024);
716 for (a, b) in back.embedding.iter().zip(&node.embedding) {
717 assert_eq!(a.to_bits(), b.to_bits());
718 }
719 }
720
721 #[test]
722 fn op_roundtrip() {
723 let op = RecordOp::Write {
724 locator: 5,
725 id: MemoryId::new(),
726 extent: Extent {
727 start: 10,
728 blocks: 3,
729 },
730 payload: vec![1, 2, 3],
731 remainder: Some(Extent {
732 start: 13,
733 blocks: 2,
734 }),
735 freed: None,
736 };
737 assert_eq!(RecordOp::decode(true, 5, &op.encode()).unwrap(), op);
738 let free = RecordOp::Free {
739 locator: 5,
740 extent: Extent {
741 start: 10,
742 blocks: 3,
743 },
744 };
745 assert_eq!(RecordOp::decode(false, 5, &free.encode()).unwrap(), free);
746 }
747
748 #[test]
749 fn free_space_best_fit_and_coalesce() {
750 let mut f = FreeSpace::default();
751 f.insert(Extent {
752 start: 10,
753 blocks: 4,
754 });
755 f.insert(Extent {
756 start: 20,
757 blocks: 2,
758 });
759 let p = f.place(2, 100);
761 assert_eq!(p.extent.start, 20);
762 assert!(p.remainder.is_none());
763 assert_eq!(f.commit(&p, 100), 100);
764 let p = f.place(3, 100);
766 assert_eq!(p.extent.start, 10);
767 assert_eq!(
768 p.remainder,
769 Some(Extent {
770 start: 13,
771 blocks: 1
772 })
773 );
774 f.commit(&p, 100);
775 let p = f.place(5, 100);
777 assert_eq!(p.extent.start, 100);
778 assert_eq!(f.commit(&p, 100), 105);
779 f.insert(Extent {
781 start: 11,
782 blocks: 2,
783 });
784 f.insert(Extent {
785 start: 10,
786 blocks: 1,
787 });
788 assert_eq!(f.by_start.get(&10), Some(&4));
789 assert_eq!(f.free_blocks(), 4);
790 }
791
792 #[test]
793 fn free_space_from_live_rejects_overlap() {
794 let e = |start, blocks| Extent { start, blocks };
795 let f = FreeSpace::from_live(vec![e(5, 2), e(1, 2)], 10).unwrap();
796 assert_eq!(f.free_blocks(), 5);
797 assert!(FreeSpace::from_live(vec![e(1, 3), e(2, 2)], 10).is_err());
798 assert!(FreeSpace::from_live(vec![e(8, 3)], 10).is_err());
799 }
800}