Fix HoldReason and FreezeReason variant count

substrate/frame/support/procedural/src/construct_runtime/expand/composite_helper.rs

@@ -68,3 +68,21 @@ pub(crate) fn expand_variant(
 		}
 	}
 }
+
+pub(crate) fn expand_variant_count(
+	composite_name: &str,
+	path: &PalletPath,
+	instance: Option<&Ident>,
+) -> TokenStream {
+	let composite_name = quote::format_ident!("{}", composite_name);
+
+	if let Some(inst) = instance {
+		quote! {
+			#path::#composite_name::<#path::#inst>::VARIANT_COUNT
+		}
+	} else {
+		quote! {
+			#path::#composite_name::VARIANT_COUNT
+		}
+	}
+}

substrate/frame/support/procedural/src/construct_runtime/expand/freeze_reason.rs

@@ -23,6 +23,7 @@ use quote::quote;
 pub fn expand_outer_freeze_reason(pallet_decls: &[Pallet], scrate: &TokenStream) -> TokenStream {
 	let mut conversion_fns = Vec::new();
 	let mut freeze_reason_variants = Vec::new();
+	let mut freeze_reason_variants_count = Vec::new();
 	for decl in pallet_decls {
 		if let Some(_) = decl.find_part("FreezeReason") {
 			let variant_name = &decl.name;
@@ -44,9 +45,14 @@ pub fn expand_outer_freeze_reason(pallet_decls: &[Pallet], scrate: &TokenStream)
 				instance,
 				variant_name,
 			));
+
+			freeze_reason_variants_count.push(composite_helper::expand_variant_count(
+				"FreezeReason",
+				path,
+				instance,
+			));
 		}
 	}
-	let freeze_reason_variants_count = freeze_reason_variants.len() as u32;
 
 	quote! {
 		/// A reason for placing a freeze on funds.
@@ -61,7 +67,7 @@ pub fn expand_outer_freeze_reason(pallet_decls: &[Pallet], scrate: &TokenStream)
 		}
 
 		impl #scrate::traits::VariantCount for RuntimeFreezeReason {
-			const VARIANT_COUNT: u32 = #freeze_reason_variants_count;
+			const VARIANT_COUNT: u32 = #( #freeze_reason_variants_count )+*;
 		}
 
 		#( #conversion_fns )*

substrate/frame/support/procedural/src/construct_runtime/expand/hold_reason.rs

@@ -23,6 +23,7 @@ use quote::quote;
 pub fn expand_outer_hold_reason(pallet_decls: &[Pallet], scrate: &TokenStream) -> TokenStream {
 	let mut conversion_fns = Vec::new();
 	let mut hold_reason_variants = Vec::new();
+	let mut hold_reason_variants_count = Vec::new();
 	for decl in pallet_decls {
 		if let Some(_) = decl.find_part("HoldReason") {
 			let variant_name = &decl.name;
@@ -44,9 +45,14 @@ pub fn expand_outer_hold_reason(pallet_decls: &[Pallet], scrate: &TokenStream) -
 				instance,
 				variant_name,
 			));
+
+			hold_reason_variants_count.push(composite_helper::expand_variant_count(
+				"HoldReason",
+				path,
+				instance,
+			));
 		}
 	}
-	let hold_reason_variants_count = hold_reason_variants.len() as u32;
 
 	quote! {
 		/// A reason for placing a hold on funds.
@@ -61,7 +67,7 @@ pub fn expand_outer_hold_reason(pallet_decls: &[Pallet], scrate: &TokenStream) -
 		}
 
 		impl #scrate::traits::VariantCount for RuntimeHoldReason {
-			const VARIANT_COUNT: u32 = #hold_reason_variants_count;
+			const VARIANT_COUNT: u32 = #( #hold_reason_variants_count )+*;
 		}
 
 		#( #conversion_fns )*

substrate/frame/support/procedural/src/lib.rs

