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