1#![no_std]
2
3extern crate alloc;
4
5use ab_aligned_buffer::{OwnedAlignedBuffer, SharedAlignedBuffer};
6use ab_contracts_common::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 }) {
600 if read_write {
601 return None;
602 }
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 take_mut::take(slot, |slot| match slot {
824 SlotState::Original(_buffer) => {
825 unreachable!("Slot can't be in Original state after being accessed")
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}