diff --git a/frame/examples/basic/src/lib.rs b/frame/examples/basic/src/lib.rs index d5045157dade7..77f1d9da9ef39 100644 --- a/frame/examples/basic/src/lib.rs +++ b/frame/examples/basic/src/lib.rs @@ -610,6 +610,7 @@ pub mod pallet { // A map that has enumerable entries. #[pallet::storage] #[pallet::getter(fn bar)] + #[pallet::proof_size = MaxEncodedLen] pub(super) type Bar = StorageMap<_, Blake2_128Concat, T::AccountId, T::Balance>; // this one uses the query kind: `ValueQuery`, we'll demonstrate the usage of 'mutate' API. diff --git a/frame/support/procedural/src/pallet/parse/storage.rs b/frame/support/procedural/src/pallet/parse/storage.rs index 8b551ab31d6c3..14427993dff34 100644 --- a/frame/support/procedural/src/pallet/parse/storage.rs +++ b/frame/support/procedural/src/pallet/parse/storage.rs @@ -29,6 +29,9 @@ mod keyword { syn::custom_keyword!(storage_prefix); syn::custom_keyword!(unbounded); syn::custom_keyword!(whitelist_storage); + syn::custom_keyword!(proof_size); + syn::custom_keyword!(MaxEncodedLen); + syn::custom_keyword!(Measured); syn::custom_keyword!(OptionQuery); syn::custom_keyword!(ResultQuery); syn::custom_keyword!(ValueQuery); @@ -39,11 +42,19 @@ mod keyword { /// * `#[pallet::storage_prefix = "CustomName"]` /// * `#[pallet::unbounded]` /// * `#[pallet::whitelist_storage] +/// * `#[pallet::proof_size = MaxEncodedLen] pub enum PalletStorageAttr { Getter(syn::Ident, proc_macro2::Span), StorageName(syn::LitStr, proc_macro2::Span), Unbounded(proc_macro2::Span), WhitelistStorage(proc_macro2::Span), + ProofSize(ProofSizeAttribute, proc_macro2::Span), +} + +#[derive(PartialEq, Debug)] +pub enum ProofSizeAttribute { + MaxEncodedLen, + Measured, } impl PalletStorageAttr { @@ -52,7 +63,8 @@ impl PalletStorageAttr { Self::Getter(_, span) | Self::StorageName(_, span) | Self::Unbounded(span) | - Self::WhitelistStorage(span) => *span, + Self::WhitelistStorage(span) | + Self::ProofSize(_, span) => *span, } } } @@ -93,6 +105,23 @@ impl syn::parse::Parse for PalletStorageAttr { } else if lookahead.peek(keyword::whitelist_storage) { content.parse::()?; Ok(Self::WhitelistStorage(attr_span)) + } else if lookahead.peek(keyword::proof_size) { + content.parse::()?; + content.parse::()?; + let lookahead = content.lookahead1(); + + if lookahead.peek(keyword::MaxEncodedLen) { + content.parse::().expect("Checked above"); + Ok(Self::ProofSize(ProofSizeAttribute::MaxEncodedLen, attr_span)) + } else if lookahead.peek(keyword::Measured) { + content.parse::().expect("Checked above"); + Ok(Self::ProofSize(ProofSizeAttribute::Measured, attr_span)) + } else { + Err(syn::Error::new( + attr_span, + format!("Invalid value for the proof_size attribute: {:?}", lookahead.error()), + )) + } } else { Err(lookahead.error()) } @@ -104,14 +133,16 @@ struct PalletStorageAttrInfo { rename_as: Option, unbounded: bool, whitelisted: bool, + proof_size: Option, } impl PalletStorageAttrInfo { fn from_attrs(attrs: Vec) -> syn::Result { let mut getter = None; let mut rename_as = None; - let mut unbounded = false; - let mut whitelisted = false; + let mut proof_size = None; + let (mut unbounded, mut whitelisted) = (false, false); + for attr in attrs { match attr { PalletStorageAttr::Getter(ident, ..) if getter.is_none() => getter = Some(ident), @@ -119,6 +150,8 @@ impl PalletStorageAttrInfo { rename_as = Some(name), PalletStorageAttr::Unbounded(..) if !unbounded => unbounded = true, PalletStorageAttr::WhitelistStorage(..) if !whitelisted => whitelisted = true, + PalletStorageAttr::ProofSize(mode, ..) if proof_size.is_none() => + proof_size = Some(mode), attr => return Err(syn::Error::new( attr.attr_span(), @@ -127,7 +160,7 @@ impl PalletStorageAttrInfo { } } - Ok(PalletStorageAttrInfo { getter, rename_as, unbounded, whitelisted }) + Ok(PalletStorageAttrInfo { getter, rename_as, unbounded, whitelisted, proof_size }) } } @@ -185,6 +218,8 @@ pub struct StorageDef { pub unbounded: bool, /// Whether or not reads to this storage key will be ignored by benchmarking pub whitelisted: bool, + /// How the proof size should be calculated by the PoV benchmarking. + pub proof_size: Option, } /// The parsed generic from the @@ -687,9 +722,12 @@ impl StorageDef { }; let attrs: Vec = helper::take_item_pallet_attrs(&mut item.attrs)?; - let PalletStorageAttrInfo { getter, rename_as, mut unbounded, whitelisted } = + let PalletStorageAttrInfo { getter, rename_as, mut unbounded, whitelisted, proof_size } = PalletStorageAttrInfo::from_attrs(attrs)?; + if unbounded && proof_size == Some(ProofSizeAttribute::MaxEncodedLen) { + return Err(syn::Error::new(item.span(), "Storage item cannot be 'unbounded' ahd have 'MaxEncodedLen' proof size at the same time.")) + } // set all storages to be unbounded if dev_mode is enabled unbounded |= dev_mode; let cfg_attrs = helper::get_item_cfg_attrs(&item.attrs); @@ -832,6 +870,7 @@ impl StorageDef { named_generics, unbounded, whitelisted, + proof_size, }) } }