1use core::mem::MaybeUninit;
25
26use crate::address::Address;
27use crate::error::ProgramError;
28
29pub const MAX_SEGMENT_BORROWS: usize = 16;
36
37#[derive(Clone, Copy, PartialEq, Eq, Debug)]
39#[repr(u8)]
40pub enum AccessKind {
41 Read = 0,
43 Write = 1,
45}
46
47#[inline(always)]
57fn address_fingerprint(address: &Address) -> u64 {
58 let bytes = address.as_array();
59 u64::from_le_bytes([
60 bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
61 ])
62}
63
64#[inline(always)]
66fn address_eq(a: &Address, b: &Address) -> bool {
67 a.as_array() == b.as_array()
68}
69
70#[inline(always)]
71fn borrow_eq(a: &SegmentBorrow, b: &SegmentBorrow) -> bool {
72 a.key_fp == b.key_fp
73 && address_eq(&a.key, &b.key)
74 && a.offset == b.offset
75 && a.size == b.size
76 && a.kind == b.kind
77}
78
79#[derive(Clone, Copy, Debug)]
86pub struct SegmentBorrow {
87 pub key_fp: u64,
89 pub key: Address,
94 pub offset: u32,
96 pub size: u32,
98 pub kind: AccessKind,
100}
101
102#[inline(always)]
104const fn ranges_overlap(a_off: u32, a_size: u32, b_off: u32, b_size: u32) -> bool {
105 let a_end = a_off as u64 + a_size as u64;
106 let b_end = b_off as u64 + b_size as u64;
107 !(a_end <= b_off as u64 || b_end <= a_off as u64)
109}
110
111pub struct SegmentBorrowRegistry {
129 entries: [MaybeUninit<SegmentBorrow>; MAX_SEGMENT_BORROWS],
130 len: u8,
131}
132
133impl SegmentBorrowRegistry {
134 #[inline(always)]
136 pub const fn new() -> Self {
137 const EMPTY: MaybeUninit<SegmentBorrow> = MaybeUninit::uninit();
138 Self {
139 entries: [EMPTY; MAX_SEGMENT_BORROWS],
140 len: 0,
141 }
142 }
143
144 #[inline(always)]
146 pub const fn len(&self) -> usize {
147 self.len as usize
148 }
149
150 #[inline(always)]
152 pub const fn is_empty(&self) -> bool {
153 self.len == 0
154 }
155
156 #[inline(always)]
161 pub fn register_leased_read(
162 &mut self,
163 key: &Address,
164 offset: u32,
165 size: u32,
166 ) -> Result<SegmentBorrow, ProgramError> {
167 let borrow = SegmentBorrow {
168 key_fp: address_fingerprint(key),
169 key: *key,
170 offset,
171 size,
172 kind: AccessKind::Read,
173 };
174 self.register(borrow)?;
175 Ok(borrow)
176 }
177
178 #[inline(always)]
180 pub fn register_leased_write(
181 &mut self,
182 key: &Address,
183 offset: u32,
184 size: u32,
185 ) -> Result<SegmentBorrow, ProgramError> {
186 let borrow = SegmentBorrow {
187 key_fp: address_fingerprint(key),
188 key: *key,
189 offset,
190 size,
191 kind: AccessKind::Write,
192 };
193 self.register(borrow)?;
194 Ok(borrow)
195 }
196
197 #[inline(always)]
203 pub fn register(&mut self, new: SegmentBorrow) -> Result<(), ProgramError> {
204 let len = self.len as usize;
205 if len >= MAX_SEGMENT_BORROWS {
206 return Err(ProgramError::AccountBorrowFailed);
207 }
208
209 let mut i = 0;
214 while i < len {
215 let existing = unsafe { self.entries.get_unchecked(i).assume_init_ref() };
218 if existing.key_fp == new.key_fp
219 && address_eq(&existing.key, &new.key)
220 && ranges_overlap(existing.offset, existing.size, new.offset, new.size)
221 {
222 match (existing.kind, new.kind) {
223 (AccessKind::Read, AccessKind::Read) => {}
224 _ => return Err(ProgramError::AccountBorrowFailed),
225 }
226 }
227 i += 1;
228 }
229
230 unsafe { self.entries.get_unchecked_mut(len).write(new) };
233 self.len = (len + 1) as u8;
234 Ok(())
235 }
236
237 #[inline(always)]
239 pub fn register_read(
240 &mut self,
241 key: &Address,
242 offset: u32,
243 size: u32,
244 ) -> Result<(), ProgramError> {
245 self.register(SegmentBorrow {
246 key_fp: address_fingerprint(key),
247 key: *key,
248 offset,
249 size,
250 kind: AccessKind::Read,
251 })
252 }
253
254 #[inline(always)]
256 pub fn register_write(
257 &mut self,
258 key: &Address,
259 offset: u32,
260 size: u32,
261 ) -> Result<(), ProgramError> {
262 self.register(SegmentBorrow {
263 key_fp: address_fingerprint(key),
264 key: *key,
265 offset,
266 size,
267 kind: AccessKind::Write,
268 })
269 }
270
271 #[inline(always)]
276 pub fn release(&mut self, borrow: &SegmentBorrow) -> bool {
277 let len = self.len as usize;
278 let mut i = 0;
279 while i < len {
280 let existing = unsafe { self.entries.get_unchecked(i).assume_init_ref() };
282 if borrow_eq(existing, borrow) {
283 let new_len = len - 1;
285 self.len = new_len as u8;
286 if i < new_len {
287 let last = unsafe { self.entries.get_unchecked(new_len).assume_init() };
290 unsafe { self.entries.get_unchecked_mut(i).write(last) };
292 }
293 return true;
294 }
295 i += 1;
296 }
297 false
298 }
299
300 #[doc(hidden)]
312 #[inline(always)]
313 pub unsafe fn release_last_registered(&mut self, borrow: &SegmentBorrow) -> bool {
314 let len = self.len as usize;
315 if len == 0 {
316 return false;
317 }
318 let last = unsafe { *self.entries.get_unchecked(len - 1).assume_init_ref() };
320 if !borrow_eq(&last, borrow) {
321 return self.release(borrow);
322 }
323 self.len = (len - 1) as u8;
324 true
325 }
326
327 #[inline(always)]
329 pub fn clear(&mut self) {
330 self.len = 0;
331 }
332
333 #[inline(always)]
338 pub fn would_conflict(&self, proposed: &SegmentBorrow) -> bool {
339 let len = self.len as usize;
340 let mut i = 0;
341 while i < len {
342 let existing = unsafe { self.entries.get_unchecked(i).assume_init_ref() };
344 if existing.key_fp == proposed.key_fp
345 && address_eq(&existing.key, &proposed.key)
346 && ranges_overlap(
347 existing.offset,
348 existing.size,
349 proposed.offset,
350 proposed.size,
351 )
352 {
353 match (existing.kind, proposed.kind) {
354 (AccessKind::Read, AccessKind::Read) => {}
355 _ => return true,
356 }
357 }
358 i += 1;
359 }
360 false
361 }
362
363 #[inline(always)]
378 pub fn register_guard(
379 &mut self,
380 borrow: SegmentBorrow,
381 ) -> Result<SegmentBorrowGuard<'_>, ProgramError> {
382 self.register(borrow)?;
383 Ok(SegmentBorrowGuard {
384 registry: self,
385 borrow,
386 })
387 }
388
389 #[inline(always)]
391 pub fn register_guard_read(
392 &mut self,
393 key: &Address,
394 offset: u32,
395 size: u32,
396 ) -> Result<SegmentBorrowGuard<'_>, ProgramError> {
397 let borrow = SegmentBorrow {
398 key_fp: address_fingerprint(key),
399 key: *key,
400 offset,
401 size,
402 kind: AccessKind::Read,
403 };
404 self.register_guard(borrow)
405 }
406
407 #[inline(always)]
409 pub fn register_guard_write(
410 &mut self,
411 key: &Address,
412 offset: u32,
413 size: u32,
414 ) -> Result<SegmentBorrowGuard<'_>, ProgramError> {
415 let borrow = SegmentBorrow {
416 key_fp: address_fingerprint(key),
417 key: *key,
418 offset,
419 size,
420 kind: AccessKind::Write,
421 };
422 self.register_guard(borrow)
423 }
424
425 #[inline]
430 pub fn for_each<F: FnMut(&SegmentBorrow)>(&self, mut f: F) {
431 let len = self.len as usize;
432 let mut i = 0;
433 while i < len {
434 f(unsafe { self.entries.get_unchecked(i).assume_init_ref() });
436 i += 1;
437 }
438 }
439
440 #[inline]
442 pub fn find_exact(
443 &self,
444 key: &Address,
445 offset: u32,
446 size: u32,
447 kind: AccessKind,
448 ) -> Option<&SegmentBorrow> {
449 let fp = address_fingerprint(key);
450 let len = self.len as usize;
451 let mut i = 0;
452 while i < len {
453 let e = unsafe { self.entries.get_unchecked(i).assume_init_ref() };
455 if e.key_fp == fp
456 && address_eq(&e.key, key)
457 && e.offset == offset
458 && e.size == size
459 && e.kind == kind
460 {
461 return Some(e);
462 }
463 i += 1;
464 }
465 None
466 }
467}
468
469pub struct SegmentBorrowGuard<'a> {
475 registry: &'a mut SegmentBorrowRegistry,
476 borrow: SegmentBorrow,
477}
478
479impl<'a> SegmentBorrowGuard<'a> {
480 #[inline(always)]
482 pub fn kind(&self) -> AccessKind {
483 self.borrow.kind
484 }
485
486 #[inline(always)]
488 pub fn offset(&self) -> u32 {
489 self.borrow.offset
490 }
491
492 #[inline(always)]
494 pub fn size(&self) -> u32 {
495 self.borrow.size
496 }
497}
498
499impl<'a> Drop for SegmentBorrowGuard<'a> {
500 fn drop(&mut self) {
501 self.registry.release(&self.borrow);
502 }
503}
504
505#[cfg(kani)]
506mod kani_proofs {
507 use super::*;
508
509 #[kani::proof]
510 fn range_overlap_is_symmetric_for_arbitrary_u32s() {
511 let a_off: u32 = kani::any();
512 let a_size: u32 = kani::any();
513 let b_off: u32 = kani::any();
514 let b_size: u32 = kani::any();
515
516 assert_eq!(
517 ranges_overlap(a_off, a_size, b_off, b_size),
518 ranges_overlap(b_off, b_size, a_off, a_size)
519 );
520 }
521
522 #[kani::proof]
523 fn overlapping_write_blocks_same_account_accesses() {
524 let offset: u32 = kani::any();
525 let size: u32 = kani::any();
526 let delta: u32 = kani::any();
527 kani::assume(offset <= 1024);
528 kani::assume(size > 0 && size <= 64);
529 kani::assume(delta < size);
530
531 let key = Address::new([7u8; 32]);
532 let probe_offset = offset + delta;
533 let mut reg = SegmentBorrowRegistry::new();
534
535 assert!(reg.register_write(&key, offset, size).is_ok());
536 assert!(reg.register_read(&key, probe_offset, 1).is_err());
537 assert!(reg.register_write(&key, probe_offset, 1).is_err());
538 assert_eq!(reg.len(), 1);
539 }
540
541 #[kani::proof]
542 fn overlapping_reads_are_shared_for_same_account() {
543 let offset: u32 = kani::any();
544 let size: u32 = kani::any();
545 let delta: u32 = kani::any();
546 kani::assume(offset <= 1024);
547 kani::assume(size > 0 && size <= 64);
548 kani::assume(delta < size);
549
550 let key = Address::new([8u8; 32]);
551 let probe_offset = offset + delta;
552 let mut reg = SegmentBorrowRegistry::new();
553
554 assert!(reg.register_read(&key, offset, size).is_ok());
555 assert!(reg.register_read(&key, probe_offset, 1).is_ok());
556 assert_eq!(reg.len(), 2);
557 }
558
559 #[kani::proof]
560 fn fingerprint_collision_different_addresses_do_not_conflict() {
561 let key_a = Address::new([9u8; 32]);
562 let mut key_b_bytes = [9u8; 32];
563 key_b_bytes[8] = 10;
564 let key_b = Address::new(key_b_bytes);
565 let mut reg = SegmentBorrowRegistry::new();
566
567 assert_eq!(address_fingerprint(&key_a), address_fingerprint(&key_b));
568 assert_ne!(key_a.as_array(), key_b.as_array());
569 assert!(reg.register_write(&key_a, 0, 8).is_ok());
570 assert!(reg.register_write(&key_b, 0, 8).is_ok());
571 assert_eq!(reg.len(), 2);
572 }
573
574 #[kani::proof]
575 fn release_removes_exact_borrow_and_preserves_others() {
576 let key = Address::new([11u8; 32]);
577 let mut reg = SegmentBorrowRegistry::new();
578
579 let first = reg.register_leased_read(&key, 0, 8).unwrap();
580 let second = reg.register_leased_write(&key, 8, 8).unwrap();
581 assert!(reg.release(&first));
582
583 assert_eq!(reg.len(), 1);
584 assert!(reg.find_exact(&key, 0, 8, AccessKind::Read).is_none());
585 assert!(reg.find_exact(&key, 8, 8, AccessKind::Write).is_some());
586 assert!(reg.release(&second));
587 assert!(reg.is_empty());
588 }
589}
590
591#[cfg(test)]
594mod tests {
595 use super::*;
596 use crate::Address;
597
598 fn test_addr(seed: u8) -> Address {
599 Address::new([seed; 32])
600 }
601
602 #[test]
603 fn read_read_same_range_allowed() {
604 let mut reg = SegmentBorrowRegistry::new();
605 let key = test_addr(1);
606 assert!(reg.register_read(&key, 0, 8).is_ok());
607 assert!(reg.register_read(&key, 0, 8).is_ok());
608 assert_eq!(reg.len(), 2);
609 }
610
611 #[test]
612 fn read_write_same_range_rejected() {
613 let mut reg = SegmentBorrowRegistry::new();
614 let key = test_addr(1);
615 assert!(reg.register_read(&key, 0, 8).is_ok());
616 assert!(reg.register_write(&key, 0, 8).is_err());
617 }
618
619 #[test]
620 fn write_write_same_range_rejected() {
621 let mut reg = SegmentBorrowRegistry::new();
622 let key = test_addr(1);
623 assert!(reg.register_write(&key, 0, 8).is_ok());
624 assert!(reg.register_write(&key, 0, 8).is_err());
625 }
626
627 #[test]
628 fn write_read_same_range_rejected() {
629 let mut reg = SegmentBorrowRegistry::new();
630 let key = test_addr(1);
631 assert!(reg.register_write(&key, 0, 8).is_ok());
632 assert!(reg.register_read(&key, 0, 8).is_err());
633 }
634
635 #[test]
636 fn non_overlapping_write_write_allowed() {
637 let mut reg = SegmentBorrowRegistry::new();
638 let key = test_addr(1);
639 assert!(reg.register_write(&key, 0, 8).is_ok());
641 assert!(reg.register_write(&key, 8, 32).is_ok());
642 }
643
644 #[test]
645 fn partially_overlapping_rejected() {
646 let mut reg = SegmentBorrowRegistry::new();
647 let key = test_addr(1);
648 assert!(reg.register_write(&key, 0, 16).is_ok());
650 assert!(reg.register_write(&key, 8, 16).is_err());
651 }
652
653 #[test]
654 fn different_accounts_always_allowed() {
655 let mut reg = SegmentBorrowRegistry::new();
656 assert!(reg.register_write(&test_addr(1), 0, 8).is_ok());
657 assert!(reg.register_write(&test_addr(2), 0, 8).is_ok());
658 }
659
660 #[test]
661 fn release_then_reacquire() {
662 let mut reg = SegmentBorrowRegistry::new();
663 let key = test_addr(1);
664 let borrow = SegmentBorrow {
665 key_fp: address_fingerprint(&key),
666 key,
667 offset: 0,
668 size: 8,
669 kind: AccessKind::Write,
670 };
671 assert!(reg.register(borrow).is_ok());
672 assert!(reg.register_write(&key, 0, 8).is_err()); assert!(reg.release(&borrow));
674 assert!(reg.register_write(&key, 0, 8).is_ok()); }
676
677 #[test]
678 fn release_last_registered_falls_back_to_exact_release() {
679 let mut reg = SegmentBorrowRegistry::new();
680 let key = test_addr(1);
681 let first = reg.register_leased_read(&key, 0, 8).unwrap();
682 let second = reg.register_leased_write(&key, 8, 8).unwrap();
683
684 assert!(unsafe { reg.release_last_registered(&first) });
687 assert_eq!(reg.len(), 1);
688 assert!(reg.find_exact(&key, 0, 8, AccessKind::Read).is_none());
689 assert!(reg.find_exact(&key, 8, 8, AccessKind::Write).is_some());
690
691 assert!(unsafe { reg.release_last_registered(&second) });
694 assert!(reg.is_empty());
695 }
696
697 #[test]
698 fn capacity_limit() {
699 let mut reg = SegmentBorrowRegistry::new();
700 for i in 0..MAX_SEGMENT_BORROWS {
701 assert!(reg.register_read(&test_addr(1), i as u32 * 8, 8).is_ok());
702 }
703 assert!(reg.register_read(&test_addr(1), 256, 8).is_err());
705 }
706
707 #[test]
708 fn would_conflict_does_not_mutate() {
709 let mut reg = SegmentBorrowRegistry::new();
710 let key = test_addr(1);
711 assert!(reg.register_write(&key, 0, 8).is_ok());
712 let proposed = SegmentBorrow {
713 key_fp: address_fingerprint(&key),
714 key,
715 offset: 0,
716 size: 8,
717 kind: AccessKind::Write,
718 };
719 assert!(reg.would_conflict(&proposed));
720 assert_eq!(reg.len(), 1); }
722
723 #[test]
724 fn adjacent_ranges_no_conflict() {
725 let mut reg = SegmentBorrowRegistry::new();
726 let key = test_addr(1);
727 assert!(reg.register_write(&key, 0, 8).is_ok());
729 assert!(reg.register_write(&key, 8, 8).is_ok());
730 }
731
732 #[test]
741 fn guard_auto_releases_write_on_drop() {
742 let mut reg = SegmentBorrowRegistry::new();
743 let key = test_addr(1);
744 {
745 let _guard = reg.register_guard_write(&key, 0, 8).unwrap();
746 }
748 assert_eq!(reg.len(), 0);
750 assert!(reg.register_write(&key, 0, 8).is_ok());
752 }
753
754 #[test]
755 fn guard_auto_releases_read_on_drop() {
756 let mut reg = SegmentBorrowRegistry::new();
757 let key = test_addr(1);
758 {
759 let _guard = reg.register_guard_read(&key, 0, 8).unwrap();
760 }
761 assert_eq!(reg.len(), 0);
762 assert!(reg.register_write(&key, 0, 8).is_ok());
764 }
765
766 #[test]
767 fn sequential_guards_reuse_slot() {
768 let mut reg = SegmentBorrowRegistry::new();
769 let key = test_addr(1);
770 for _ in 0..4 {
771 let _guard = reg.register_guard_write(&key, 0, 8).unwrap();
772 }
774 assert_eq!(reg.len(), 0);
775 }
776
777 #[test]
778 fn guard_accessors() {
779 let mut reg = SegmentBorrowRegistry::new();
780 let key = test_addr(1);
781 let guard = reg.register_guard_write(&key, 16, 32).unwrap();
782 assert_eq!(guard.kind(), AccessKind::Write);
783 assert_eq!(guard.offset(), 16);
784 assert_eq!(guard.size(), 32);
785 }
786
787 #[test]
788 fn guard_then_manual_register_ok() {
789 let mut reg = SegmentBorrowRegistry::new();
790 let key = test_addr(1);
791 {
792 let _guard = reg.register_guard_write(&key, 0, 8).unwrap();
793 }
794 assert!(reg.register_read(&key, 0, 8).is_ok());
796 assert_eq!(reg.len(), 1);
797 }
798}