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 inflector::Inflector;
24use proc_macro2::TokenStream as TokenStream2;
25use quote::{format_ident, quote, ToTokens};
26use syn::{parse_quote_spanned, spanned::Spanned};
27
28impl TaskEnumDef {
29	/// Since we optionally allow users to manually specify a `#[pallet::task_enum]`, in the
30	/// event they _don't_ specify one (which is actually the most common behavior) we have to
31	/// generate one based on the existing [`TasksDef`]. This method performs that generation.
32	pub fn generate(tasks: &TasksDef, def: &Def) -> Self {
33		// We use the span of the attribute to indicate that the error comes from code generated
34		// for the specific section, otherwise the item impl.
35		let span = tasks
36			.tasks_attr
37			.as_ref()
38			.map_or_else(|| tasks.item_impl.span(), |attr| attr.span());
39
40		let type_decl_bounded_generics = def.type_decl_bounded_generics(span);
41
42		let variants = if tasks.tasks_attr.is_some() {
43			tasks
44				.tasks
45				.iter()
46				.map(|task| {
47					let ident = &task.item.sig.ident;
48					let ident =
49						format_ident!("{}", ident.to_string().to_class_case(), span = ident.span());
50
51					let args = task.item.sig.inputs.iter().collect::<Vec<_>>();
52
53					if args.is_empty() {
54						quote!(#ident)
55					} else {
56						quote!(#ident {
57							#(#args),*
58						})
59					}
60				})
61				.collect::<Vec<_>>()
62		} else {
63			Vec::new()
64		};
65
66		parse_quote_spanned! { span =>
67			/// Auto-generated enum that encapsulates all tasks defined by this pallet.
68			///
69			/// Conceptually similar to the [`Call`] enum, but for tasks. This is only
70			/// generated if there are tasks present in this pallet.
71			#[pallet::task_enum]
72			pub enum Task<#type_decl_bounded_generics> {
73				#(
74					#variants,
75				)*
76			}
77		}
78	}
79}
80
81impl TaskEnumDef {
82	fn expand_to_tokens(&self, def: &Def) -> TokenStream2 {
83		if let Some(attr) = &self.attr {
84			let ident = &self.item_enum.ident;
85			let vis = &self.item_enum.vis;
86			let attrs = &self.item_enum.attrs;
87			let generics = &self.item_enum.generics;
88			let variants = &self.item_enum.variants;
89			let frame_support = &def.frame_support;
90			let type_use_generics = &def.type_use_generics(attr.span());
91			let type_impl_generics = &def.type_impl_generics(attr.span());
92
93			// `item_enum` is short-hand / generated enum
94			quote! {
95				#(#attrs)*
96				#[derive(
97					#frame_support::CloneNoBound,
98					#frame_support::EqNoBound,
99					#frame_support::PartialEqNoBound,
100					#frame_support::pallet_prelude::Encode,
101					#frame_support::pallet_prelude::Decode,
102					#frame_support::pallet_prelude::DecodeWithMemTracking,
103					#frame_support::pallet_prelude::TypeInfo,
104				)]
105				#[codec(encode_bound())]
106				#[codec(decode_bound())]
107				#[scale_info(skip_type_params(#type_use_generics))]
108				#vis enum #ident #generics {
109					#variants
110					#[doc(hidden)]
111					#[codec(skip)]
112					__Ignore(core::marker::PhantomData<(#type_use_generics)>, #frame_support::Never),
113				}
114
115				impl<#type_impl_generics> core::fmt::Debug for #ident<#type_use_generics> {
116					fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
117						f.debug_struct(stringify!(#ident)).field("value", self).finish()
118					}
119				}
120			}
121		} else {
122			// `item_enum` is a manually specified enum (no attribute)
123			self.item_enum.to_token_stream()
124		}
125	}
126}
127
128impl TasksDef {
129	fn expand_to_tokens(&self, def: &Def) -> TokenStream2 {
130		let frame_support = &def.frame_support;
131		let enum_ident = syn::Ident::new("Task", self.enum_ident.span());
132		let enum_arguments = &self.enum_arguments;
133		let enum_use = quote!(#enum_ident #enum_arguments);
134
135		let task_fn_idents = self
136			.tasks
137			.iter()
138			.map(|task| {
139				format_ident!(
140					"{}",
141					&task.item.sig.ident.to_string().to_class_case(),
142					span = task.item.sig.ident.span()
143				)
144			})
145			.collect::<Vec<_>>();
146		let task_indices = self.tasks.iter().map(|task| &task.index_attr.meta.index);
147		let task_conditions = self.tasks.iter().map(|task| &task.condition_attr.meta.expr);
148		let task_weights = self.tasks.iter().map(|task| &task.weight_attr.meta.expr);
149		let task_iters = self.tasks.iter().map(|task| &task.list_attr.meta.expr);
150
151		let task_fn_impls = self.tasks.iter().map(|task| {
152			let mut task_fn_impl = task.item.clone();
153			task_fn_impl.attrs = vec![];
154			task_fn_impl
155		});
156
157		let task_fn_names = self.tasks.iter().map(|task| &task.item.sig.ident);
158		let task_arg_names = self.tasks.iter().map(|task| &task.arg_names).collect::<Vec<_>>();
159
160		let impl_generics = &self.item_impl.generics;
161		quote! {
162			impl #impl_generics #enum_use
163			{
164				#(#task_fn_impls)*
165			}
166
167			impl #impl_generics #frame_support::traits::Task for #enum_use
168			{
169				type Enumeration = #frame_support::__private::IntoIter<#enum_use>;
170
171				fn iter() -> Self::Enumeration {
172					let mut all_tasks = #frame_support::__private::vec![];
173					#(all_tasks
174						.extend(#task_iters.map(|(#(#task_arg_names),*)| #enum_ident::#task_fn_idents { #(#task_arg_names: #task_arg_names.clone()),* })
175						.collect::<#frame_support::__private::Vec<_>>());
176					)*
177					all_tasks.into_iter()
178				}
179
180				fn task_index(&self) -> u32 {
181					match self.clone() {
182						#(#enum_ident::#task_fn_idents { .. } => #task_indices,)*
183						Task::__Ignore(_, _) => unreachable!(),
184					}
185				}
186
187				fn is_valid(&self) -> bool {
188					match self.clone() {
189						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => (#task_conditions)(#(#task_arg_names),* ),)*
190						Task::__Ignore(_, _) => unreachable!(),
191					}
192				}
193
194				fn run(&self) -> Result<(), #frame_support::pallet_prelude::DispatchError> {
195					match self.clone() {
196						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => {
197							<#enum_use>::#task_fn_names(#( #task_arg_names, )* )
198						},)*
199						Task::__Ignore(_, _) => unreachable!(),
200					}
201				}
202
203				#[allow(unused_variables)]
204				fn weight(&self) -> #frame_support::pallet_prelude::Weight {
205					match self.clone() {
206						#(#enum_ident::#task_fn_idents { #(#task_arg_names),* } => #task_weights,)*
207						Task::__Ignore(_, _) => unreachable!(),
208					}
209				}
210			}
211		}
212	}
213}
214
215/// Generate code related to tasks.
216pub fn expand_tasks(def: &Def) -> TokenStream2 {
217	let Some(tasks_def) = &def.tasks else {
218		return quote!();
219	};
220
221	let default_task_enum = TaskEnumDef::generate(&tasks_def, def);
222
223	let task_enum = def.task_enum.as_ref().unwrap_or_else(|| &default_task_enum);
224
225	let tasks_expansion = tasks_def.expand_to_tokens(def);
226	let task_enum_expansion = task_enum.expand_to_tokens(def);
227
228	quote! {
229		#tasks_expansion
230		#task_enum_expansion
231	}
232}
frame_support_procedural/pallet/expand/tasks.rs

frame_support_procedural/pallet/expand/
tasks.rs
