WIP

Signed-off-by: Oliver Tale-Yazdi <[email protected]>

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<T: Config> = StorageMap<_, Blake2_128Concat, T::AccountId, T::Balance>;
 
 	// this one uses the query kind: `ValueQuery`, we'll demonstrate the usage of 'mutate' API.

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::<keyword::whitelist_storage>()?;
 			Ok(Self::WhitelistStorage(attr_span))
+		} else if lookahead.peek(keyword::proof_size) {
+			content.parse::<keyword::proof_size>()?;
+			content.parse::<syn::Token![=]>()?;
+			let lookahead = content.lookahead1();
+
+			if lookahead.peek(keyword::MaxEncodedLen) {
+				content.parse::<keyword::MaxEncodedLen>().expect("Checked above");
+				Ok(Self::ProofSize(ProofSizeAttribute::MaxEncodedLen, attr_span))
+			} else if lookahead.peek(keyword::Measured) {
+				content.parse::<keyword::Measured>().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,21 +133,25 @@ struct PalletStorageAttrInfo {
 	rename_as: Option<syn::LitStr>,
 	unbounded: bool,
 	whitelisted: bool,
+	proof_size: Option<ProofSizeAttribute>,
 }
 
 impl PalletStorageAttrInfo {
 	fn from_attrs(attrs: Vec<PalletStorageAttr>) -> syn::Result<Self> {
 		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),
 				PalletStorageAttr::StorageName(name, ..) if rename_as.is_none() =>
 					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<ProofSizeAttribute>,
 }
 
 /// The parsed generic from the
@@ -687,9 +722,12 @@ impl StorageDef {
 		};
 
 		let attrs: Vec<PalletStorageAttr> = 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,
 		})
 	}
 }