1extern crate alloc;
37use alloc::vec;
38use alloc::vec::Vec;
39
40use grafos_std::error::{FabricError, Result};
41use grafos_std::mem::{MemBuilder, MemLease};
42
43use serde::{de::DeserializeOwned, Serialize};
44
45const HEADER_SIZE: u64 = 32;
46
47const BUCKET_EMPTY: u8 = 0;
48const BUCKET_OCCUPIED: u8 = 1;
49const BUCKET_TOMBSTONE: u8 = 2;
50
51pub struct FabricHashMap<K, V> {
85 lease: MemLease,
86 count: u64,
87 bucket_count: u64,
88 key_stride: u64,
89 val_stride: u64,
90 bucket_size: u64,
91 _marker: core::marker::PhantomData<(K, V)>,
92}
93
94impl<K: Serialize + DeserializeOwned + PartialEq, V: Serialize + DeserializeOwned>
95 FabricHashMap<K, V>
96{
97 pub fn new(lease: MemLease, key_stride: usize, val_stride: usize) -> Result<Self> {
110 let arena = lease.mem().arena_size()?;
111 let bucket_size = Self::compute_bucket_size(key_stride as u64, val_stride as u64);
112 if arena < HEADER_SIZE + bucket_size {
113 return Err(FabricError::CapacityExceeded);
114 }
115 let bucket_count = (arena - HEADER_SIZE) / bucket_size;
116 let map = FabricHashMap {
117 lease,
118 count: 0,
119 bucket_count,
120 key_stride: key_stride as u64,
121 val_stride: val_stride as u64,
122 bucket_size,
123 _marker: core::marker::PhantomData,
124 };
125 map.write_header()?;
126 let zero_bucket = vec![0u8; bucket_size as usize];
128 for i in 0..bucket_count {
129 let offset = HEADER_SIZE + i * bucket_size;
130 map.lease.mem().write(offset, &zero_bucket)?;
131 }
132 Ok(map)
133 }
134
135 pub fn from_lease(lease: MemLease) -> Result<Self> {
149 let arena = lease.mem().arena_size()?;
150 if arena < HEADER_SIZE {
151 return Err(FabricError::CapacityExceeded);
152 }
153 let hdr = lease.mem().read(0, HEADER_SIZE as u32)?;
154 let count = u64::from_le_bytes([
155 hdr[0], hdr[1], hdr[2], hdr[3], hdr[4], hdr[5], hdr[6], hdr[7],
156 ]);
157 let bucket_count = u64::from_le_bytes([
158 hdr[8], hdr[9], hdr[10], hdr[11], hdr[12], hdr[13], hdr[14], hdr[15],
159 ]);
160 let key_stride = u64::from_le_bytes([
161 hdr[16], hdr[17], hdr[18], hdr[19], hdr[20], hdr[21], hdr[22], hdr[23],
162 ]);
163 let val_stride = u64::from_le_bytes([
164 hdr[24], hdr[25], hdr[26], hdr[27], hdr[28], hdr[29], hdr[30], hdr[31],
165 ]);
166 let bucket_size = Self::compute_bucket_size(key_stride, val_stride);
167 if arena < HEADER_SIZE + bucket_count * bucket_size {
168 return Err(FabricError::CapacityExceeded);
169 }
170 Ok(FabricHashMap {
171 lease,
172 count,
173 bucket_count,
174 key_stride,
175 val_stride,
176 bucket_size,
177 _marker: core::marker::PhantomData,
178 })
179 }
180
181 pub fn with_capacity(
192 bucket_count: usize,
193 key_stride: usize,
194 val_stride: usize,
195 ) -> Result<Self> {
196 let bucket_size = Self::compute_bucket_size(key_stride as u64, val_stride as u64);
197 let needed = HEADER_SIZE + (bucket_count as u64) * bucket_size;
198 let lease = MemBuilder::new().min_bytes(needed).acquire()?;
199 Self::new(lease, key_stride, val_stride)
200 }
201
202 pub fn insert(&mut self, key: &K, value: &V) -> Result<Option<V>> {
215 let key_bytes = postcard::to_allocvec(key).map_err(|_| FabricError::IoError(-1))?;
216 let val_bytes = postcard::to_allocvec(value).map_err(|_| FabricError::IoError(-1))?;
217
218 if key_bytes.len() as u64 > self.key_stride || val_bytes.len() as u64 > self.val_stride {
219 return Err(FabricError::CapacityExceeded);
220 }
221
222 let hash = Self::hash_bytes(&key_bytes);
223 let mut idx = (hash % self.bucket_count) as usize;
224 let mut first_tombstone: Option<usize> = None;
225
226 for _ in 0..self.bucket_count {
227 let bucket = self.read_bucket(idx)?;
228 match bucket.occupied {
229 BUCKET_EMPTY => {
230 if (self.count + 1) * 4 > self.bucket_count * 3 {
232 return Err(FabricError::CapacityExceeded);
233 }
234 let target = first_tombstone.unwrap_or(idx);
235 self.write_bucket(target, BUCKET_OCCUPIED, hash, &key_bytes, &val_bytes)?;
236 self.count += 1;
237 self.write_header()?;
238 return Ok(None);
239 }
240 BUCKET_TOMBSTONE => {
241 if first_tombstone.is_none() {
242 first_tombstone = Some(idx);
243 }
244 }
245 BUCKET_OCCUPIED => {
246 if bucket.hash == hash {
247 let existing_key: K =
248 postcard::from_bytes(&bucket.key_data[..bucket.key_len as usize])
249 .map_err(|_| FabricError::IoError(-1))?;
250 if &existing_key == key {
251 let old_val: V =
252 postcard::from_bytes(&bucket.val_data[..bucket.val_len as usize])
253 .map_err(|_| FabricError::IoError(-1))?;
254 self.write_bucket(idx, BUCKET_OCCUPIED, hash, &key_bytes, &val_bytes)?;
255 return Ok(Some(old_val));
256 }
257 }
258 }
259 _ => {}
260 }
261 idx = (idx + 1) % self.bucket_count as usize;
262 }
263
264 if let Some(target) = first_tombstone {
266 if (self.count + 1) * 4 > self.bucket_count * 3 {
267 return Err(FabricError::CapacityExceeded);
268 }
269 self.write_bucket(target, BUCKET_OCCUPIED, hash, &key_bytes, &val_bytes)?;
270 self.count += 1;
271 self.write_header()?;
272 return Ok(None);
273 }
274
275 Err(FabricError::CapacityExceeded)
276 }
277
278 pub fn get(&self, key: &K) -> Result<Option<V>> {
286 let key_bytes = postcard::to_allocvec(key).map_err(|_| FabricError::IoError(-1))?;
287 let hash = Self::hash_bytes(&key_bytes);
288 let mut idx = (hash % self.bucket_count) as usize;
289
290 for _ in 0..self.bucket_count {
291 let bucket = self.read_bucket(idx)?;
292 match bucket.occupied {
293 BUCKET_EMPTY => return Ok(None),
294 BUCKET_OCCUPIED => {
295 if bucket.hash == hash {
296 let existing_key: K =
297 postcard::from_bytes(&bucket.key_data[..bucket.key_len as usize])
298 .map_err(|_| FabricError::IoError(-1))?;
299 if &existing_key == key {
300 let val: V =
301 postcard::from_bytes(&bucket.val_data[..bucket.val_len as usize])
302 .map_err(|_| FabricError::IoError(-1))?;
303 return Ok(Some(val));
304 }
305 }
306 }
307 _ => {} }
309 idx = (idx + 1) % self.bucket_count as usize;
310 }
311
312 Ok(None)
313 }
314
315 pub fn remove(&mut self, key: &K) -> Result<Option<V>> {
324 let key_bytes = postcard::to_allocvec(key).map_err(|_| FabricError::IoError(-1))?;
325 let hash = Self::hash_bytes(&key_bytes);
326 let mut idx = (hash % self.bucket_count) as usize;
327
328 for _ in 0..self.bucket_count {
329 let bucket = self.read_bucket(idx)?;
330 match bucket.occupied {
331 BUCKET_EMPTY => return Ok(None),
332 BUCKET_OCCUPIED => {
333 if bucket.hash == hash {
334 let existing_key: K =
335 postcard::from_bytes(&bucket.key_data[..bucket.key_len as usize])
336 .map_err(|_| FabricError::IoError(-1))?;
337 if &existing_key == key {
338 let val: V =
339 postcard::from_bytes(&bucket.val_data[..bucket.val_len as usize])
340 .map_err(|_| FabricError::IoError(-1))?;
341 let zero_key = vec![0u8; 0];
343 let zero_val = vec![0u8; 0];
344 self.write_bucket(idx, BUCKET_TOMBSTONE, 0, &zero_key, &zero_val)?;
345 self.count -= 1;
346 self.write_header()?;
347 return Ok(Some(val));
348 }
349 }
350 }
351 _ => {} }
353 idx = (idx + 1) % self.bucket_count as usize;
354 }
355
356 Ok(None)
357 }
358
359 pub fn contains_key(&self, key: &K) -> Result<bool> {
363 Ok(self.get(key)?.is_some())
364 }
365
366 pub fn len(&self) -> usize {
368 self.count as usize
369 }
370
371 pub fn is_empty(&self) -> bool {
373 self.count == 0
374 }
375
376 pub fn lease_id(&self) -> u128 {
379 self.lease.lease_id()
380 }
381
382 pub fn expires_at_unix_secs(&self) -> u64 {
385 self.lease.expires_at_unix_secs()
386 }
387
388 pub fn iter(&self) -> FabricHashMapIter<'_, K, V> {
393 FabricHashMapIter {
394 map: self,
395 bucket_idx: 0,
396 }
397 }
398
399 fn compute_bucket_size(key_stride: u64, val_stride: u64) -> u64 {
400 1 + 8 + 4 + key_stride + 4 + val_stride
402 }
403
404 fn bucket_offset(&self, idx: usize) -> u64 {
405 HEADER_SIZE + (idx as u64) * self.bucket_size
406 }
407
408 fn read_bucket(&self, idx: usize) -> Result<BucketData> {
409 let offset = self.bucket_offset(idx);
410 let raw = self.lease.mem().read(offset, self.bucket_size as u32)?;
411 let occupied = raw[0];
412 let hash = u64::from_le_bytes([
413 raw[1], raw[2], raw[3], raw[4], raw[5], raw[6], raw[7], raw[8],
414 ]);
415 let key_len = u32::from_le_bytes([raw[9], raw[10], raw[11], raw[12]]);
416 let key_start = 13usize;
417 let key_end = key_start + self.key_stride as usize;
418 let key_data = raw[key_start..key_end].to_vec();
419 let val_len_start = key_end;
420 let val_len = u32::from_le_bytes([
421 raw[val_len_start],
422 raw[val_len_start + 1],
423 raw[val_len_start + 2],
424 raw[val_len_start + 3],
425 ]);
426 let val_start = val_len_start + 4;
427 let val_end = val_start + self.val_stride as usize;
428 let val_data = raw[val_start..val_end].to_vec();
429
430 Ok(BucketData {
431 occupied,
432 hash,
433 key_len,
434 key_data,
435 val_len,
436 val_data,
437 })
438 }
439
440 fn write_bucket(
441 &self,
442 idx: usize,
443 occupied: u8,
444 hash: u64,
445 key_bytes: &[u8],
446 val_bytes: &[u8],
447 ) -> Result<()> {
448 let mut buf = vec![0u8; self.bucket_size as usize];
449 buf[0] = occupied;
450 buf[1..9].copy_from_slice(&hash.to_le_bytes());
451 let key_len = key_bytes.len() as u32;
452 buf[9..13].copy_from_slice(&key_len.to_le_bytes());
453 buf[13..13 + key_bytes.len()].copy_from_slice(key_bytes);
454 let val_len_start = 13 + self.key_stride as usize;
455 let val_len = val_bytes.len() as u32;
456 buf[val_len_start..val_len_start + 4].copy_from_slice(&val_len.to_le_bytes());
457 buf[val_len_start + 4..val_len_start + 4 + val_bytes.len()].copy_from_slice(val_bytes);
458 let offset = self.bucket_offset(idx);
459 self.lease.mem().write(offset, &buf)
460 }
461
462 fn write_header(&self) -> Result<()> {
463 let mut hdr = [0u8; HEADER_SIZE as usize];
464 hdr[0..8].copy_from_slice(&self.count.to_le_bytes());
465 hdr[8..16].copy_from_slice(&self.bucket_count.to_le_bytes());
466 hdr[16..24].copy_from_slice(&self.key_stride.to_le_bytes());
467 hdr[24..32].copy_from_slice(&self.val_stride.to_le_bytes());
468 self.lease.mem().write(0, &hdr)
469 }
470
471 fn hash_bytes(data: &[u8]) -> u64 {
473 let mut hash: u64 = 0xcbf29ce484222325;
474 for &b in data {
475 hash ^= b as u64;
476 hash = hash.wrapping_mul(0x100000001b3);
477 }
478 hash
479 }
480}
481
482struct BucketData {
483 occupied: u8,
484 hash: u64,
485 key_len: u32,
486 key_data: Vec<u8>,
487 val_len: u32,
488 val_data: Vec<u8>,
489}
490
491pub struct FabricHashMapIter<'a, K, V> {
498 map: &'a FabricHashMap<K, V>,
499 bucket_idx: u64,
500}
501
502impl<'a, K: Serialize + DeserializeOwned + PartialEq, V: Serialize + DeserializeOwned> Iterator
503 for FabricHashMapIter<'a, K, V>
504{
505 type Item = Result<(K, V)>;
506
507 fn next(&mut self) -> Option<Self::Item> {
508 while self.bucket_idx < self.map.bucket_count {
509 let idx = self.bucket_idx as usize;
510 self.bucket_idx += 1;
511 match self.map.read_bucket(idx) {
512 Ok(bucket) if bucket.occupied == BUCKET_OCCUPIED => {
513 let key: K =
514 match postcard::from_bytes(&bucket.key_data[..bucket.key_len as usize]) {
515 Ok(k) => k,
516 Err(_) => return Some(Err(FabricError::IoError(-1))),
517 };
518 let val: V =
519 match postcard::from_bytes(&bucket.val_data[..bucket.val_len as usize]) {
520 Ok(v) => v,
521 Err(_) => return Some(Err(FabricError::IoError(-1))),
522 };
523 return Some(Ok((key, val)));
524 }
525 Err(e) => return Some(Err(e)),
526 _ => continue,
527 }
528 }
529 None
530 }
531}
532
533#[cfg(test)]
534mod tests {
535 use super::*;
536 use alloc::string::String;
537 use grafos_std::host;
538
539 fn setup(arena_size: u64) -> MemLease {
540 host::reset_mock();
541 host::mock_set_fbmu_arena_size(arena_size);
542 grafos_std::mem::MemBuilder::new()
543 .acquire()
544 .expect("acquire")
545 }
546
547 #[test]
548 fn insert_get_roundtrip() {
549 let lease = setup(65536);
550 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
551
552 assert!(map.is_empty());
553 assert_eq!(map.insert(&1, &100).expect("insert"), None);
554 assert_eq!(map.insert(&2, &200).expect("insert"), None);
555 assert_eq!(map.insert(&3, &300).expect("insert"), None);
556
557 assert_eq!(map.len(), 3);
558 assert_eq!(map.get(&1).expect("get"), Some(100));
559 assert_eq!(map.get(&2).expect("get"), Some(200));
560 assert_eq!(map.get(&3).expect("get"), Some(300));
561 assert_eq!(map.get(&4).expect("get"), None);
562 }
563
564 #[test]
565 fn insert_overwrites_existing() {
566 let lease = setup(65536);
567 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
568
569 assert_eq!(map.insert(&1, &100).expect("insert"), None);
570 assert_eq!(map.insert(&1, &999).expect("insert"), Some(100));
571 assert_eq!(map.len(), 1);
572 assert_eq!(map.get(&1).expect("get"), Some(999));
573 }
574
575 #[test]
576 fn remove_returns_value() {
577 let lease = setup(65536);
578 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
579
580 map.insert(&1, &100).expect("insert");
581 map.insert(&2, &200).expect("insert");
582
583 assert_eq!(map.remove(&1).expect("remove"), Some(100));
584 assert_eq!(map.len(), 1);
585 assert_eq!(map.get(&1).expect("get"), None);
586 assert_eq!(map.get(&2).expect("get"), Some(200));
587 }
588
589 #[test]
590 fn remove_nonexistent_returns_none() {
591 let lease = setup(65536);
592 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
593 assert_eq!(map.remove(&999).expect("remove"), None);
594 }
595
596 #[test]
597 fn contains_key() {
598 let lease = setup(65536);
599 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
600 map.insert(&42, &1).expect("insert");
601 assert!(map.contains_key(&42).expect("contains"));
602 assert!(!map.contains_key(&99).expect("contains"));
603 }
604
605 #[test]
606 fn string_keys() {
607 let lease = setup(65536);
608 let mut map: FabricHashMap<String, u64> = FabricHashMap::new(lease, 32, 16).expect("new");
609
610 map.insert(&String::from("alpha"), &1).expect("insert");
611 map.insert(&String::from("beta"), &2).expect("insert");
612 map.insert(&String::from("gamma"), &3).expect("insert");
613
614 assert_eq!(map.get(&String::from("beta")).expect("get"), Some(2));
615 }
616
617 #[test]
618 fn collision_handling() {
619 let lease = setup(296);
623 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
624 assert_eq!(map.bucket_count, 8);
625
626 for i in 0..6u32 {
628 map.insert(&i, &(i * 10)).expect("insert");
629 }
630 assert_eq!(map.len(), 6);
631
632 for i in 0..6u32 {
634 assert_eq!(map.get(&i).expect("get"), Some(i * 10));
635 }
636 }
637
638 #[test]
639 fn iter_all_entries() {
640 let lease = setup(65536);
641 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
642
643 map.insert(&1, &10).expect("insert");
644 map.insert(&2, &20).expect("insert");
645 map.insert(&3, &30).expect("insert");
646
647 let mut entries: Vec<(u32, u32)> = map.iter().map(|r| r.expect("iter")).collect();
648 entries.sort_by_key(|&(k, _)| k);
649 assert_eq!(entries, alloc::vec![(1, 10), (2, 20), (3, 30)]);
650 }
651
652 #[test]
653 fn insert_after_remove_reuses_tombstone() {
654 let lease = setup(65536);
655 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
656
657 map.insert(&1, &100).expect("insert");
658 map.remove(&1).expect("remove");
659 assert_eq!(map.len(), 0);
660
661 map.insert(&1, &200).expect("re-insert");
662 assert_eq!(map.get(&1).expect("get"), Some(200));
663 assert_eq!(map.len(), 1);
664 }
665
666 #[test]
667 fn with_capacity_works() {
668 host::reset_mock();
669 host::mock_set_fbmu_arena_size(65536);
670 let map: FabricHashMap<u32, u32> =
671 FabricHashMap::with_capacity(64, 8, 8).expect("with_capacity");
672 assert!(map.bucket_count >= 64);
673 }
674
675 #[test]
676 fn multi_bucket_distribution() {
677 host::reset_mock();
681 host::mock_set_fbmu_arena_size(65536);
682 let mut map: FabricHashMap<u32, u32> =
683 FabricHashMap::with_capacity(64, 8, 8).expect("with_capacity");
684
685 let n = 40u32; for i in 0..n {
687 assert_eq!(map.insert(&i, &(i * 100)).expect("insert"), None);
688 }
689 assert_eq!(map.len(), n as usize);
690
691 for i in 0..n {
693 assert_eq!(map.get(&i).expect("get"), Some(i * 100));
694 }
695
696 let mut occupied_indices: Vec<usize> = Vec::new();
700 for idx in 0..map.bucket_count as usize {
701 let bucket = map.read_bucket(idx).expect("read_bucket");
702 if bucket.occupied == BUCKET_OCCUPIED {
703 occupied_indices.push(idx);
704 }
705 }
706 assert_eq!(occupied_indices.len(), n as usize);
707
708 let min_idx = *occupied_indices.first().unwrap();
711 let max_idx = *occupied_indices.last().unwrap();
712 let span = max_idx - min_idx + 1;
713 assert!(
714 span > map.bucket_count as usize / 2,
715 "expected occupied buckets to span > half the range, got span={} out of {}",
716 span,
717 map.bucket_count
718 );
719 }
720
721 #[test]
722 fn load_factor_limit_enforced() {
723 let bucket_size = FabricHashMap::<u32, u32>::compute_bucket_size(8, 8);
726 let arena = HEADER_SIZE + 8 * bucket_size;
727 let lease = setup(arena);
728 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
729 assert_eq!(map.bucket_count, 8);
730
731 for i in 0..6u32 {
732 map.insert(&i, &(i * 10)).expect("insert");
733 }
734 assert_eq!(map.len(), 6);
735 assert_eq!(
736 map.insert(&6, &60).unwrap_err(),
737 FabricError::CapacityExceeded
738 );
739 }
740
741 #[test]
742 fn remove_then_get_across_probe_chain() {
743 host::reset_mock();
747 host::mock_set_fbmu_arena_size(65536);
748 let mut map: FabricHashMap<u32, u32> =
749 FabricHashMap::with_capacity(64, 8, 8).expect("with_capacity");
750
751 for i in 0..30u32 {
753 map.insert(&i, &(i + 1000)).expect("insert");
754 }
755
756 for i in (0..30u32).step_by(2) {
758 assert_eq!(map.remove(&i).expect("remove"), Some(i + 1000));
759 }
760 assert_eq!(map.len(), 15);
761
762 for i in (1..30u32).step_by(2) {
764 assert_eq!(
765 map.get(&i).expect("get after removals"),
766 Some(i + 1000),
767 "key {} missing after removing even keys",
768 i
769 );
770 }
771
772 for i in (0..30u32).step_by(2) {
774 assert_eq!(map.get(&i).expect("get removed"), None);
775 }
776 }
777
778 #[test]
779 fn iter_skips_tombstones() {
780 let lease = setup(65536);
781 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
782
783 map.insert(&1, &10).expect("insert");
784 map.insert(&2, &20).expect("insert");
785 map.insert(&3, &30).expect("insert");
786 map.remove(&2).expect("remove");
787
788 let mut entries: Vec<(u32, u32)> = map.iter().map(|r| r.expect("iter")).collect();
789 entries.sort_by_key(|&(k, _)| k);
790 assert_eq!(entries, alloc::vec![(1, 10), (3, 30)]);
791 }
792
793 #[test]
794 fn map_lease_accessors() {
795 let lease = setup(65536);
796 let map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
797 assert_ne!(map.lease_id(), 0);
798 assert!(map.expires_at_unix_secs() > 0);
799 }
800
801 #[test]
802 fn insert_reuses_tombstoned_bucket() {
803 let lease = setup(65536);
806 let mut map: FabricHashMap<u32, u32> = FabricHashMap::new(lease, 8, 8).expect("new");
807
808 map.insert(&42, &100).expect("insert");
809 map.remove(&42).expect("remove");
810 assert_eq!(map.len(), 0);
811
812 map.insert(&42, &200).expect("re-insert");
813 assert_eq!(map.len(), 1);
814 assert_eq!(map.get(&42).expect("get"), Some(200));
815
816 let entries: Vec<(u32, u32)> = map.iter().map(|r| r.expect("iter")).collect();
818 assert_eq!(entries.len(), 1);
819 }
820}