1use alloy::primitives::U256;
7use proc_macro2::TokenStream;
8use quote::{format_ident, quote};
9use syn::{Ident, Type};
10
11use crate::{FieldInfo, FieldKind};
12
13pub(crate) struct PackingConstants(String);
15
16impl PackingConstants {
17 pub(crate) fn new(name: &Ident) -> Self {
19 Self(const_name(name))
20 }
21
22 pub(crate) fn slot(&self) -> Ident {
24 format_ident!("{}", &self.0)
25 }
26
27 pub(crate) fn location(&self) -> Ident {
29 let span = proc_macro2::Span::call_site();
30 Ident::new(&format!("{}_LOC", self.0), span)
31 }
32
33 pub(crate) fn offset(&self) -> Ident {
35 let span = proc_macro2::Span::call_site();
36 Ident::new(&format!("{}_OFFSET", self.0), span)
37 }
38
39 pub(crate) fn into_tuple(self) -> (Ident, Ident) {
41 (self.slot(), self.offset())
42 }
43}
44
45pub(crate) fn const_name(name: &Ident) -> String {
47 name.to_string().to_uppercase()
48}
49
50#[derive(Debug, Clone)]
52pub(crate) enum SlotAssignment {
53 Manual(U256),
55 Auto {
57 base_slot: U256,
59 },
60}
61
62impl SlotAssignment {
63 pub(crate) fn ref_slot(&self) -> &U256 {
64 match self {
65 Self::Manual(slot) => slot,
66 Self::Auto { base_slot } => base_slot,
67 }
68 }
69}
70
71#[derive(Debug)]
73pub(crate) struct LayoutField<'a> {
74 pub name: &'a Ident,
76 pub ty: &'a Type,
78 pub kind: FieldKind<'a>,
80 pub assigned_slot: SlotAssignment,
82}
83
84pub(crate) fn allocate_slots(fields: &[FieldInfo]) -> syn::Result<Vec<LayoutField<'_>>> {
94 let mut result = Vec::with_capacity(fields.len());
95 let mut current_base_slot = U256::ZERO;
96
97 for field in fields.iter() {
98 let kind = classify_field_type(&field.ty)?;
99
100 let assigned_slot = if let Some(explicit) = field.slot {
102 SlotAssignment::Manual(explicit)
103 } else if let Some(new_base) = field.base_slot {
104 current_base_slot = new_base;
106 SlotAssignment::Auto {
107 base_slot: new_base,
108 }
109 } else {
110 SlotAssignment::Auto {
111 base_slot: current_base_slot,
112 }
113 };
114
115 result.push(LayoutField {
116 name: &field.name,
117 ty: &field.ty,
118 kind,
119 assigned_slot,
120 });
121 }
122
123 Ok(result)
124}
125
126pub(crate) fn gen_constants_from_ir(fields: &[LayoutField<'_>], gen_location: bool) -> TokenStream {
132 let mut constants = TokenStream::new();
133 let mut current_base_slot: Option<&LayoutField<'_>> = None;
134
135 for field in fields {
136 let ty = field.ty;
137 let consts = PackingConstants::new(field.name);
138 let (loc_const, (slot_const, offset_const)) = (consts.location(), consts.into_tuple());
139 let slots_to_end = quote! {
140 ::alloy::primitives::U256::from_limbs([<#ty as crate::storage::StorableType>::SLOTS as u64, 0, 0, 0])
141 .saturating_sub(::alloy::primitives::U256::ONE)
142 };
143
144 let bytes_expr = quote! { <#ty as crate::storage::StorableType>::BYTES };
146
147 let (slot_expr, offset_expr) = match &field.assigned_slot {
149 SlotAssignment::Manual(manual_slot) => {
151 let hex_value = format!("{manual_slot}_U256");
152 let slot_lit = syn::LitInt::new(&hex_value, proc_macro2::Span::call_site());
153 let slot_expr = quote! {
156 ::alloy::primitives::uint!(#slot_lit)
157 .checked_add(#slots_to_end).expect("slot overflow")
158 .saturating_sub(#slots_to_end)
159 };
160 (slot_expr, quote! { 0 })
161 }
162 SlotAssignment::Auto { base_slot, .. } => {
164 let output = if let Some(current_base) = current_base_slot
165 && current_base.assigned_slot.ref_slot() == field.assigned_slot.ref_slot()
166 {
167 let (prev_slot, prev_offset) =
169 PackingConstants::new(current_base.name).into_tuple();
170 gen_slot_packing_logic(
171 current_base.ty,
172 field.ty,
173 quote! { #prev_slot },
174 quote! { #prev_offset },
175 )
176 } else {
177 let limbs = *base_slot.as_limbs();
179
180 let slot_expr = quote! {
183 ::alloy::primitives::U256::from_limbs([#(#limbs),*])
184 .checked_add(#slots_to_end).expect("slot overflow")
185 .saturating_sub(#slots_to_end)
186 };
187 (slot_expr, quote! { 0 })
188 };
189 current_base_slot = Some(field);
191 output
192 }
193 };
194
195 constants.extend(quote! {
197 pub const #slot_const: ::alloy::primitives::U256 = #slot_expr;
198 pub const #offset_const: usize = #offset_expr;
199 });
200
201 if gen_location {
205 constants.extend(quote! {
206 pub const #loc_const: crate::storage::packing::FieldLocation =
207 crate::storage::packing::FieldLocation::new(#slot_const.as_limbs()[0] as usize, #offset_const, #bytes_expr);
208 });
209 }
210
211 #[cfg(debug_assertions)]
213 {
214 let bytes_const = format_ident!("{slot_const}_BYTES");
215 constants.extend(quote! { pub const #bytes_const: usize = #bytes_expr; });
216 }
217 }
218
219 constants
220}
221
222pub(crate) fn classify_field_type(ty: &Type) -> syn::Result<FieldKind<'_>> {
228 use crate::utils::extract_mapping_types;
229
230 if let Some((key_ty, value_ty)) = extract_mapping_types(ty) {
232 return Ok(FieldKind::Mapping {
233 key: key_ty,
234 value: value_ty,
235 });
236 }
237
238 Ok(FieldKind::Direct(ty))
240}
241
242pub(crate) fn get_neighbor_slot_refs<T, F>(
249 idx: usize,
250 fields: &[T],
251 packing: &Ident,
252 get_name: F,
253 use_full_slot: bool,
254) -> (Option<TokenStream>, Option<TokenStream>)
255where
256 F: Fn(&T) -> &Ident,
257{
258 let prev_slot_ref = if idx > 0 {
259 let prev_name = get_name(&fields[idx - 1]);
260 if use_full_slot {
261 let prev_slot = PackingConstants::new(prev_name).slot();
262 Some(quote! { #packing::#prev_slot })
263 } else {
264 let prev_loc = PackingConstants::new(prev_name).location();
265 Some(quote! { #packing::#prev_loc.offset_slots })
266 }
267 } else {
268 None
269 };
270
271 let next_slot_ref = if idx + 1 < fields.len() {
272 let next_name = get_name(&fields[idx + 1]);
273 if use_full_slot {
274 let next_slot = PackingConstants::new(next_name).slot();
275 Some(quote! { #packing::#next_slot })
276 } else {
277 let next_loc = PackingConstants::new(next_name).location();
278 Some(quote! { #packing::#next_loc.offset_slots })
279 }
280 } else {
281 None
282 };
283
284 (prev_slot_ref, next_slot_ref)
285}
286
287pub(crate) fn gen_slot_packing_logic(
294 prev_ty: &Type,
295 curr_ty: &Type,
296 prev_slot_expr: TokenStream,
297 prev_offset_expr: TokenStream,
298) -> (TokenStream, TokenStream) {
299 let prev_layout_slots = quote! {
301 ::alloy::primitives::U256::from_limbs([<#prev_ty as crate::storage::StorableType>::SLOTS as u64, 0, 0, 0])
302 };
303 let curr_slots_to_end = quote! {
304 ::alloy::primitives::U256::from_limbs([<#curr_ty as crate::storage::StorableType>::SLOTS as u64, 0, 0, 0])
305 .saturating_sub(::alloy::primitives::U256::ONE)
306 };
307
308 let can_pack_expr = quote! {
310 #prev_offset_expr
311 + <#prev_ty as crate::storage::StorableType>::BYTES
312 + <#curr_ty as crate::storage::StorableType>::BYTES <= 32
313 };
314
315 let slot_expr = quote! {{
316 if #can_pack_expr {
317 #prev_slot_expr
318 } else {
319 #prev_slot_expr
322 .checked_add(#prev_layout_slots).expect("slot overflow")
323 .checked_add(#curr_slots_to_end).expect("slot overflow")
324 .saturating_sub(#curr_slots_to_end)
325 }
326 }};
327
328 let offset_expr = quote! {{
329 if #can_pack_expr { #prev_offset_expr + <#prev_ty as crate::storage::StorableType>::BYTES } else { 0 }
330 }};
331
332 (slot_expr, offset_expr)
333}
334
335pub(crate) fn gen_layout_ctx_expr(
342 ty: &Type,
343 is_manual_slot: bool,
344 slot_const_ref: TokenStream,
345 offset_const_ref: TokenStream,
346 prev_slot_const_ref: Option<TokenStream>,
347 next_slot_const_ref: Option<TokenStream>,
348) -> TokenStream {
349 if !is_manual_slot && (prev_slot_const_ref.is_some() || next_slot_const_ref.is_some()) {
350 let prev_check = prev_slot_const_ref.map(|prev| quote! { #slot_const_ref == #prev });
352 let next_check = next_slot_const_ref.map(|next| quote! { #slot_const_ref == #next });
353
354 let shares_slot_check = match (prev_check, next_check) {
355 (Some(prev), Some(next)) => quote! { (#prev || #next) },
356 (Some(prev), None) => prev,
357 (None, Some(next)) => next,
358 (None, None) => unreachable!(),
359 };
360
361 quote! {
362 {
363 if #shares_slot_check && <#ty as crate::storage::StorableType>::IS_PACKABLE {
364 crate::storage::LayoutCtx::packed(#offset_const_ref)
365 } else {
366 crate::storage::LayoutCtx::FULL
367 }
368 }
369 }
370 } else {
371 quote! { crate::storage::LayoutCtx::FULL }
372 }
373}
374
375pub(crate) fn gen_collision_check_fn(
381 idx: usize,
382 field: &LayoutField<'_>,
383 all_fields: &[LayoutField<'_>],
384) -> (Ident, TokenStream) {
385 fn gen_slot_count_expr(ty: &Type) -> TokenStream {
386 quote! { ::alloy::primitives::U256::from_limbs([<#ty as crate::storage::StorableType>::SLOTS as u64, 0, 0, 0]) }
387 }
388
389 let check_fn_name = format_ident!("__check_collision_{}", field.name);
390 let consts = PackingConstants::new(field.name);
391 let (slot_const, offset_const) = consts.into_tuple();
392 let (field_name, field_ty) = (field.name, field.ty);
393
394 let mut checks = TokenStream::new();
395
396 for (other_idx, other_field) in all_fields.iter().enumerate() {
398 if other_idx == idx {
399 continue;
400 }
401
402 let other_consts = PackingConstants::new(other_field.name);
403 let (other_slot_const, other_offset_const) = other_consts.into_tuple();
404 let other_name = other_field.name;
405 let other_ty = other_field.ty;
406
407 let current_count_expr = gen_slot_count_expr(field.ty);
409 let other_count_expr = gen_slot_count_expr(other_field.ty);
410
411 checks.extend(quote! {
414 {
415 let slot = #slot_const;
416 let slot_end = slot.checked_add(#current_count_expr).expect("slot range overflow");
417 let other_slot = #other_slot_const;
418 let other_slot_end = other_slot.checked_add(#other_count_expr).expect("slot range overflow");
419
420 let no_overlap = if slot == other_slot {
424 let byte_end = #offset_const + <#field_ty as crate::storage::StorableType>::BYTES;
425 let other_byte_end = #other_offset_const + <#other_ty as crate::storage::StorableType>::BYTES;
426 byte_end <= #other_offset_const || other_byte_end <= #offset_const
427 } else {
428 slot_end.le(&other_slot) || other_slot_end.le(&slot)
429 };
430
431 debug_assert!(
432 no_overlap,
433 "Storage slot collision: field `{}` (slot {:?}, offset {}) overlaps with field `{}` (slot {:?}, offset {})",
434 stringify!(#field_name),
435 slot,
436 #offset_const,
437 stringify!(#other_name),
438 other_slot,
439 #other_offset_const
440 );
441 }
442 });
443 }
444
445 let check_fn = quote! {
446 #[cfg(debug_assertions)]
447 #[inline(always)]
448 #[allow(non_snake_case)]
449 fn #check_fn_name() {
450 #checks
451 }
452 };
453
454 (check_fn_name, check_fn)
455}
456
457#[cfg(test)]
458mod tests {
459 use super::*;
460 use crate::utils::extract_mapping_types;
461 use syn::parse_quote;
462
463 fn field(name: &str, ty: Type, slot: Option<u64>, base_slot: Option<u64>) -> FieldInfo {
464 FieldInfo {
465 name: Ident::new(name, proc_macro2::Span::call_site()),
466 ty,
467 slot: slot.map(U256::from),
468 base_slot: base_slot.map(U256::from),
469 }
470 }
471
472 #[test]
473 fn allocate_slots_edge_cases() {
474 assert!(allocate_slots(&[]).unwrap().is_empty());
476
477 let f = [field("a", parse_quote!(U256), None, None)];
479 let r = allocate_slots(&f).unwrap();
480 assert!(
481 matches!(r[0].assigned_slot, SlotAssignment::Auto { base_slot } if base_slot == U256::ZERO)
482 );
483
484 let f = [field("a", parse_quote!(U256), Some(5), None)];
486 let r = allocate_slots(&f).unwrap();
487 assert!(matches!(r[0].assigned_slot, SlotAssignment::Manual(s) if s == U256::from(5)));
488
489 let fields = vec![
491 field("a", parse_quote!(u8), None, None),
492 field("b", parse_quote!(u16), None, None),
493 field("c", parse_quote!(u32), None, None),
494 ];
495 for r in allocate_slots(&fields).unwrap() {
496 assert_eq!(*r.assigned_slot.ref_slot(), U256::ZERO);
497 }
498
499 let fields = vec![
501 field("a", parse_quote!(u8), None, None),
502 field("b", parse_quote!(U256), Some(10), None),
503 field("c", parse_quote!(u8), None, None),
504 ];
505 let r = allocate_slots(&fields).unwrap();
506 assert_eq!(*r[0].assigned_slot.ref_slot(), U256::ZERO);
507 assert!(matches!(r[1].assigned_slot, SlotAssignment::Manual(_)));
508 assert_eq!(*r[2].assigned_slot.ref_slot(), U256::ZERO);
509
510 let fields = vec![
512 field("a", parse_quote!(u8), None, None),
513 field("b", parse_quote!(u8), None, Some(5)),
514 field("c", parse_quote!(u8), None, None),
515 ];
516 let r = allocate_slots(&fields).unwrap();
517 assert_eq!(*r[0].assigned_slot.ref_slot(), U256::ZERO);
518 assert_eq!(*r[1].assigned_slot.ref_slot(), U256::from(5));
519 assert_eq!(*r[2].assigned_slot.ref_slot(), U256::from(5));
520
521 let fields = vec![
523 field("a", parse_quote!(u8), None, None),
524 field("b", parse_quote!(u8), None, Some(3)),
525 field("c", parse_quote!(u8), None, None),
526 field("d", parse_quote!(u8), None, Some(10)),
527 field("e", parse_quote!(u8), None, None),
528 ];
529 let r = allocate_slots(&fields).unwrap();
530 let expected = [0u64, 3, 3, 10, 10];
531 for (i, exp) in expected.iter().enumerate() {
532 assert_eq!(*r[i].assigned_slot.ref_slot(), U256::from(*exp));
533 }
534
535 let fields = vec![
537 field("a", parse_quote!(u8), None, None),
538 field("b", parse_quote!(U256), Some(50), None),
539 field("c", parse_quote!(u8), None, Some(20)),
540 field("d", parse_quote!(u8), None, None),
541 field("e", parse_quote!(U256), Some(99), None),
542 field("f", parse_quote!(u8), None, None),
543 ];
544 let r = allocate_slots(&fields).unwrap();
545 assert_eq!(*r[0].assigned_slot.ref_slot(), U256::ZERO);
546 assert!(matches!(r[1].assigned_slot, SlotAssignment::Manual(s) if s == U256::from(50)));
547 assert_eq!(*r[2].assigned_slot.ref_slot(), U256::from(20));
548 assert_eq!(*r[3].assigned_slot.ref_slot(), U256::from(20));
549 assert!(matches!(r[4].assigned_slot, SlotAssignment::Manual(s) if s == U256::from(99)));
550 assert_eq!(*r[5].assigned_slot.ref_slot(), U256::from(20));
551
552 let fields = vec![
554 field("a", parse_quote!(u8), None, Some(0)),
555 field("b", parse_quote!(u8), None, None),
556 ];
557 let r = allocate_slots(&fields).unwrap();
558 assert_eq!(*r[0].assigned_slot.ref_slot(), U256::ZERO);
559 assert_eq!(*r[1].assigned_slot.ref_slot(), U256::ZERO);
560
561 let f = [field("x", parse_quote!(U256), Some(u64::MAX), None)];
563 let r = allocate_slots(&f).unwrap();
564 assert_eq!(*r[0].assigned_slot.ref_slot(), U256::from(u64::MAX));
565
566 let fields = vec![
568 field("alpha", parse_quote!(u64), None, None),
569 field("beta", parse_quote!(U256), Some(7), None),
570 ];
571 let r = allocate_slots(&fields).unwrap();
572 assert_eq!(r[0].name.to_string(), "alpha");
573 assert_eq!(r[1].name.to_string(), "beta");
574
575 let fields = vec![
577 field("name", parse_quote!(String), None, None),
578 field("balances", parse_quote!(Mapping<Address, U256>), None, None),
579 field("supply", parse_quote!(U256), None, None),
580 ];
581 let r = allocate_slots(&fields).unwrap();
582 assert!(matches!(r[0].kind, FieldKind::Direct(_)));
583 assert!(matches!(r[1].kind, FieldKind::Mapping { .. }));
584 assert!(matches!(r[2].kind, FieldKind::Direct(_)));
585 }
586
587 #[test]
588 fn classify_field_type_variations() {
589 let assert_direct = |ty: Type| {
590 assert!(matches!(
591 classify_field_type(&ty).unwrap(),
592 FieldKind::Direct(_)
593 ));
594 };
595 let assert_mapping = |ty: Type| {
596 assert!(matches!(
597 classify_field_type(&ty).unwrap(),
598 FieldKind::Mapping { .. }
599 ));
600 };
601
602 assert_direct(parse_quote!(u64));
604 assert_direct(parse_quote!(U256));
605 assert_direct(parse_quote!(String));
606
607 assert_direct(parse_quote!(Vec<u8>));
609 assert_direct(parse_quote!(Option<u32>));
610
611 assert_mapping(parse_quote!(Mapping<Address, U256>));
613
614 if let FieldKind::Mapping { value, .. } =
616 classify_field_type(&parse_quote!(Mapping<Address, Mapping<Address, U256>>)).unwrap()
617 {
618 assert!(extract_mapping_types(value).is_some());
619 } else {
620 panic!("expected nested Mapping");
621 }
622 }
623
624 #[test]
625 fn packing_constants_variations() {
626 let span = proc_macro2::Span::call_site();
627
628 let c = PackingConstants::new(&Ident::new("total_supply", span));
630 assert_eq!(c.location().to_string(), "TOTAL_SUPPLY_LOC");
631 let (slot, offset) = c.into_tuple();
632 assert_eq!(slot.to_string(), "TOTAL_SUPPLY");
633 assert_eq!(offset.to_string(), "TOTAL_SUPPLY_OFFSET");
634
635 let c = PackingConstants::new(&Ident::new("x", span));
637 assert_eq!(c.slot().to_string(), "X");
638 assert_eq!(c.offset().to_string(), "X_OFFSET");
639 assert_eq!(c.location().to_string(), "X_LOC");
640
641 let c = PackingConstants::new(&Ident::new("balance", span));
643 assert_eq!(c.slot().to_string(), "BALANCE");
644 assert_eq!(c.location().to_string(), "BALANCE_LOC");
645
646 assert_eq!(const_name(&Ident::new("my_field", span)), "MY_FIELD");
648 assert_eq!(const_name(&Ident::new("x", span)), "X");
649 }
650
651 #[test]
652 fn slot_assignment_ref_slot_variations() {
653 let assert_ref_slot = |assignment: SlotAssignment, expected: U256| {
654 assert_eq!(*assignment.ref_slot(), expected);
655 };
656
657 assert_ref_slot(SlotAssignment::Manual(U256::from(42)), U256::from(42));
659 assert_ref_slot(SlotAssignment::Manual(U256::ZERO), U256::ZERO);
660 assert_ref_slot(
661 SlotAssignment::Manual(U256::from(u64::MAX)),
662 U256::from(u64::MAX),
663 );
664
665 assert_ref_slot(
667 SlotAssignment::Auto {
668 base_slot: U256::from(7),
669 },
670 U256::from(7),
671 );
672 assert_ref_slot(
673 SlotAssignment::Auto {
674 base_slot: U256::ZERO,
675 },
676 U256::ZERO,
677 );
678 assert_ref_slot(
679 SlotAssignment::Auto {
680 base_slot: U256::from(u64::MAX),
681 },
682 U256::from(u64::MAX),
683 );
684 }
685}