1#![no_std]
2
3extern crate alloc;
4
5use ab_aligned_buffer::{OwnedAlignedBuffer, SharedAlignedBuffer};
6use ab_core_primitives::address::Address;
7use alloc::boxed::Box;
8use smallvec::SmallVec;
9use tracing::debug;
10
11const INLINE_SIZE: usize = 8;
16const NEW_CONTRACTS_INLINE: usize = 2;
18
19#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
21pub struct SlotKey {
22 pub owner: Address,
24 pub contract: Address,
26}
27
28#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
29pub struct SlotIndex(usize);
30
31impl From<SlotIndex> for usize {
32 #[inline(always)]
33 fn from(value: SlotIndex) -> Self {
34 value.0
35 }
36}
37
38#[derive(Debug, Clone)]
39pub enum Slot {
40 ReadOnly {
41 key: SlotKey,
42 buffer: SharedAlignedBuffer,
43 },
44 ReadWrite {
45 key: SlotKey,
46 buffer: SharedAlignedBuffer,
47 },
48}
49
50impl Slot {
51 fn is_null_contract(&self) -> bool {
52 let slot_key = match self {
53 Slot::ReadOnly { key, .. } => key,
54 Slot::ReadWrite { key, .. } => key,
55 };
56
57 slot_key.contract == Address::NULL
58 }
59}
60
61#[derive(Debug, Clone)]
62enum SlotState {
63 Original(SharedAlignedBuffer),
65 OriginalReadOnly(SharedAlignedBuffer),
67 Modified(SharedAlignedBuffer),
69 ModifiedReadOnly(SharedAlignedBuffer),
71 OriginalReadWrite {
73 buffer: OwnedAlignedBuffer,
74 previous: SharedAlignedBuffer,
76 },
77 ModifiedReadWrite {
79 buffer: OwnedAlignedBuffer,
80 previous: SharedAlignedBuffer,
82 },
83}
84
85#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
86struct SlotAccess {
87 slot_index: SlotIndex,
88 read_write: bool,
90}
91
92#[derive(Debug, Clone)]
93struct Inner {
94 slots: SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
95 slot_access: SmallVec<[SlotAccess; INLINE_SIZE]>,
96 new_contracts: SmallVec<[Address; NEW_CONTRACTS_INLINE]>,
102}
103
104#[derive(Debug, Clone)]
106pub struct Slots(Box<Inner>);
107
108impl Slots {
109 #[inline(always)]
116 pub fn new<I>(slots: I) -> Self
117 where
118 I: IntoIterator<Item = Slot>,
119 {
120 let slots = slots
121 .into_iter()
122 .filter_map(|slot| {
123 if slot.is_null_contract() {
127 return None;
128 }
129
130 Some(match slot {
131 Slot::ReadOnly { key, buffer } => {
132 (key, SlotState::OriginalReadOnly(buffer))
134 }
135 Slot::ReadWrite { key, buffer } => (key, SlotState::Original(buffer)),
136 })
137 })
138 .collect();
139
140 let inner = Inner {
141 slots,
142 slot_access: SmallVec::new(),
143 new_contracts: SmallVec::new(),
144 };
145
146 Self(Box::new(inner))
147 }
148
149 #[inline(always)]
154 pub fn new_nested_rw(&mut self) -> NestedSlots<'_> {
155 NestedSlots(NestedSlotsInner::ReadWrite {
156 inner: &mut self.0,
157 parent_slot_access_len: 0,
158 original_parent: true,
159 })
160 }
161
162 #[inline(always)]
164 pub fn new_nested_ro(&self) -> NestedSlots<'_> {
165 NestedSlots(NestedSlotsInner::ReadOnly { inner: &self.0 })
166 }
167
168 #[must_use]
176 #[inline(always)]
177 pub fn add_new_contract(&mut self, owner: Address) -> bool {
178 let new_contracts = &mut self.0.new_contracts;
179
180 if new_contracts.contains(&owner) {
181 debug!(?owner, "Not adding new contract duplicate");
182 return false;
183 }
184
185 new_contracts.push(owner);
186 true
187 }
188
189 #[inline]
191 pub fn iter(&self) -> impl ExactSizeIterator<Item = (&SlotKey, &SharedAlignedBuffer)> + '_ {
192 self.0.slots.iter().map(|(slot_key, slot)| match slot {
193 SlotState::Original(buffer) => (slot_key, buffer),
194 SlotState::OriginalReadOnly(_) => {
195 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
196 }
197 SlotState::Modified(buffer) => (slot_key, buffer),
198 SlotState::ModifiedReadOnly(_) => {
199 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
200 }
201 SlotState::OriginalReadWrite { .. } => {
202 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
203 }
204 SlotState::ModifiedReadWrite { .. } => {
205 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
206 }
207 })
208 }
209
210 #[inline]
212 pub fn iter_modified(&self) -> impl Iterator<Item = (&SlotKey, &SharedAlignedBuffer)> + '_ {
213 self.0
214 .slots
215 .iter()
216 .filter_map(|(slot_key, slot)| match slot {
217 SlotState::Original(_) => None,
218 SlotState::OriginalReadOnly(_) => unreachable!(
219 "Only original and modified slots can exist at the `Slots` level; qed"
220 ),
221 SlotState::Modified(buffer) => Some((slot_key, buffer)),
222 SlotState::ModifiedReadOnly(_) => unreachable!(
223 "Only original and modified slots can exist at the `Slots` level; qed"
224 ),
225 SlotState::OriginalReadWrite { .. } => unreachable!(
226 "Only original and modified slots can exist at the `Slots` level; qed"
227 ),
228 SlotState::ModifiedReadWrite { .. } => unreachable!(
229 "Only original and modified slots can exist at the `Slots` level; qed"
230 ),
231 })
232 }
233
234 #[inline]
236 pub fn into_slots(self) -> impl ExactSizeIterator<Item = (SlotKey, SharedAlignedBuffer)> {
237 self.0.slots.into_iter().map(|(slot_key, slot)| match slot {
238 SlotState::Original(buffer) => (slot_key, buffer),
239 SlotState::OriginalReadOnly(_) => {
240 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
241 }
242 SlotState::Modified(buffer) => (slot_key, buffer),
243 SlotState::ModifiedReadOnly(_) => {
244 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
245 }
246 SlotState::OriginalReadWrite { .. } => {
247 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
248 }
249 SlotState::ModifiedReadWrite { .. } => {
250 unreachable!("Only original and modified slots can exist at the `Slots` level; qed")
251 }
252 })
253 }
254}
255
256#[derive(Debug)]
258enum NestedSlotsInner<'a> {
259 ReadWrite {
261 inner: &'a mut Inner,
262 parent_slot_access_len: usize,
263 original_parent: bool,
264 },
265 ReadOnly { inner: &'a Inner },
267}
268
269#[derive(Debug)]
270pub struct NestedSlots<'a>(NestedSlotsInner<'a>);
271
272impl<'a> Drop for NestedSlots<'a> {
273 #[inline(always)]
274 fn drop(&mut self) {
275 let (inner, parent_slot_access_len, original_parent) = match &mut self.0 {
276 NestedSlotsInner::ReadWrite {
277 inner,
278 parent_slot_access_len,
279 original_parent,
280 } => (&mut **inner, *parent_slot_access_len, *original_parent),
281 NestedSlotsInner::ReadOnly { .. } => {
282 return;
284 }
285 };
286
287 let slots = &mut inner.slots;
288 let slot_access = &mut inner.slot_access;
289
290 for slot_access in slot_access.drain(parent_slot_access_len..) {
292 let slot = &mut slots
293 .get_mut(usize::from(slot_access.slot_index))
294 .expect("Accessed slot exists; qed")
295 .1;
296
297 take_mut::take(slot, |slot| match slot {
298 SlotState::Original(_buffer) => {
299 unreachable!("Slot can't be in Original state after being accessed")
300 }
301 SlotState::OriginalReadOnly(buffer) => SlotState::Original(buffer),
302 SlotState::Modified(buffer) => SlotState::Modified(buffer),
303 SlotState::ModifiedReadOnly(buffer) => SlotState::Modified(buffer),
304 SlotState::OriginalReadWrite { buffer, .. }
305 | SlotState::ModifiedReadWrite { buffer, .. } => {
306 SlotState::Modified(buffer.into_shared())
307 }
308 })
309 }
310
311 if original_parent {
312 inner
316 .slots
317 .retain(|(slot_key, _slot)| slot_key.contract != Address::NULL);
318 }
319 }
320}
321
322impl<'a> NestedSlots<'a> {
323 #[inline(always)]
324 fn inner_ro(&self) -> &Inner {
325 match &self.0 {
326 NestedSlotsInner::ReadWrite { inner, .. } => inner,
327 NestedSlotsInner::ReadOnly { inner } => inner,
328 }
329 }
330
331 #[inline(always)]
332 fn inner_rw(&mut self) -> Option<&mut Inner> {
333 match &mut self.0 {
334 NestedSlotsInner::ReadWrite { inner, .. } => Some(inner),
335 NestedSlotsInner::ReadOnly { .. } => None,
336 }
337 }
338
339 #[inline(always)]
346 pub fn new_nested_rw<'b>(&'b mut self) -> Option<NestedSlots<'b>>
347 where
348 'a: 'b,
349 {
350 let inner = match &mut self.0 {
351 NestedSlotsInner::ReadWrite { inner, .. } => &mut **inner,
352 NestedSlotsInner::ReadOnly { .. } => {
353 return None;
354 }
355 };
356
357 let parent_slot_access_len = inner.slot_access.len();
358
359 Some(NestedSlots(NestedSlotsInner::ReadWrite {
360 inner,
361 parent_slot_access_len,
362 original_parent: false,
363 }))
364 }
365
366 #[inline(always)]
368 pub fn new_nested_ro<'b>(&'b self) -> NestedSlots<'b>
369 where
370 'a: 'b,
371 {
372 let inner = match &self.0 {
373 NestedSlotsInner::ReadWrite { inner, .. } => &**inner,
374 NestedSlotsInner::ReadOnly { inner } => &**inner,
375 };
376
377 NestedSlots(NestedSlotsInner::ReadOnly { inner })
378 }
379
380 #[must_use]
388 #[inline(always)]
389 pub fn add_new_contract(&mut self, owner: Address) -> bool {
390 let Some(inner) = self.inner_rw() else {
391 debug!(?owner, "`add_new_contract` access violation");
392 return false;
393 };
394
395 let new_contracts = &mut inner.new_contracts;
396
397 if new_contracts.contains(&owner) {
398 debug!(?owner, "Not adding new contract duplicate");
399 return false;
400 }
401
402 new_contracts.push(owner);
403 true
404 }
405
406 #[inline(always)]
413 pub fn get_code(&self, owner: Address) -> Option<SharedAlignedBuffer> {
414 let result = self.get_code_internal(owner);
415
416 if result.is_none() {
417 debug!(?owner, "`get_code` access violation");
418 }
419
420 result
421 }
422
423 #[inline(always)]
424 fn get_code_internal(&self, owner: Address) -> Option<SharedAlignedBuffer> {
425 let inner = self.inner_ro();
426 let slots = &inner.slots;
427 let slot_access = &inner.slot_access;
428
429 let contract = Address::SYSTEM_CODE;
430
431 let slot_index = slots.iter().position(|(slot_key, _slot)| {
432 slot_key.owner == owner && slot_key.contract == contract
433 })?;
434 let slot_index = SlotIndex(slot_index);
435
436 if slot_access
438 .iter()
439 .any(|slot_access| slot_access.slot_index == slot_index && slot_access.read_write)
440 {
441 return None;
442 }
443
444 let buffer = match &slots
445 .get(usize::from(slot_index))
446 .expect("Just found; qed")
447 .1
448 {
449 SlotState::Original(buffer)
450 | SlotState::OriginalReadOnly(buffer)
451 | SlotState::Modified(buffer)
452 | SlotState::ModifiedReadOnly(buffer) => buffer,
453 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => {
454 return None;
455 }
456 };
457
458 Some(buffer.clone())
459 }
460
461 #[inline(always)]
465 pub fn use_ro(&mut self, slot_key: SlotKey) -> Option<&SharedAlignedBuffer> {
466 let inner_rw = match &mut self.0 {
467 NestedSlotsInner::ReadWrite { inner, .. } => &mut **inner,
468 NestedSlotsInner::ReadOnly { inner } => {
469 let result = Self::use_ro_internal_read_only(
471 slot_key,
472 &inner.slots,
473 &inner.slot_access,
474 &inner.new_contracts,
475 );
476
477 if result.is_none() {
478 debug!(?slot_key, "`use_ro` access violation");
479 }
480
481 return result;
482 }
483 };
484
485 let result = Self::use_ro_internal(
486 slot_key,
487 &mut inner_rw.slots,
488 &mut inner_rw.slot_access,
489 &inner_rw.new_contracts,
490 );
491
492 if result.is_none() {
493 debug!(?slot_key, "`use_ro` access violation");
494 }
495
496 result
497 }
498
499 #[inline(always)]
500 fn use_ro_internal<'b>(
501 slot_key: SlotKey,
502 slots: &'b mut SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
503 slot_access: &mut SmallVec<[SlotAccess; INLINE_SIZE]>,
504 new_contracts: &[Address],
505 ) -> Option<&'b SharedAlignedBuffer> {
506 let maybe_slot_index = slots
507 .iter()
508 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
509 .map(SlotIndex);
510
511 if let Some(slot_index) = maybe_slot_index {
512 if let Some(read_write) = slot_access.iter().find_map(|slot_access| {
514 (slot_access.slot_index == slot_index).then_some(slot_access.read_write)
515 }) {
516 if read_write {
517 return None;
518 }
519 } else {
520 slot_access.push(SlotAccess {
521 slot_index,
522 read_write: false,
523 });
524 }
525
526 let slot = &mut slots
527 .get_mut(usize::from(slot_index))
528 .expect("Just found; qed")
529 .1;
530
531 match slot {
533 SlotState::Original(buffer) => {
534 let buffer = buffer.clone();
535 *slot = SlotState::OriginalReadOnly(buffer);
536 let SlotState::OriginalReadOnly(buffer) = slot else {
537 unreachable!("Just inserted; qed");
538 };
539 Some(buffer)
540 }
541 SlotState::OriginalReadOnly(buffer) | SlotState::ModifiedReadOnly(buffer) => {
542 Some(buffer)
543 }
544 SlotState::Modified(buffer) => {
545 let buffer = buffer.clone();
546 *slot = SlotState::ModifiedReadOnly(buffer);
547 let SlotState::ModifiedReadOnly(buffer) = slot else {
548 unreachable!("Just inserted; qed");
549 };
550 Some(buffer)
551 }
552 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => None,
553 }
554 } else {
555 if !(slot_key.contract == Address::NULL
559 || new_contracts
560 .iter()
561 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
562 {
563 return None;
564 }
565
566 slot_access.push(SlotAccess {
567 slot_index: SlotIndex(slots.len()),
568 read_write: false,
569 });
570
571 let slot = SlotState::OriginalReadOnly(SharedAlignedBuffer::default());
572 slots.push((slot_key, slot));
573 let slot = &slots.last().expect("Just inserted; qed").1;
574 let SlotState::OriginalReadOnly(buffer) = slot else {
575 unreachable!("Just inserted; qed");
576 };
577
578 Some(buffer)
579 }
580 }
581
582 #[inline(always)]
584 fn use_ro_internal_read_only<'b>(
585 slot_key: SlotKey,
586 slots: &'b SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
587 slot_access: &SmallVec<[SlotAccess; INLINE_SIZE]>,
588 new_contracts: &[Address],
589 ) -> Option<&'b SharedAlignedBuffer> {
590 let maybe_slot_index = slots
591 .iter()
592 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
593 .map(SlotIndex);
594
595 if let Some(slot_index) = maybe_slot_index {
596 if let Some(read_write) = slot_access.iter().find_map(|slot_access| {
598 (slot_access.slot_index == slot_index).then_some(slot_access.read_write)
599 }) && read_write
600 {
601 return None;
602 }
603
604 let slot = &slots
605 .get(usize::from(slot_index))
606 .expect("Just found; qed")
607 .1;
608
609 match slot {
611 SlotState::Original(buffer)
612 | SlotState::OriginalReadOnly(buffer)
613 | SlotState::ModifiedReadOnly(buffer)
614 | SlotState::Modified(buffer) => Some(buffer),
615 SlotState::OriginalReadWrite { .. } | SlotState::ModifiedReadWrite { .. } => None,
616 }
617 } else {
618 if !(slot_key.contract == Address::NULL
622 || new_contracts
623 .iter()
624 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
625 {
626 return None;
627 }
628
629 Some(SharedAlignedBuffer::empty_ref())
630 }
631 }
632
633 #[inline(always)]
641 pub fn use_rw(
642 &mut self,
643 slot_key: SlotKey,
644 capacity: u32,
645 ) -> Option<(SlotIndex, &mut OwnedAlignedBuffer)> {
646 let inner = self.inner_rw()?;
647 let slots = &mut inner.slots;
648 let slot_access = &mut inner.slot_access;
649 let new_contracts = &inner.new_contracts;
650
651 let result = Self::use_rw_internal(slot_key, capacity, slots, slot_access, new_contracts);
652
653 if result.is_none() {
654 debug!(?slot_key, "`use_rw` access violation");
655 }
656
657 result
658 }
659
660 #[inline(always)]
661 fn use_rw_internal<'b>(
662 slot_key: SlotKey,
663 capacity: u32,
664 slots: &'b mut SmallVec<[(SlotKey, SlotState); INLINE_SIZE]>,
665 slot_access: &mut SmallVec<[SlotAccess; INLINE_SIZE]>,
666 new_contracts: &[Address],
667 ) -> Option<(SlotIndex, &'b mut OwnedAlignedBuffer)> {
668 let maybe_slot_index = slots
669 .iter()
670 .position(|(slot_key_candidate, _slot)| slot_key_candidate == &slot_key)
671 .map(SlotIndex);
672
673 if let Some(slot_index) = maybe_slot_index {
674 if slot_access
676 .iter()
677 .any(|slot_access| slot_access.slot_index == slot_index)
678 {
679 return None;
680 }
681
682 slot_access.push(SlotAccess {
683 slot_index,
684 read_write: true,
685 });
686
687 let slot = &mut slots
688 .get_mut(usize::from(slot_index))
689 .expect("Just found; qed")
690 .1;
691
692 let buffer = match slot {
694 SlotState::OriginalReadOnly(_buffer) | SlotState::ModifiedReadOnly(_buffer) => {
695 return None;
696 }
697 SlotState::Original(buffer) => {
698 let mut new_buffer =
699 OwnedAlignedBuffer::with_capacity(capacity.max(buffer.len()));
700 new_buffer.copy_from_slice(buffer.as_slice());
701
702 *slot = SlotState::OriginalReadWrite {
703 buffer: new_buffer,
704 previous: buffer.clone(),
705 };
706 let SlotState::OriginalReadWrite { buffer, .. } = slot else {
707 unreachable!("Just inserted; qed");
708 };
709 buffer
710 }
711 SlotState::Modified(buffer) => {
712 let mut new_buffer =
713 OwnedAlignedBuffer::with_capacity(capacity.max(buffer.len()));
714 new_buffer.copy_from_slice(buffer.as_slice());
715
716 *slot = SlotState::ModifiedReadWrite {
717 buffer: new_buffer,
718 previous: buffer.clone(),
719 };
720 let SlotState::ModifiedReadWrite { buffer, .. } = slot else {
721 unreachable!("Just inserted; qed");
722 };
723 buffer
724 }
725 SlotState::OriginalReadWrite { buffer, .. }
726 | SlotState::ModifiedReadWrite { buffer, .. } => {
727 buffer.ensure_capacity(capacity);
728 buffer
729 }
730 };
731
732 Some((slot_index, buffer))
733 } else {
734 if !(slot_key.contract == Address::NULL
738 || new_contracts
739 .iter()
740 .any(|candidate| candidate == slot_key.owner || candidate == slot_key.contract))
741 {
742 return None;
743 }
744
745 let slot_index = SlotIndex(slots.len());
746 slot_access.push(SlotAccess {
747 slot_index,
748 read_write: true,
749 });
750
751 let slot = SlotState::OriginalReadWrite {
752 buffer: OwnedAlignedBuffer::with_capacity(capacity),
753 previous: SharedAlignedBuffer::default(),
754 };
755 slots.push((slot_key, slot));
756 let slot = &mut slots.last_mut().expect("Just inserted; qed").1;
757 let SlotState::OriginalReadWrite { buffer, .. } = slot else {
758 unreachable!("Just inserted; qed");
759 };
760
761 Some((slot_index, buffer))
762 }
763 }
764
765 pub fn access_used_rw(&mut self, slot_index: SlotIndex) -> Option<&mut OwnedAlignedBuffer> {
773 let maybe_slot = self
774 .inner_rw()?
775 .slots
776 .get_mut(usize::from(slot_index))
777 .map(|(_slot_key, slot)| slot);
778
779 let Some(slot) = maybe_slot else {
780 debug!(?slot_index, "`access_used_rw` access violation (not found)");
781 return None;
782 };
783
784 match slot {
786 SlotState::Original(_buffer)
787 | SlotState::OriginalReadOnly(_buffer)
788 | SlotState::Modified(_buffer)
789 | SlotState::ModifiedReadOnly(_buffer) => {
790 debug!(?slot_index, "`access_used_rw` access violation (read only)");
791 None
792 }
793 SlotState::OriginalReadWrite { buffer, .. }
794 | SlotState::ModifiedReadWrite { buffer, .. } => Some(buffer),
795 }
796 }
797
798 #[cold]
800 pub fn reset(&mut self) {
801 let (inner, parent_slot_access_len) = match &mut self.0 {
802 NestedSlotsInner::ReadWrite {
803 inner,
804 parent_slot_access_len,
805 original_parent: _,
806 } => (&mut **inner, parent_slot_access_len),
807 NestedSlotsInner::ReadOnly { .. } => {
808 return;
810 }
811 };
812
813 let slots = &mut inner.slots;
814 let slot_access = &mut inner.slot_access;
815
816 for slot_access in slot_access.drain(*parent_slot_access_len..) {
818 let slot = &mut slots
819 .get_mut(usize::from(slot_access.slot_index))
820 .expect("Accessed slot exists; qed")
821 .1;
822 take_mut::take(slot, |slot| match slot {
823 SlotState::Original(_buffer) => {
824 unreachable!("Slot can't be in Original state after being accessed")
825 }
826 SlotState::OriginalReadOnly(buffer) => SlotState::Original(buffer),
827 SlotState::Modified(buffer) => SlotState::Modified(buffer),
828 SlotState::ModifiedReadOnly(buffer) => SlotState::Modified(buffer),
829 SlotState::OriginalReadWrite { previous, .. } => SlotState::Original(previous),
830 SlotState::ModifiedReadWrite { previous, .. } => SlotState::Modified(previous),
831 });
832 }
833
834 *parent_slot_access_len = 0;
835 }
836}