@@ -1506,18 +1506,77 @@ pub fn origin(_: TokenStream, _: TokenStream) -> TokenStream {
 ///
 /// ```ignore
 /// Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Encode, Decode, MaxEncodedLen, TypeInfo,
-/// RuntimeDebug
+/// RuntimeDebug,
+/// CompositeEnumVariantCount,
 /// ```
 ///
 /// For ease of usage, when no `#[derive]` attributes are found for the enum under
 /// `#[pallet::composite_enum]`, the aforementioned traits are automatically derived for it. The
 /// inverse is also true: if there are any `#[derive]` attributes found for the enum, then no traits
-/// will automatically be derived for it.
+/// will automatically be derived for it (this implies that you need to provide the `frame_support::traits::VariantCount` implementation).
 #[proc_macro_attribute]
 pub fn composite_enum(_: TokenStream, _: TokenStream) -> TokenStream {
 	pallet_macro_stub()
 }
 
+/// Derive macro for implementing [`frame_support::traits::VariantCount`] for enum type.
+///
+/// This macro works only for the enum which does not have variants or variants do not have fields.
+///
+/// This macro is automatically derived for `#[pallet::composite_enum]` enums.
+#[proc_macro_derive(CompositeEnumVariantCount)]
+pub fn derive_composite_enum_variant_count(input: TokenStream) -> TokenStream {
+	let input: syn::DeriveInput = match syn::parse(input) {
+		Ok(input) => input,
+		Err(e) => return e.to_compile_error().into(),
+	};
+
+	// only `Enum` type is expected here (see also `CompositeDef::try_from`)
+	let data_enum = match &input.data {
+		syn::Data::Enum(ref data_enum) => data_enum,
+		_ => return syn::Error::new(
+			input.ident.span(),
+			"Invalid type: expected enum item",
+		).to_compile_error().into()
+	};
+
+		// check variants: composite enums support only field-less enum variants. This is because fields can introduce too many possibilities, making it challenging to compute a fixed variant count.
+		// The only exception is `__Ignore` which can be added by `CompositeDef::try_from`.
+		for variant in &data_enum.variants {
+			match variant.fields {
+				syn::Fields::Named(_) | syn::Fields::Unnamed(_) => {
+					let variant_name = variant.ident.to_string();
+					if variant_name != "__Ignore" {
+						return syn::Error::new(
+							variant.ident.span(),
+							"The composite enum does not support variants with fields!",
+						).to_compile_error().into()
+					}
+				},
+				syn::Fields::Unit => (),
+			}
+		}
+	// get variants count
+	let variants_count = data_enum.variants.len() as u32;
+
+	// add `VariantCount` implementation for composite enum
+	let frame_support = match generate_access_from_frame_or_crate("frame-support") {
+		Ok(c) => c,
+		Err(e) => return e.into_compile_error().into(),
+	};
+	let frame_support = &frame_support;
+
+	let name = &input.ident;
+	let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+
+	quote! {
+		impl #impl_generics #frame_support::traits::VariantCount for #name #ty_generics #where_clause {
+			const VARIANT_COUNT: u32 = #variants_count;
+		}
+	}
+	.into()
+}
+
 ///
 /// ---
 ///

substrate/frame/support/procedural/src/pallet/parse/composite.rs

@@ -138,6 +138,7 @@ impl CompositeDef {
 					#scrate::__private::codec::Encode, #scrate::__private::codec::Decode, #scrate::__private::codec::MaxEncodedLen,
 					#scrate::__private::scale_info::TypeInfo,
 					#scrate::__private::RuntimeDebug,
+					#scrate::CompositeEnumVariantCount,
 				)]
 			};
 			item.attrs.push(derive_attr);

substrate/frame/support/src/lib.rs

@@ -497,7 +497,7 @@ pub fn debug(data: &impl sp_std::fmt::Debug) {
 
 #[doc(inline)]
 pub use frame_support_procedural::{
-	construct_runtime, match_and_insert, transactional, PalletError, RuntimeDebugNoBound,
+	construct_runtime, match_and_insert, transactional, CompositeEnumVariantCount, PalletError, RuntimeDebugNoBound,
 };
 
 #[doc(hidden)]
@@ -851,7 +851,7 @@ pub mod pallet_prelude {
 			BuildGenesisConfig, ConstU32, EnsureOrigin, Get, GetDefault, GetStorageVersion, Hooks,
 			IsType, PalletInfoAccess, StorageInfoTrait, StorageVersion, Task, TypedGet,
 		},
-		Blake2_128, Blake2_128Concat, Blake2_256, CloneNoBound, DebugNoBound, EqNoBound, Identity,
+		Blake2_128, Blake2_128Concat, Blake2_256, CloneNoBound, CompositeEnumVariantCount, DebugNoBound, EqNoBound, Identity,
 		PartialEqNoBound, RuntimeDebugNoBound, Twox128, Twox256, Twox64Concat,
 	};
 	pub use codec::{Decode, Encode, MaxEncodedLen};