1#![expect(
2 clippy::rest_pattern_accessible_field,
3 reason = "Intentionally not needing other fields, too verbose otherwise"
4)]
5#![no_std]
6
7extern crate alloc;
8
9use ab_aligned_buffer::{OwnedAlignedBuffer, SharedAlignedBuffer};
10use ab_core_primitives::address::Address;
11use alloc::boxed::Box;
12use replace_with::replace_with_or_abort;
13use smallvec::SmallVec;
14use tracing::debug;
15
16const INLINE_SIZE: usize = 8;
21const NEW_CONTRACTS_INLINE: usize = 2;
23
24#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
26pub struct SlotKey {
27 pub owner: Address,
29 pub contract: Address,
31}
32
33#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
35pub struct SlotIndex(usize);
36
37impl From<SlotIndex> for usize {
38 #[inline(always)]
39 fn from(value: SlotIndex) -> Self {
40 value.0
41 }
42}
43
44#[derive(Debug, Clone)]
45pub enum Slot {
46 ReadOnly {
47 key: SlotKey,
48 buffer: SharedAlignedBuffer,
49 },
50 ReadWrite {
51 key: SlotKey,
52 buffer: SharedAlignedBuffer,
53 },
54}
55
56impl Slot {
57 fn is_null_contract(&self) -> bool {
58 let slot_key = match self {
59 Slot::ReadOnly { key, .. } => key,
60 Slot::ReadWrite { key, .. } => key,
61 };
62
63 slot_key.contract == Address::NULL
64 }
65}
66
67#[derive(Debug, Clone)]
68enum SlotState {
69 Original(SharedAlignedBuffer),
71 OriginalReadOnly(SharedAlignedBuffer),
73 Modified(SharedAlignedBuffer),
75 ModifiedReadOnly(SharedAlignedBuffer),
77 OriginalReadWrite {
79 buffer: OwnedAlignedBuffer,
80 previous: SharedAlignedBuffer,
82 },
83 ModifiedReadWrite {
85 buffer: OwnedAlignedBuffer,
86 previous: SharedAlignedBuffer,
88 },
89}
90
91#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
92struct SlotAccess {
93 slot_index: SlotIndex,
94 read_write: bool,
96}
97
98#[derive(Debug, Clone)]
99struct Inner {
100 slots: SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
101 slot_access: SmallVec<[SlotAccess; INLINE_SIZE]>,
102 new_contracts: SmallVec<[Address; NEW_CONTRACTS_INLINE]>,
108}
109
110#[derive(Debug, Clone)]
112pub struct Slots(Box<Inner>);
113
114impl Slots {
115 #[inline(always)]
122 pub fn new<I>(slots: I) -> Self
123 where
124 I: IntoIterator<Item = Slot>,
125 {
126 let slots = slots
127 .into_iter()
128 .filter_map(|slot| {
129 if slot.is_null_contract() {
133 return None;
134 }
135
136 Some(match slot {
137 Slot::ReadOnly { key, buffer } => {
138 (key, SlotState::OriginalReadOnly(buffer))
140 }
141 Slot::ReadWrite { key, buffer } => (key, SlotState::Original(buffer)),
142 })
143 })
144 .collect();
145
146 let inner = Inner {
147 slots,
148 slot_access: SmallVec::new(),
149 new_contracts: SmallVec::new(),
150 };
151
152 Self(Box::new(inner))
153 }
154
155 #[inline(always)]
160 pub fn new_nested_rw(&mut self) -> NestedSlots<'_> {
161 NestedSlots(NestedSlotsInner::ReadWrite {
162 inner: &mut self.0,
163 parent_slot_access_len: 0,
164 original_parent: true,
165 })
166 }
167
168 #[inline(always)]
170 pub fn new_nested_ro(&self) -> NestedSlots<'_> {
171 NestedSlots(NestedSlotsInner::ReadOnly { inner: &self.0 })
172 }
173
174 #[must_use]
182 #[inline(always)]
183 pub fn add_new_contract(&mut self, owner: Address) -> bool {
184 let new_contracts = &mut self.0.new_contracts;
185
186 if new_contracts.contains(&owner) {
187 debug!(?owner, "Not adding a new contract duplicate");
188 return false;
189 }
190
191 new_contracts.push(owner);
192 true
193 }
194
195 #[inline]
197 pub fn iter(&self) -> impl ExactSizeIterator<Item = (&SlotKey, &SharedAlignedBuffer)> + '_ {
198 self.0.slots.iter().map(|(slot_key, slot)| match slot {
199 SlotState::Original(buffer) => (slot_key, buffer),
200 SlotState::OriginalReadOnly(_) => {
201 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
202 }
203 SlotState::Modified(buffer) => (slot_key, buffer),
204 SlotState::ModifiedReadOnly(_) => {
205 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
206 }
207 SlotState::OriginalReadWrite { .. } => {
208 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
209 }
210 SlotState::ModifiedReadWrite { .. } => {
211 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
212 }
213 })
214 }
215
216 #[inline]
218 pub fn iter_modified(&self) -> impl Iterator<Item = (&SlotKey, &SharedAlignedBuffer)> + '_ {
219 self.0
220 .slots
221 .iter()
222 .filter_map(|(slot_key, slot)| match slot {
223 SlotState::Original(_) => None,
224 SlotState::OriginalReadOnly(_) => unreachable!(
225 "Only original and modified slots can exist at the `Slots` level; qed"
226 ),
227 SlotState::Modified(buffer) => Some((slot_key, buffer)),
228 SlotState::ModifiedReadOnly(_) => unreachable!(
229 "Only original and modified slots can exist at the `Slots` level; qed"
230 ),
231 SlotState::OriginalReadWrite { .. } => unreachable!(
232 "Only original and modified slots can exist at the `Slots` level; qed"
233 ),
234 SlotState::ModifiedReadWrite { .. } => unreachable!(
235 "Only original and modified slots can exist at the `Slots` level; qed"
236 ),
237 })
238 }
239
240 #[inline]
242 pub fn into_slots(self) -> impl ExactSizeIterator<Item = (SlotKey, SharedAlignedBuffer)> {
243 self.0.slots.into_iter().map(|(slot_key, slot)| match slot {
244 SlotState::Original(buffer) => (slot_key, buffer),
245 SlotState::OriginalReadOnly(_) => {
246 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
247 }
248 SlotState::Modified(buffer) => (slot_key, buffer),
249 SlotState::ModifiedReadOnly(_) => {
250 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
251 }
252 SlotState::OriginalReadWrite { .. } => {
253 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
254 }
255 SlotState::ModifiedReadWrite { .. } => {
256 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
257 }
258 })
259 }
260}
261
262#[derive(Debug)]
264enum NestedSlotsInner<'a> {
265 ReadWrite {
267 inner: &'a mut Inner,
268 parent_slot_access_len: usize,
269 original_parent: bool,
270 },
271 ReadOnly { inner: &'a Inner },
273}
274
275#[derive(Debug)]
276pub struct NestedSlots<'a>(NestedSlotsInner<'a>);
277
278impl Drop for NestedSlots<'_> {
279 #[inline(always)]
280 fn drop(&mut self) {
281 let (inner, parent_slot_access_len, original_parent) = match &mut self.0 {
282 NestedSlotsInner::ReadWrite {
283 inner,
284 parent_slot_access_len,
285 original_parent,
286 } => (&mut **inner, *parent_slot_access_len, *original_parent),
287 NestedSlotsInner::ReadOnly { .. } => {
288 return;
290 }
291 };
292
293 let slots = &mut inner.slots;
294 let slot_access = &mut inner.slot_access;
295
296 for slot_access in slot_access.drain(parent_slot_access_len..) {
298 let slot = &mut slots
299 .get_mut(usize::from(slot_access.slot_index))
300 .expect("Accessed slot exists; qed")
301 .1;
302
303 replace_with_or_abort(slot, |slot| match slot {
304 SlotState::Original(_buffer) => {
305 unreachable!("Slot can't be in `Original` state after being accessed; qed")
306 }
307 SlotState::OriginalReadOnly(buffer) => SlotState::Original(buffer),
308 SlotState::Modified(buffer) => SlotState::Modified(buffer),
309 SlotState::ModifiedReadOnly(buffer) => SlotState::Modified(buffer),
310 SlotState::OriginalReadWrite { buffer, .. }
311 | SlotState::ModifiedReadWrite { buffer, .. } => {
312 SlotState::Modified(buffer.into_shared())
313 }
314 });
315 }
316
317 if original_parent {
318 inner
322 .slots
323 .retain(|(slot_key, _slot)| slot_key.contract != Address::NULL);
324 }
325 }
326}
327
328impl<'a> NestedSlots<'a> {
329 #[inline(always)]
330 fn inner_ro(&self) -> &Inner {
331 match &self.0 {
332 NestedSlotsInner::ReadWrite { inner, .. } => inner,
333 NestedSlotsInner::ReadOnly { inner } => inner,
334 }
335 }
336
337 #[inline(always)]
338 fn inner_rw(&mut self) -> Option<&mut Inner> {
339 match &mut self.0 {
340 NestedSlotsInner::ReadWrite { inner, .. } => Some(inner),
341 NestedSlotsInner::ReadOnly { .. } => None,
342 }
343 }
344
345 #[inline(always)]
352 pub fn new_nested_rw<'b>(&'b mut self) -> Option<NestedSlots<'b>>
353 where
354 'a: 'b,
355 {
356 let inner = match &mut self.0 {
357 NestedSlotsInner::ReadWrite { inner, .. } => &mut **inner,
358 NestedSlotsInner::ReadOnly { .. } => {
359 return None;
360 }
361 };
362
363 let parent_slot_access_len = inner.slot_access.len();
364
365 Some(NestedSlots(NestedSlotsInner::ReadWrite {
366 inner,
367 parent_slot_access_len,
368 original_parent: false,
369 }))
370 }
371
372 #[inline(always)]
374 pub fn new_nested_ro<'b>(&'b self) -> NestedSlots<'b>
375 where
376 'a: 'b,
377 {
378 let inner = match &self.0 {
379 NestedSlotsInner::ReadWrite { inner, .. } => &**inner,
380 NestedSlotsInner::ReadOnly { inner } => &**inner,
381 };
382
383 NestedSlots(NestedSlotsInner::ReadOnly { inner })
384 }
385
386 #[must_use]
394 #[inline(always)]
395 pub fn add_new_contract(&mut self, owner: Address) -> bool {
396 let Some(inner) = self.inner_rw() else {
397 debug!(?owner, "`add_new_contract` access violation");
398 return false;
399 };
400
401 let new_contracts = &mut inner.new_contracts;
402
403 if new_contracts.contains(&owner) {
404 debug!(?owner, "Not adding a new contract duplicate");
405 return false;
406 }
407
408 new_contracts.push(owner);
409 true
410 }
411
412 #[inline(always)]
419 pub fn get_code(&self, owner: Address) -> Option<SharedAlignedBuffer> {
420 let result = self.get_code_internal(owner);
421
422 if result.is_none() {
423 debug!(?owner, "`get_code` access violation");
424 }
425
426 result
427 }
428
429 #[inline(always)]
430 fn get_code_internal(&self, owner: Address) -> Option<SharedAlignedBuffer> {
431 let inner = self.inner_ro();
432 let slots = &inner.slots;
433 let slot_access = &inner.slot_access;
434
435 let contract = Address::SYSTEM_CODE;
436
437 let slot_index = slots.iter().position(|(slot_key, _slot)| {
438 slot_key.owner == owner && slot_key.contract == contract
439 })?;
440 let slot_index = SlotIndex(slot_index);
441
442 if slot_access
444 .iter()
445 .any(|slot_access| slot_access.slot_index == slot_index && slot_access.read_write)
446 {
447 return None;
448 }
449
450 let buffer = match &slots
451 .get(usize::from(slot_index))
452 .expect("Just found; qed")
453 .1
454 {
455 SlotState::Original(buffer)
456 | SlotState::OriginalReadOnly(buffer)
457 | SlotState::Modified(buffer)
458 | SlotState::ModifiedReadOnly(buffer) => buffer,
459 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => {
460 return None;
461 }
462 };
463
464 Some(buffer.clone())
465 }
466
467 #[inline(always)]
471 pub fn use_ro(&mut self, slot_key: SlotKey) -> Option<&SharedAlignedBuffer> {
472 let inner_rw = match &mut self.0 {
473 NestedSlotsInner::ReadWrite { inner, .. } => &mut **inner,
474 NestedSlotsInner::ReadOnly { inner } => {
475 let result = Self::use_ro_internal_read_only(
477 slot_key,
478 &inner.slots,
479 &inner.slot_access,
480 &inner.new_contracts,
481 );
482
483 if result.is_none() {
484 debug!(?slot_key, "`use_ro` access violation");
485 }
486
487 return result;
488 }
489 };
490
491 let result = Self::use_ro_internal(
492 slot_key,
493 &mut inner_rw.slots,
494 &mut inner_rw.slot_access,
495 &inner_rw.new_contracts,
496 );
497
498 if result.is_none() {
499 debug!(?slot_key, "`use_ro` access violation");
500 }
501
502 result
503 }
504
505 #[inline(always)]
506 fn use_ro_internal<'b>(
507 slot_key: SlotKey,
508 slots: &'b mut SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
509 slot_access: &mut SmallVec<[SlotAccess; INLINE_SIZE]>,
510 new_contracts: &[Address],
511 ) -> Option<&'b SharedAlignedBuffer> {
512 let maybe_slot_index = slots
513 .iter()
514 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
515 .map(SlotIndex);
516
517 if let Some(slot_index) = maybe_slot_index {
518 if let Some(read_write) = slot_access.iter().find_map(|slot_access| {
520 (slot_access.slot_index == slot_index).then_some(slot_access.read_write)
521 }) {
522 if read_write {
523 return None;
524 }
525 } else {
526 slot_access.push(SlotAccess {
527 slot_index,
528 read_write: false,
529 });
530 }
531
532 let slot = &mut slots
533 .get_mut(usize::from(slot_index))
534 .expect("Just found; qed")
535 .1;
536
537 match slot {
539 SlotState::Original(buffer) => {
540 let buffer = buffer.clone();
541 *slot = SlotState::OriginalReadOnly(buffer);
542 let SlotState::OriginalReadOnly(buffer) = slot else {
543 unreachable!("Just inserted; qed");
544 };
545 Some(buffer)
546 }
547 SlotState::OriginalReadOnly(buffer) | SlotState::ModifiedReadOnly(buffer) => {
548 Some(buffer)
549 }
550 SlotState::Modified(buffer) => {
551 let buffer = buffer.clone();
552 *slot = SlotState::ModifiedReadOnly(buffer);
553 let SlotState::ModifiedReadOnly(buffer) = slot else {
554 unreachable!("Just inserted; qed");
555 };
556 Some(buffer)
557 }
558 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => None,
559 }
560 } else {
561 if !(slot_key.contract == Address::NULL
565 || new_contracts
566 .iter()
567 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
568 {
569 return None;
570 }
571
572 slot_access.push(SlotAccess {
573 slot_index: SlotIndex(slots.len()),
574 read_write: false,
575 });
576
577 let slot = SlotState::OriginalReadOnly(SharedAlignedBuffer::default());
578 slots.push((slot_key, slot));
579 let slot = &slots.last().expect("Just inserted; qed").1;
580 let SlotState::OriginalReadOnly(buffer) = slot else {
581 unreachable!("Just inserted; qed");
582 };
583
584 Some(buffer)
585 }
586 }
587
588 #[inline(always)]
590 fn use_ro_internal_read_only<'b>(
591 slot_key: SlotKey,
592 slots: &'b SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
593 slot_access: &SmallVec<[SlotAccess; INLINE_SIZE]>,
594 new_contracts: &[Address],
595 ) -> Option<&'b SharedAlignedBuffer> {
596 let maybe_slot_index = slots
597 .iter()
598 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
599 .map(SlotIndex);
600
601 if let Some(slot_index) = maybe_slot_index {
602 if let Some(read_write) = slot_access.iter().find_map(|slot_access| {
604 (slot_access.slot_index == slot_index).then_some(slot_access.read_write)
605 }) && read_write
606 {
607 return None;
608 }
609
610 let slot = &slots
611 .get(usize::from(slot_index))
612 .expect("Just found; qed")
613 .1;
614
615 match slot {
617 SlotState::Original(buffer)
618 | SlotState::OriginalReadOnly(buffer)
619 | SlotState::ModifiedReadOnly(buffer)
620 | SlotState::Modified(buffer) => Some(buffer),
621 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => None,
622 }
623 } else {
624 if !(slot_key.contract == Address::NULL
628 || new_contracts
629 .iter()
630 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
631 {
632 return None;
633 }
634
635 Some(SharedAlignedBuffer::empty_ref())
636 }
637 }
638
639 #[inline(always)]
648 pub fn use_rw(
649 &mut self,
650 slot_key: SlotKey,
651 capacity: u32,
652 ) -> Option<(SlotIndex, &mut OwnedAlignedBuffer)> {
653 let inner = self.inner_rw()?;
654 let slots = &mut inner.slots;
655 let slot_access = &mut inner.slot_access;
656 let new_contracts = &inner.new_contracts;
657
658 let result = Self::use_rw_internal(slot_key, capacity, slots, slot_access, new_contracts);
659
660 if result.is_none() {
661 debug!(?slot_key, "`use_rw` access violation");
662 }
663
664 result
665 }
666
667 #[inline(always)]
668 fn use_rw_internal<'b>(
669 slot_key: SlotKey,
670 capacity: u32,
671 slots: &'b mut SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
672 slot_access: &mut SmallVec<[SlotAccess; INLINE_SIZE]>,
673 new_contracts: &[Address],
674 ) -> Option<(SlotIndex, &'b mut OwnedAlignedBuffer)> {
675 let maybe_slot_index = slots
676 .iter()
677 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
678 .map(SlotIndex);
679
680 if let Some(slot_index) = maybe_slot_index {
681 if slot_access
683 .iter()
684 .any(|slot_access| slot_access.slot_index == slot_index)
685 {
686 return None;
687 }
688
689 slot_access.push(SlotAccess {
690 slot_index,
691 read_write: true,
692 });
693
694 let slot = &mut slots
695 .get_mut(usize::from(slot_index))
696 .expect("Just found; qed")
697 .1;
698
699 let buffer = match slot {
701 SlotState::OriginalReadOnly(_buffer) | SlotState::ModifiedReadOnly(_buffer) => {
702 return None;
703 }
704 SlotState::Original(buffer) => {
705 let mut new_buffer =
706 OwnedAlignedBuffer::with_capacity(capacity.max(buffer.len()));
707 new_buffer.copy_from_slice(buffer.as_slice());
708
709 *slot = SlotState::OriginalReadWrite {
710 buffer: new_buffer,
711 previous: buffer.clone(),
712 };
713 let SlotState::OriginalReadWrite { buffer, .. } = slot else {
714 unreachable!("Just inserted; qed");
715 };
716 buffer
717 }
718 SlotState::Modified(buffer) => {
719 let mut new_buffer =
720 OwnedAlignedBuffer::with_capacity(capacity.max(buffer.len()));
721 new_buffer.copy_from_slice(buffer.as_slice());
722
723 *slot = SlotState::ModifiedReadWrite {
724 buffer: new_buffer,
725 previous: buffer.clone(),
726 };
727 let SlotState::ModifiedReadWrite { buffer, .. } = slot else {
728 unreachable!("Just inserted; qed");
729 };
730 buffer
731 }
732 SlotState::OriginalReadWrite { buffer, .. }
733 | SlotState::ModifiedReadWrite { buffer, .. } => {
734 buffer.ensure_capacity(capacity);
735 buffer
736 }
737 };
738
739 Some((slot_index, buffer))
740 } else {
741 if !(slot_key.contract == Address::NULL
745 || new_contracts
746 .iter()
747 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
748 {
749 return None;
750 }
751
752 let slot_index = SlotIndex(slots.len());
753 slot_access.push(SlotAccess {
754 slot_index,
755 read_write: true,
756 });
757
758 let slot = SlotState::OriginalReadWrite {
759 buffer: OwnedAlignedBuffer::with_capacity(capacity),
760 previous: SharedAlignedBuffer::default(),
761 };
762 slots.push((slot_key, slot));
763 let slot = &mut slots.last_mut().expect("Just inserted; qed").1;
764 let SlotState::OriginalReadWrite { buffer, .. } = slot else {
765 unreachable!("Just inserted; qed");
766 };
767
768 Some((slot_index, buffer))
769 }
770 }
771
772 pub fn access_used_rw(&mut self, slot_index: SlotIndex) -> Option<&mut OwnedAlignedBuffer> {
780 let maybe_slot = self
781 .inner_rw()?
782 .slots
783 .get_mut(usize::from(slot_index))
784 .map(|(_slot_key, slot)| slot);
785
786 let Some(slot) = maybe_slot else {
787 debug!(?slot_index, "`access_used_rw` access violation (not found)");
788 return None;
789 };
790
791 match slot {
793 SlotState::Original(_buffer)
794 | SlotState::OriginalReadOnly(_buffer)
795 | SlotState::Modified(_buffer)
796 | SlotState::ModifiedReadOnly(_buffer) => {
797 debug!(?slot_index, "`access_used_rw` access violation (read only)");
798 None
799 }
800 SlotState::OriginalReadWrite { buffer, .. }
801 | SlotState::ModifiedReadWrite { buffer, .. } => Some(buffer),
802 }
803 }
804
805 #[cold]
807 pub fn reset(&mut self) {
808 let (inner, parent_slot_access_len) = match &mut self.0 {
809 NestedSlotsInner::ReadWrite {
810 inner,
811 parent_slot_access_len,
812 original_parent: _,
813 } => (&mut **inner, parent_slot_access_len),
814 NestedSlotsInner::ReadOnly { .. } => {
815 return;
817 }
818 };
819
820 let slots = &mut inner.slots;
821 let slot_access = &mut inner.slot_access;
822
823 for slot_access in slot_access.drain(*parent_slot_access_len..) {
825 let slot = &mut slots
826 .get_mut(usize::from(slot_access.slot_index))
827 .expect("Accessed slot exists; qed")
828 .1;
829 replace_with_or_abort(slot, |slot| match slot {
830 SlotState::Original(_buffer) => {
831 unreachable!("Slot can't be in `Original` state after being accessed; qed");
832 }
833 SlotState::OriginalReadOnly(buffer) => SlotState::Original(buffer),
834 SlotState::Modified(buffer) => SlotState::Modified(buffer),
835 SlotState::ModifiedReadOnly(buffer) => SlotState::Modified(buffer),
836 SlotState::OriginalReadWrite { previous, .. } => SlotState::Original(previous),
837 SlotState::ModifiedReadWrite { previous, .. } => SlotState::Modified(previous),
838 });
839 }
840
841 *parent_slot_access_len = 0;
842 }
843}