tempo_precompiles_macros/
storable_primitives.rs1use proc_macro2::TokenStream;
4use quote::quote;
5
6pub(crate) const RUST_INT_SIZES: &[usize] = &[8, 16, 32, 64, 128];
7pub(crate) const ALLOY_INT_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256];
8
9#[derive(Debug, Clone)]
13enum StorableConversionStrategy {
14 U256, UnsignedRust,
16 UnsignedAlloy(proc_macro2::Ident),
17 SignedRust(proc_macro2::Ident),
18 SignedAlloy(proc_macro2::Ident),
19 FixedBytes(usize),
20}
21
22#[derive(Debug, Clone)]
24enum StorageKeyStrategy {
25 Simple, WithSize(usize), SignedRaw(usize), AsSlice, }
30
31#[derive(Debug, Clone)]
33struct TypeConfig {
34 type_path: TokenStream,
35 byte_count: usize,
36 storable_strategy: StorableConversionStrategy,
37 storage_key_strategy: StorageKeyStrategy,
38}
39
40fn gen_storable_layout_impl(type_path: &TokenStream, byte_count: usize) -> TokenStream {
44 quote! {
45 impl StorableType for #type_path {
46 const LAYOUT: Layout = Layout::Bytes(#byte_count);
47 type Handler = crate::storage::Slot<Self>;
48
49 fn handle(slot: U256, ctx: LayoutCtx, address: ::alloy::primitives::Address) -> Self::Handler {
50 crate::storage::Slot::new_with_ctx(slot, ctx, address)
51 }
52 }
53 }
54}
55
56fn gen_storage_key_impl(type_path: &TokenStream, strategy: &StorageKeyStrategy) -> TokenStream {
58 let conversion = match strategy {
59 StorageKeyStrategy::Simple => quote! { self.to_be_bytes() },
60 StorageKeyStrategy::WithSize(size) => quote! { self.to_be_bytes::<#size>() },
61 StorageKeyStrategy::SignedRaw(size) => quote! { self.into_raw().to_be_bytes::<#size>() },
62 StorageKeyStrategy::AsSlice => quote! { self.as_slice() },
63 };
64
65 quote! {
66 impl StorageKey for #type_path {
67 #[inline]
68 fn as_storage_bytes(&self) -> impl AsRef<[u8]> {
69 #conversion
70 }
71 }
72 }
73}
74
75fn gen_packable_impl(
78 type_path: &TokenStream,
79 strategy: &StorableConversionStrategy,
80) -> TokenStream {
81 match strategy {
82 StorableConversionStrategy::U256 => {
83 quote! {
84 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
85 impl Packable for #type_path {
86 #[inline]
87 fn to_word(&self) -> U256 {
88 *self
89 }
90
91 #[inline]
92 fn from_word(word: U256) -> crate::error::Result<Self> {
93 Ok(word)
94 }
95 }
96 }
97 }
98 StorableConversionStrategy::UnsignedRust => {
99 quote! {
100 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
101 impl Packable for #type_path {
102 #[inline]
103 fn to_word(&self) -> U256 {
104 ::alloy::primitives::U256::from(*self)
105 }
106
107 #[inline]
108 fn from_word(word: U256) -> crate::error::Result<Self> {
109 word.try_into().map_err(|_| crate::error::TempoPrecompileError::under_overflow())
110 }
111 }
112 }
113 }
114 StorableConversionStrategy::UnsignedAlloy(ty) => {
115 quote! {
116 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
117 impl Packable for #type_path {
118 #[inline]
119 fn to_word(&self) -> ::alloy::primitives::U256 {
120 ::alloy::primitives::U256::from(*self)
121 }
122
123 #[inline]
124 fn from_word(word: ::alloy::primitives::U256) -> crate::error::Result<Self> {
125 if word > ::alloy::primitives::U256::from(::alloy::primitives::#ty::MAX) {
127 return Err(crate::error::TempoPrecompileError::under_overflow());
128 }
129 Ok(word.to::<Self>())
130 }
131 }
132 }
133 }
134 StorableConversionStrategy::SignedRust(unsigned_type) => {
135 quote! {
136 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
137 impl Packable for #type_path {
138 #[inline]
139 fn to_word(&self) -> U256 {
140 ::alloy::primitives::U256::from(*self as #unsigned_type)
142 }
143
144 #[inline]
145 fn from_word(word: U256) -> crate::error::Result<Self> {
146 let unsigned: #unsigned_type = word.try_into()
148 .map_err(|_| crate::error::TempoPrecompileError::under_overflow())?;
149 Ok(unsigned as Self)
150 }
151 }
152 }
153 }
154 StorableConversionStrategy::SignedAlloy(unsigned_type) => {
155 quote! {
156 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
157 impl Packable for #type_path {
158 #[inline]
159 fn to_word(&self) -> ::alloy::primitives::U256 {
160 ::alloy::primitives::U256::from(self.into_raw())
162 }
163
164 #[inline]
165 fn from_word(word: ::alloy::primitives::U256) -> crate::error::Result<Self> {
166 if word > ::alloy::primitives::U256::from(::alloy::primitives::#unsigned_type::MAX) {
168 return Err(crate::error::TempoPrecompileError::under_overflow());
169 }
170 let unsigned_val = word.to::<::alloy::primitives::#unsigned_type>();
172 Ok(Self::from_raw(unsigned_val))
173 }
174 }
175 }
176 }
177 StorableConversionStrategy::FixedBytes(size) => {
178 quote! {
179 impl crate::storage::types::sealed::OnlyPrimitives for #type_path {}
180 impl Packable for #type_path {
181 #[inline]
182 fn to_word(&self) -> ::alloy::primitives::U256 {
183 let mut bytes = [0u8; 32];
184 bytes[..#size].copy_from_slice(&self[..]);
185 ::alloy::primitives::U256::from_be_bytes(bytes)
186 }
187
188 #[inline]
189 fn from_word(word: ::alloy::primitives::U256) -> crate::error::Result<Self> {
190 let bytes = word.to_be_bytes::<32>();
191 let mut fixed_bytes = [0u8; #size];
192 fixed_bytes.copy_from_slice(&bytes[..#size]);
193 Ok(Self::from(fixed_bytes))
194 }
195 }
196 }
197 }
198 }
199}
200
201fn gen_complete_impl_set(config: &TypeConfig) -> TokenStream {
205 let storable_type_impl = gen_storable_layout_impl(&config.type_path, config.byte_count);
206 let packable_impl = gen_packable_impl(&config.type_path, &config.storable_strategy);
207 let storage_key_impl = gen_storage_key_impl(&config.type_path, &config.storage_key_strategy);
208
209 quote! {
210 #storable_type_impl
211 #packable_impl
212 #storage_key_impl
213 }
214}
215
216pub(crate) fn gen_storable_rust_ints() -> TokenStream {
218 let mut impls = Vec::with_capacity(RUST_INT_SIZES.len() * 2);
219
220 for size in RUST_INT_SIZES {
221 let unsigned_type = quote::format_ident!("u{}", size);
222 let signed_type = quote::format_ident!("i{}", size);
223 let byte_count = size / 8;
224
225 let unsigned_config = TypeConfig {
227 type_path: quote! { #unsigned_type },
228 byte_count,
229 storable_strategy: StorableConversionStrategy::UnsignedRust,
230 storage_key_strategy: StorageKeyStrategy::Simple,
231 };
232 impls.push(gen_complete_impl_set(&unsigned_config));
233
234 let signed_config = TypeConfig {
236 type_path: quote! { #signed_type },
237 byte_count,
238 storable_strategy: StorableConversionStrategy::SignedRust(unsigned_type.clone()),
239 storage_key_strategy: StorageKeyStrategy::Simple,
240 };
241 impls.push(gen_complete_impl_set(&signed_config));
242 }
243
244 quote! {
245 #(#impls)*
246 }
247}
248
249fn gen_alloy_integers() -> Vec<TokenStream> {
251 let mut impls = Vec::with_capacity(ALLOY_INT_SIZES.len() * 2);
252
253 for &size in ALLOY_INT_SIZES {
254 let unsigned_type = quote::format_ident!("U{}", size);
255 let signed_type = quote::format_ident!("I{}", size);
256 let byte_count = size / 8;
257
258 let unsigned_config = TypeConfig {
260 type_path: quote! { ::alloy::primitives::#unsigned_type },
261 byte_count,
262 storable_strategy: if size == 256 {
263 StorableConversionStrategy::U256
264 } else {
265 StorableConversionStrategy::UnsignedAlloy(unsigned_type.clone())
266 },
267 storage_key_strategy: StorageKeyStrategy::WithSize(byte_count),
268 };
269 impls.push(gen_complete_impl_set(&unsigned_config));
270
271 let signed_config = TypeConfig {
273 type_path: quote! { ::alloy::primitives::#signed_type },
274 byte_count,
275 storable_strategy: StorableConversionStrategy::SignedAlloy(unsigned_type.clone()),
276 storage_key_strategy: StorageKeyStrategy::SignedRaw(byte_count),
277 };
278 impls.push(gen_complete_impl_set(&signed_config));
279 }
280
281 impls
282}
283
284fn gen_fixed_bytes(sizes: &[usize]) -> Vec<TokenStream> {
286 let mut impls = Vec::with_capacity(sizes.len());
287
288 for &size in sizes {
289 let config = TypeConfig {
291 type_path: quote! { ::alloy::primitives::FixedBytes<#size> },
292 byte_count: size,
293 storable_strategy: StorableConversionStrategy::FixedBytes(size),
294 storage_key_strategy: StorageKeyStrategy::AsSlice,
295 };
296 impls.push(gen_complete_impl_set(&config));
297 }
298
299 impls
300}
301
302pub(crate) fn gen_storable_alloy_bytes() -> TokenStream {
304 let sizes: Vec<usize> = (1..=32).collect();
305 let impls = gen_fixed_bytes(&sizes);
306
307 quote! {
308 #(#impls)*
309 }
310}
311
312pub(crate) fn gen_storable_alloy_ints() -> TokenStream {
314 let impls = gen_alloy_integers();
315
316 quote! {
317 #(#impls)*
318 }
319}
320
321#[derive(Debug, Clone)]
325struct ArrayConfig {
326 elem_type: TokenStream,
327 array_size: usize,
328 elem_byte_count: usize,
329 elem_is_packable: bool,
330}
331
332fn is_packable(byte_count: usize) -> bool {
334 byte_count < 32 && 32 % byte_count == 0
335}
336
337fn gen_array_impl(config: &ArrayConfig) -> TokenStream {
339 let ArrayConfig {
340 elem_type,
341 array_size,
342 elem_byte_count,
343 elem_is_packable,
344 } = config;
345
346 let slot_count = if *elem_is_packable {
348 (*array_size * elem_byte_count).div_ceil(32)
350 } else {
351 *array_size
353 };
354
355 let load_impl = if *elem_is_packable {
356 gen_packed_array_load(array_size, elem_byte_count)
357 } else {
358 gen_unpacked_array_load(array_size)
359 };
360
361 let store_impl = if *elem_is_packable {
362 gen_packed_array_store(array_size, elem_byte_count)
363 } else {
364 gen_unpacked_array_store()
365 };
366
367 quote! {
368 impl crate::storage::StorableType for [#elem_type; #array_size] {
370 const LAYOUT: crate::storage::Layout = crate::storage::Layout::Slots(#slot_count);
372
373 type Handler = crate::storage::types::array::ArrayHandler<#elem_type, #array_size>;
374
375 fn handle(slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx, address: ::alloy::primitives::Address) -> Self::Handler {
376 debug_assert_eq!(ctx, crate::storage::LayoutCtx::FULL, "Arrays cannot be packed");
377 Self::Handler::new(slot, address)
378 }
379 }
380
381 impl crate::storage::Storable for [#elem_type; #array_size] {
383 #[inline]
384 fn load<S: crate::storage::StorageOps>(storage: &S, slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx) -> crate::error::Result<Self> {
385 debug_assert_eq!(
386 ctx, crate::storage::LayoutCtx::FULL,
387 "Arrays can only be loaded with LayoutCtx::FULL"
388 );
389
390 use crate::storage::packing::{calc_element_slot, calc_element_offset, extract_packed_value};
391 let base_slot = slot;
392 #load_impl
393 }
394
395 #[inline]
396 fn store<S: crate::storage::StorageOps>(&self, storage: &mut S, slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx) -> crate::error::Result<()> {
397 debug_assert_eq!(
398 ctx, crate::storage::LayoutCtx::FULL,
399 "Arrays can only be stored with LayoutCtx::FULL"
400 );
401
402 use crate::storage::packing::{calc_element_slot, calc_element_offset, insert_packed_value};
403 let base_slot = slot;
404 #store_impl
405 }
406
407 }
409
410 impl crate::storage::StorageKey for [#elem_type; #array_size] {
412 #[inline]
413 fn as_storage_bytes(&self) -> impl AsRef<[u8]> {
414 use crate::storage::StorageKey;
415 let mut bytes = Vec::with_capacity(#array_size * <#elem_type as crate::storage::StorableType>::BYTES);
416 for elem in self.iter() {
417 bytes.extend_from_slice(elem.as_storage_bytes().as_ref());
418 }
419 bytes
420 }
421 }
422 }
423}
424
425fn gen_packed_array_load(array_size: &usize, elem_byte_count: &usize) -> TokenStream {
427 quote! {
428 let mut result = [Default::default(); #array_size];
429 for i in 0..#array_size {
430 let slot_idx = calc_element_slot(i, #elem_byte_count);
431 let offset = calc_element_offset(i, #elem_byte_count);
432 let slot_addr = base_slot + U256::from(slot_idx);
433 let slot_value = storage.load(slot_addr)?;
434 result[i] = extract_packed_value(slot_value, offset, #elem_byte_count)?;
435 }
436 Ok(result)
437 }
438}
439
440fn gen_packed_array_store(array_size: &usize, elem_byte_count: &usize) -> TokenStream {
442 quote! {
443 let slot_count = (#array_size * #elem_byte_count).div_ceil(32);
445
446 for slot_idx in 0..slot_count {
448 let slot_addr = base_slot + U256::from(slot_idx);
449 let mut slot_value = U256::ZERO;
450
451 for i in 0..#array_size {
453 let elem_slot = calc_element_slot(i, #elem_byte_count);
454 if elem_slot == slot_idx {
455 let offset = calc_element_offset(i, #elem_byte_count);
456 slot_value = insert_packed_value(slot_value, &self[i], offset, #elem_byte_count)?;
457 }
458 }
459
460 storage.store(slot_addr, slot_value)?;
461 }
462 Ok(())
463 }
464}
465
466fn gen_unpacked_array_load(array_size: &usize) -> TokenStream {
468 quote! {
469 let mut result = [Default::default(); #array_size];
470 for i in 0..#array_size {
471 let elem_slot = base_slot + ::alloy::primitives::U256::from(i);
472 result[i] = crate::storage::Storable::load(storage, elem_slot, crate::storage::LayoutCtx::FULL)?;
473 }
474 Ok(result)
475 }
476}
477
478fn gen_unpacked_array_store() -> TokenStream {
480 quote! {
481 for (i, elem) in self.iter().enumerate() {
482 let elem_slot = base_slot + ::alloy::primitives::U256::from(i);
483 crate::storage::Storable::store(elem, storage, elem_slot, crate::storage::LayoutCtx::FULL)?;
484 }
485 Ok(())
486 }
487}
488
489fn gen_arrays_for_type(
491 elem_type: TokenStream,
492 elem_byte_count: usize,
493 sizes: &[usize],
494) -> Vec<TokenStream> {
495 let elem_is_packable = is_packable(elem_byte_count);
496
497 sizes
498 .iter()
499 .map(|&size| {
500 let config = ArrayConfig {
501 elem_type: elem_type.clone(),
502 array_size: size,
503 elem_byte_count,
504 elem_is_packable,
505 };
506 gen_array_impl(&config)
507 })
508 .collect()
509}
510
511pub(crate) fn gen_storable_arrays() -> TokenStream {
513 let mut all_impls = Vec::new();
514 let sizes: Vec<usize> = (1..=32).collect();
515
516 for &bit_size in RUST_INT_SIZES {
518 let type_ident = quote::format_ident!("u{}", bit_size);
519 let byte_count = bit_size / 8;
520 all_impls.extend(gen_arrays_for_type(
521 quote! { #type_ident },
522 byte_count,
523 &sizes,
524 ));
525 }
526
527 for &bit_size in RUST_INT_SIZES {
529 let type_ident = quote::format_ident!("i{}", bit_size);
530 let byte_count = bit_size / 8;
531 all_impls.extend(gen_arrays_for_type(
532 quote! { #type_ident },
533 byte_count,
534 &sizes,
535 ));
536 }
537
538 for &bit_size in ALLOY_INT_SIZES {
540 let type_ident = quote::format_ident!("U{}", bit_size);
541 let byte_count = bit_size / 8;
542 all_impls.extend(gen_arrays_for_type(
543 quote! { ::alloy::primitives::#type_ident },
544 byte_count,
545 &sizes,
546 ));
547 }
548
549 for &bit_size in ALLOY_INT_SIZES {
551 let type_ident = quote::format_ident!("I{}", bit_size);
552 let byte_count = bit_size / 8;
553 all_impls.extend(gen_arrays_for_type(
554 quote! { ::alloy::primitives::#type_ident },
555 byte_count,
556 &sizes,
557 ));
558 }
559
560 all_impls.extend(gen_arrays_for_type(
562 quote! { ::alloy::primitives::Address },
563 20,
564 &sizes,
565 ));
566
567 for &byte_size in &[20, 32] {
569 all_impls.extend(gen_arrays_for_type(
570 quote! { ::alloy::primitives::FixedBytes<#byte_size> },
571 byte_size,
572 &sizes,
573 ));
574 }
575
576 quote! {
577 #(#all_impls)*
578 }
579}
580
581pub(crate) fn gen_nested_arrays() -> TokenStream {
583 let mut all_impls = Vec::new();
584
585 for inner in &[2usize, 4, 8, 16] {
588 let inner_slots = inner.div_ceil(32); let max_outer = 32 / inner_slots.max(1);
590
591 for outer in 1..=max_outer.min(32) {
592 all_impls.extend(gen_arrays_for_type(
593 quote! { [u8; #inner] },
594 inner_slots * 32, &[outer],
596 ));
597 }
598 }
599
600 for inner in &[2usize, 4, 8] {
602 let inner_slots = (inner * 2).div_ceil(32);
603 let max_outer = 32 / inner_slots.max(1);
604
605 for outer in 1..=max_outer.min(16) {
606 all_impls.extend(gen_arrays_for_type(
607 quote! { [u16; #inner] },
608 inner_slots * 32,
609 &[outer],
610 ));
611 }
612 }
613
614 quote! {
615 #(#all_impls)*
616 }
617}
618
619pub(crate) fn gen_struct_arrays(struct_type: TokenStream, array_sizes: &[usize]) -> TokenStream {
637 let impls: Vec<_> = array_sizes
638 .iter()
639 .map(|&size| gen_struct_array_impl(&struct_type, size))
640 .collect();
641
642 quote! {
643 #(#impls)*
644 }
645}
646
647fn gen_struct_array_impl(struct_type: &TokenStream, array_size: usize) -> TokenStream {
649 let struct_type_str = struct_type
651 .to_string()
652 .replace("::", "_")
653 .replace(['<', '>', ' ', '[', ']', ';'], "_");
654 let mod_ident = quote::format_ident!("__array_{}_{}", struct_type_str, array_size);
655
656 let load_impl = gen_struct_array_load(struct_type, array_size);
658 let store_impl = gen_struct_array_store(struct_type);
659
660 quote! {
661 mod #mod_ident {
663 use super::*;
664 pub const ELEM_SLOTS: usize = <#struct_type as crate::storage::StorableType>::SLOTS;
665 pub const ARRAY_LEN: usize = #array_size;
666 pub const SLOT_COUNT: usize = ARRAY_LEN * ELEM_SLOTS;
667 }
668
669 impl crate::storage::StorableType for [#struct_type; #array_size] {
671 const LAYOUT: crate::storage::Layout = crate::storage::Layout::Slots(#mod_ident::SLOT_COUNT);
672
673 type Handler = crate::storage::Slot<Self>;
674
675 fn handle(slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx, address: ::alloy::primitives::Address) -> Self::Handler {
676 crate::storage::Slot::new_with_ctx(slot, ctx, address)
677 }
678 }
679
680 impl crate::storage::Storable for [#struct_type; #array_size] {
682 #[inline]
683 fn load<S: crate::storage::StorageOps>(storage: &S, slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx) -> crate::error::Result<Self> {
684 debug_assert_eq!(
685 ctx, crate::storage::LayoutCtx::FULL,
686 "Struct arrays can only be loaded with LayoutCtx::FULL"
687 );
688 let base_slot = slot;
689 #load_impl
690 }
691
692 #[inline]
693 fn store<S: crate::storage::StorageOps>(&self, storage: &mut S, slot: ::alloy::primitives::U256, ctx: crate::storage::LayoutCtx) -> crate::error::Result<()> {
694 debug_assert_eq!(
695 ctx, crate::storage::LayoutCtx::FULL,
696 "Struct arrays can only be stored with LayoutCtx::FULL"
697 );
698 let base_slot = slot;
699 #store_impl
700 }
701
702 }
704
705 impl crate::storage::StorageKey for [#struct_type; #array_size] {
707 #[inline]
708 fn as_storage_bytes(&self) -> impl AsRef<[u8]> {
709 use crate::storage::StorageKey;
710 let mut bytes = Vec::with_capacity(#array_size * <#struct_type as crate::storage::StorableType>::SLOTS * 32);
712 for elem in self.iter() {
713 bytes.extend_from_slice(elem.as_storage_bytes().as_ref());
714 }
715 bytes
716 }
717 }
718 }
719}
720
721fn gen_struct_array_load(struct_type: &TokenStream, array_size: usize) -> TokenStream {
725 quote! {
726 let mut result = [Default::default(); #array_size];
727 for i in 0..#array_size {
728 let elem_slot = base_slot.checked_add(
730 ::alloy::primitives::U256::from(i).checked_mul(
731 ::alloy::primitives::U256::from(<#struct_type as crate::storage::StorableType>::SLOTS)
732 ).ok_or(crate::error::TempoError::SlotOverflow)?
733 ).ok_or(crate::error::TempoError::SlotOverflow)?;
734
735 result[i] = <#struct_type as crate::storage::Storable>::load(storage, elem_slot, crate::storage::LayoutCtx::FULL)?;
736 }
737 Ok(result)
738 }
739}
740
741fn gen_struct_array_store(struct_type: &TokenStream) -> TokenStream {
743 quote! {
744 for (i, elem) in self.iter().enumerate() {
745 let elem_slot = base_slot.checked_add(
747 ::alloy::primitives::U256::from(i).checked_mul(
748 ::alloy::primitives::U256::from(<#struct_type as crate::storage::StorableType>::SLOTS)
749 ).ok_or(crate::error::TempoError::SlotOverflow)?
750 ).ok_or(crate::error::TempoError::SlotOverflow)?;
751
752 <#struct_type as crate::storage::Storable>::store(elem, storage, elem_slot, crate::storage::LayoutCtx::FULL)?;
753 }
754 Ok(())
755 }
756}