From 35c39c960804578396c78532c393d7385ee154b4 Mon Sep 17 00:00:00 2001
From: gupnik <17176722+gupnik@users.noreply.github.com>
Date: Mon, 4 Dec 2023 16:42:14 +0530
Subject: [PATCH] Fixes runtime type with doc parsing in derive_impl (#2594)

Step in https://github.com/paritytech/polkadot-sdk/issues/171

This PR fixes a bug in `derive_impl` causing `#[inject_runtime_type]`
attribute to be parsed incorrectly when docs are added.
---
 .../support/procedural/src/derive_impl.rs     | 41 +++++++++++++++----
 1 file changed, 32 insertions(+), 9 deletions(-)

diff --git a/substrate/frame/support/procedural/src/derive_impl.rs b/substrate/frame/support/procedural/src/derive_impl.rs
index 8b5e334f1f55..3e044053116b 100644
--- a/substrate/frame/support/procedural/src/derive_impl.rs
+++ b/substrate/frame/support/procedural/src/derive_impl.rs
@@ -46,11 +46,15 @@ pub struct PalletAttr {
 	typ: PalletAttrType,
 }
 
-fn get_first_item_pallet_attr<Attr>(item: &syn::ImplItemType) -> syn::Result<Option<Attr>>
-where
-	Attr: syn::parse::Parse,
-{
-	item.attrs.get(0).map(|a| syn::parse2(a.into_token_stream())).transpose()
+fn is_runtime_type(item: &syn::ImplItemType) -> bool {
+	item.attrs.iter().any(|attr| {
+		if let Ok(PalletAttr { typ: PalletAttrType::RuntimeType(_), .. }) =
+			parse2::<PalletAttr>(attr.into_token_stream())
+		{
+			return true
+		}
+		false
+	})
 }
 
 #[derive(Parse, Debug)]
@@ -132,10 +136,7 @@ fn combine_impls(
 				return None
 			}
 			if let ImplItem::Type(typ) = item.clone() {
-				let mut typ = typ.clone();
-				if let Ok(Some(PalletAttr { typ: PalletAttrType::RuntimeType(_), .. })) =
-					get_first_item_pallet_attr::<PalletAttr>(&mut typ)
-				{
+				if is_runtime_type(&typ) {
 					let item: ImplItem = if inject_runtime_types {
 						parse_quote! {
 							type #ident = #ident;
@@ -227,3 +228,25 @@ fn test_derive_impl_attr_args_parsing() {
 	assert!(parse2::<DeriveImplAttrArgs>(quote!()).is_err());
 	assert!(parse2::<DeriveImplAttrArgs>(quote!(Config Config)).is_err());
 }
+
+#[test]
+fn test_runtime_type_with_doc() {
+	trait TestTrait {
+		type Test;
+	}
+	#[allow(unused)]
+	struct TestStruct;
+	let p = parse2::<ItemImpl>(quote!(
+		impl TestTrait for TestStruct {
+			/// Some doc
+			#[inject_runtime_type]
+			type Test = u32;
+		}
+	))
+	.unwrap();
+	for item in p.items {
+		if let ImplItem::Type(typ) = item {
+			assert_eq!(is_runtime_type(&typ), true);
+		}
+	}
+}