[contracts] Add docs generator for the contracts API to the #[define_env] macro

frame/contracts/proc-macro/src/lib.rs

@@ -30,9 +30,12 @@ use alloc::{
 	vec::Vec,
 };
 use proc_macro::TokenStream;
-use proc_macro2::TokenStream as TokenStream2;
+use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::{quote, quote_spanned, ToTokens};
-use syn::{parse_macro_input, spanned::Spanned, Data, DeriveInput, FnArg, Ident};
+use syn::{
+	parse_macro_input, punctuated::Punctuated, spanned::Spanned, token::Comma, Data, DeriveInput,
+	FnArg, Ident,
+};
 
 /// This derives `Debug` for a struct where each field must be of some numeric type.
 /// It interprets each field as its represents some weight and formats it as times so that
@@ -383,16 +386,69 @@ fn is_valid_special_arg(idx: usize, arg: &FnArg) -> bool {
 	matches!(*pat.ty, syn::Type::Infer(_))
 }
 
+/// Expands documentation for host functions.
+fn expand_docs(def: &mut EnvDef) -> TokenStream2 {
+	let mut modules = def.host_funcs.iter().map(|f| f.module.clone()).collect::<Vec<_>>();
+	modules.sort();
+	modules.dedup();
+
+	let doc_selector = |a: &syn::Attribute| a.path.is_ident("doc");
+	let docs = modules.iter().map(|m| {
+		let funcs = def.host_funcs.iter_mut().map(|f| {
+			if *m == f.module {
+				// Remove auxiliary args: `ctx: _` and `memory: _`
+				f.item.sig.inputs = f
+					.item
+					.sig
+					.inputs
+					.iter()
+					.skip(2)
+					.map(|p| p.clone())
+					.collect::<Punctuated<FnArg, Comma>>();
+				let func_decl = f.item.sig.to_token_stream();
+				let func_docs = f.item.attrs.iter().filter(|a| doc_selector(a)).map(|d| {
+					let docs = d.to_token_stream();
+					quote! { #docs }
+				});
+				quote! {
+					#( #func_docs )*
+					#func_decl;
+				}
+			} else {
+				quote! {}
+			}
+		});
+
+		let module = Ident::new(m, Span::call_site());
+
+		quote! {
+			/// Documentation for the seal module host functions.
+			mod #module {
+				use crate::wasm::runtime::{TrapReason, ReturnCode};
+				/// Dumb trait for generating host functions docs.
+				trait Docs {
+					#( #funcs )*
+				}
+			}
+		}
+	});
+	quote! {
+		#( #docs )*
+	}
+}
+
 /// Expands environment definiton.
 /// Should generate source code for:
 ///  - implementations of the host functions to be added to the wasm runtime environment (see
 ///    `expand_impls()`).
-fn expand_env(def: &mut EnvDef) -> TokenStream2 {
+fn expand_env(def: &mut EnvDef, docs: bool) -> TokenStream2 {
 	let impls = expand_impls(def);
+	let docs = docs.then_some(expand_docs(def)).unwrap_or(TokenStream2::new());
 
 	quote! {
 		pub struct Env;
 		#impls
+		#docs
 	}
 }
 
@@ -579,18 +635,33 @@ fn expand_functions(
 ///
 /// The implementation on `()` can be used in places where no `Ext` exists, yet. This is useful
 /// when only checking whether a code can be instantiated without actually executing any code.
+///
+/// # Generating Documentation
+///
+/// Passing `doc` attribute to the macro (like `#[define_env(doc)]`) will make it also expand
+/// additional `pallet_contracts::wasm::runtime::seal0`, `pallet_contracts::wasm::runtime::seal1`,
+/// `...` modules each having its `Doc` trait containing methods holding documentation for every
+/// defined host function.
+///
+/// To build up these docs, run:
+///
+/// ```nocompile
+/// 	cargo doc --no-deps --document-private-items
+/// ```
 #[proc_macro_attribute]
 pub fn define_env(attr: TokenStream, item: TokenStream) -> TokenStream {
-	if !attr.is_empty() {
-		let msg = "Invalid `define_env` attribute macro: expected no attributes: `#[define_env]`.";
+	if !attr.is_empty() && !(attr.to_string() == "doc".to_string()) {
+		let msg = r#"Invalid `define_env` attribute macro: expected either no attributes or a single `doc` attibute:
+					 - `#[define_env]`
+					 - `#[define_env(doc)]`"#;
 		let span = TokenStream2::from(attr).span();
 		return syn::Error::new(span, msg).to_compile_error().into()
 	}
 
 	let item = syn::parse_macro_input!(item as syn::ItemMod);
 
 	match EnvDef::try_from(item) {
-		Ok(mut def) => expand_env(&mut def).into(),
+		Ok(mut def) => expand_env(&mut def, !attr.is_empty()).into(),
 		Err(e) => e.to_compile_error().into(),
 	}
 }