frame_support_procedural/pallet/expand/
tasks.rs

1//! Contains logic for expanding task-related items.
2
3// This file is part of Substrate.
4
5// Copyright (C) Parity Technologies (UK) Ltd.
6// SPDX-License-Identifier: Apache-2.0
7
8// Licensed under the Apache License, Version 2.0 (the "License");
9// you may not use this file except in compliance with the License.
10// You may obtain a copy of the License at
11//
12// 	http://www.apache.org/licenses/LICENSE-2.0
13//
14// Unless required by applicable law or agreed to in writing, software
15// distributed under the License is distributed on an "AS IS" BASIS,
16// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17// See the License for the specific language governing permissions and
18// limitations under the License.
19
20//! Home of the expansion code for the Tasks API
21
22use crate::pallet::{parse::tasks::*, Def};
23use derive_syn_parse::Parse;
24use inflector::Inflector;
25use proc_macro2::TokenStream as TokenStream2;
26use quote::{format_ident, quote, ToTokens};
27use syn::{parse_quote, spanned::Spanned, ItemEnum, ItemImpl};
28
29impl TaskEnumDef {
30	/// Since we optionally allow users to manually specify a `#[pallet::task_enum]`, in the
31	/// event they _don't_ specify one (which is actually the most common behavior) we have to
32	/// generate one based on the existing [`TasksDef`]. This method performs that generation.
33	pub fn generate(
34		tasks: &TasksDef,
35		type_decl_bounded_generics: TokenStream2,
36		type_use_generics: TokenStream2,
37	) -> Self {
38		let variants = if tasks.tasks_attr.is_some() {
39			tasks
40				.tasks
41				.iter()
42				.map(|task| {
43					let ident = &task.item.sig.ident;
44					let ident =
45						format_ident!("{}", ident.to_string().to_class_case(), span = ident.span());
46
47					let args = task.item.sig.inputs.iter().collect::<Vec<_>>();
48
49					if args.is_empty() {
50						quote!(#ident)
51					} else {
52						quote!(#ident {
53							#(#args),*
54						})
55					}
56				})
57				.collect::<Vec<_>>()
58		} else {
59			Vec::new()
60		};
61		let mut task_enum_def: TaskEnumDef = parse_quote! {
62			/// Auto-generated enum that encapsulates all tasks defined by this pallet.
63			///
64			/// Conceptually similar to the [`Call`] enum, but for tasks. This is only
65			/// generated if there are tasks present in this pallet.
66			#[pallet::task_enum]
67			pub enum Task<#type_decl_bounded_generics> {
68				#(
69					#variants,
70				)*
71			}
72		};
73		task_enum_def.type_use_generics = type_use_generics;
74		task_enum_def
75	}
76}
77
78impl ToTokens for TaskEnumDef {
79	fn to_tokens(&self, tokens: &mut TokenStream2) {
80		let item_enum = &self.item_enum;
81		let ident = &item_enum.ident;
82		let vis = &item_enum.vis;
83		let attrs = &item_enum.attrs;
84		let generics = &item_enum.generics;
85		let variants = &item_enum.variants;
86		let scrate = &self.scrate;
87		let type_use_generics = &self.type_use_generics;
88		if self.attr.is_some() {
89			// `item_enum` is short-hand / generated enum
90			tokens.extend(quote! {
91				#(#attrs)*
92				#[derive(
93					#scrate::CloneNoBound,
94					#scrate::EqNoBound,
95					#scrate::PartialEqNoBound,
96					#scrate::pallet_prelude::Encode,
97					#scrate::pallet_prelude::Decode,
98					#scrate::pallet_prelude::TypeInfo,
99				)]
100				#[codec(encode_bound())]
101				#[codec(decode_bound())]
102				#[scale_info(skip_type_params(#type_use_generics))]
103				#vis enum #ident #generics {
104					#variants
105					#[doc(hidden)]
106					#[codec(skip)]
107					__Ignore(core::marker::PhantomData<T>, #scrate::Never),
108				}
109
110				impl<T: Config> core::fmt::Debug for #ident<#type_use_generics> {
111					fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
112						f.debug_struct(stringify!(#ident)).field("value", self).finish()
113					}
114				}
115			});
116		} else {
117			// `item_enum` is a manually specified enum (no attribute)
118			tokens.extend(item_enum.to_token_stream());
119		}
120	}
121}
122
123/// Represents an already-expanded [`TasksDef`].
124#[derive(Parse)]
125pub struct ExpandedTasksDef {
126	pub task_item_impl: ItemImpl,
127	pub task_trait_impl: ItemImpl,
128}
129
130impl ToTokens for TasksDef {
131	fn to_tokens(&self, tokens: &mut TokenStream2) {
132		let scrate = &self.scrate;
133		let enum_ident = syn::Ident::new("Task", self.enum_ident.span());
134		let enum_arguments = &self.enum_arguments;
135		let enum_use = quote!(#enum_ident #enum_arguments);
136
137		let task_fn_idents = self
138			.tasks
139			.iter()
140			.map(|task| {
141				format_ident!(
142					"{}",
143					&task.item.sig.ident.to_string().to_class_case(),
144					span = task.item.sig.ident.span()
145				)
146			})
147			.collect::<Vec<_>>();
148		let task_indices = self.tasks.iter().map(|task| &task.index_attr.meta.index);
149		let task_conditions = self.tasks.iter().map(|task| &task.condition_attr.meta.expr);
150		let task_weights = self.tasks.iter().map(|task| &task.weight_attr.meta.expr);
151		let task_iters = self.tasks.iter().map(|task| &task.list_attr.meta.expr);
152
153		let task_fn_impls = self.tasks.iter().map(|task| {
154			let mut task_fn_impl = task.item.clone();
155			task_fn_impl.attrs = vec![];
156			task_fn_impl
157		});
158
159		let task_fn_names = self.tasks.iter().map(|task| &task.item.sig.ident);
160		let task_arg_names = self.tasks.iter().map(|task| &task.arg_names).collect::<Vec<_>>();
161
162		let impl_generics = &self.item_impl.generics;
163		tokens.extend(quote! {
164			impl #impl_generics #enum_use
165			{
166				#(#task_fn_impls)*
167			}
168
169			impl #impl_generics #scrate::traits::Task for #enum_use
170			{
171				type Enumeration = #scrate::__private::IntoIter<#enum_use>;
172
173				fn iter() -> Self::Enumeration {
174					let mut all_tasks = #scrate::__private::vec![];
175					#(all_tasks
176						.extend(#task_iters.map(|(#(#task_arg_names),*)| #enum_ident::#task_fn_idents { #(#task_arg_names: #task_arg_names.clone()),* })
177						.collect::<#scrate::__private::Vec<_>>());
178					)*
179					all_tasks.into_iter()
180				}
181
182				fn task_index(&self) -> u32 {
183					match self.clone() {
184						#(#enum_ident::#task_fn_idents { .. } => #task_indices,)*
185						Task::__Ignore(_, _) => unreachable!(),
186					}
187				}
188
189				fn is_valid(&self) -> bool {
190					match self.clone() {
191						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => (#task_conditions)(#(#task_arg_names),* ),)*
192						Task::__Ignore(_, _) => unreachable!(),
193					}
194				}
195
196				fn run(&self) -> Result<(), #scrate::pallet_prelude::DispatchError> {
197					match self.clone() {
198						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => {
199							<#enum_use>::#task_fn_names(#( #task_arg_names, )* )
200						},)*
201						Task::__Ignore(_, _) => unreachable!(),
202					}
203				}
204
205				#[allow(unused_variables)]
206				fn weight(&self) -> #scrate::pallet_prelude::Weight {
207					match self.clone() {
208						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => #task_weights,)*
209						Task::__Ignore(_, _) => unreachable!(),
210					}
211				}
212			}
213		});
214	}
215}
216
217/// Expands the [`TasksDef`] in the enclosing [`Def`], if present, and returns its tokens.
218///
219/// This modifies the underlying [`Def`] in addition to returning any tokens that were added.
220pub fn expand_tasks_impl(def: &mut Def) -> TokenStream2 {
221	let Some(tasks) = &mut def.tasks else { return quote!() };
222	let ExpandedTasksDef { task_item_impl, task_trait_impl } = parse_quote!(#tasks);
223	quote! {
224		#task_item_impl
225		#task_trait_impl
226	}
227}
228
229/// Represents a fully-expanded [`TaskEnumDef`].
230#[derive(Parse)]
231pub struct ExpandedTaskEnum {
232	pub item_enum: ItemEnum,
233	pub debug_impl: ItemImpl,
234}
235
236/// Modifies a [`Def`] to expand the underlying [`TaskEnumDef`] if present, and also returns
237/// its tokens. A blank [`TokenStream2`] is returned if no [`TaskEnumDef`] has been generated
238/// or defined.
239pub fn expand_task_enum(def: &mut Def) -> TokenStream2 {
240	let Some(task_enum) = &mut def.task_enum else { return quote!() };
241	let ExpandedTaskEnum { item_enum, debug_impl } = parse_quote!(#task_enum);
242	quote! {
243		#item_enum
244		#debug_impl
245	}
246}
247
248/// Modifies a [`Def`] to expand the underlying [`TasksDef`] and also generate a
249/// [`TaskEnumDef`] if applicable. The tokens for these items are returned if they are created.
250pub fn expand_tasks(def: &mut Def) -> TokenStream2 {
251	if let Some(tasks_def) = &def.tasks {
252		if def.task_enum.is_none() {
253			def.task_enum = Some(TaskEnumDef::generate(
254				&tasks_def,
255				def.type_decl_bounded_generics(tasks_def.item_impl.span()),
256				def.type_use_generics(tasks_def.item_impl.span()),
257			));
258		}
259	}
260	let tasks_extra_output = expand_tasks_impl(def);
261	let task_enum_extra_output = expand_task_enum(def);
262	quote! {
263		#tasks_extra_output
264		#task_enum_extra_output
265	}
266}
frame_support_procedural/pallet/expand/tasks.rs

frame_support_procedural/pallet/expand/
tasks.rs
