pendzl_lang_codegen/
implementation.rs1use crate::{implementations::*, internal, internal::*};
24use proc_macro2::TokenStream;
25use quote::{quote, ToTokens};
26use std::collections::HashMap;
27use syn::{Item, Path};
28
29pub fn generate(attrs: TokenStream, ink_module: TokenStream) -> TokenStream {
30 if internal::skip() {
31 return quote! {};
32 }
33 let input: TokenStream = ink_module;
34
35 let args = syn::parse2::<AttributeArgs>(attrs)
37 .expect("No default contracts to implement provided")
38 .iter()
39 .map(|arg| match arg {
40 NestedMeta::Path(method) => method.to_token_stream().to_string().replace(' ', ""),
41 _ => panic!("Expected names of pendzl traits to implement in the contract!"),
42 })
43 .collect::<Vec<String>>();
44
45 let mut module = syn::parse2::<syn::ItemMod>(input).expect("Can't parse contract module");
46 let (braces, items) = match module.clone().content {
47 Some((brace, items)) => (brace, items),
48 None => {
49 panic!(
50 "{}",
51 "out-of-line pendzl modules are not supported, use `#[implementation] mod name {{ ... }}`",
52 )
53 }
54 };
55
56 let ident = extract_storage_struct_name(&items);
58 let (map, mut items) = consume_overriders(items);
60
61 let mut imports = HashMap::<&str, syn::ItemUse>::default();
63 let mut overriden_traits = HashMap::<&str, syn::Item>::default();
65
66 let mut impl_args = ImplArgs::new(&map, &mut items, &mut imports, &mut overriden_traits, ident);
67
68 for to_default_implement in &args {
69 match to_default_implement.as_str() {
70 "PSP22" => impl_psp22(&mut impl_args),
71 "PSP22Burnable" => impl_psp22_burnable(&mut impl_args),
72 "PSP22Mintable" => impl_psp22_mintable(&mut impl_args),
73 "PSP22Vault" => impl_psp22_vault(&mut impl_args),
74 "PSP22Metadata" => impl_psp22_metadata(&mut impl_args),
75 "PSP34" => impl_psp34(&mut impl_args),
76 "PSP34Burnable" => impl_psp34_burnable(&mut impl_args),
77 "PSP34Metadata" => impl_psp34_metadata(&mut impl_args),
78 "PSP34Mintable" => impl_psp34_mintable(&mut impl_args),
79 "Ownable" => impl_ownable(&mut impl_args),
80 "AccessControl" => impl_access_control(&mut impl_args),
81 "Pausable" => impl_pausable(&mut impl_args),
82 "Vesting" => impl_vesting(&mut impl_args),
83 _ => panic!("pendzl::implementation({to_default_implement}) not implemented!"),
84 }
85 }
86
87 cleanup_imports(impl_args.imports);
88
89 let import_storage = syn::parse2::<syn::ItemUse>(quote!(
90 use pendzl::traits::StorageFieldGetter;
91 ))
92 .expect("Should parse import");
93
94 impl_args.imports.insert("PendzlStorage", import_storage);
95 impl_args.items.append(
97 &mut impl_args
98 .imports
99 .values()
100 .cloned()
101 .map(syn::Item::Use)
102 .collect(),
103 );
104
105 impl_args
107 .items
108 .append(&mut impl_args.overriden_traits.values().cloned().collect());
109
110 module.content = Some((braces, items));
111
112 quote! {
113 #module
114 }
115}
116
117fn cleanup_imports(imports: &mut HashMap<&str, syn::ItemUse>) {
119 let psp22_default_impls = vec!["PSP22Mintable", "PSP22Burnable", "PSP22Metadata"];
121 check_and_remove_import("PSP22", psp22_default_impls, imports);
122
123 let psp34_default_impls = vec!["PSP34Mintable", "PSP34Burnable", "PSP34Metadata"];
124 check_and_remove_import("PSP34", psp34_default_impls, imports);
125}
126
127fn check_and_remove_import(
128 name_to_check: &str,
129 to_check: Vec<&str>,
130 imports: &mut HashMap<&str, syn::ItemUse>,
131) {
132 if to_check.iter().any(|name| imports.contains_key(name)) {
133 imports.remove(name_to_check);
134 }
135}
136
137fn consume_overriders(items: Vec<syn::Item>) -> (OverridenFnMap, Vec<syn::Item>) {
140 let mut map = HashMap::new();
141 let mut result: Vec<syn::Item> = vec![];
142 items.into_iter().for_each(|mut item| {
143 if let Item::Fn(item_fn) = &mut item {
144 if is_attr(&item_fn.attrs, "overrider") {
145 let attr_name = "overrider";
146 let fn_name = item_fn.sig.ident.to_string();
147 let code = item_fn.block.clone();
148 let mut attributes = item_fn.attrs.clone();
149 let inputs = item_fn.sig.inputs.clone();
150 let to_remove_idx = attributes
152 .iter()
153 .position(|attr| is_attr(&[attr.clone()], attr_name))
154 .expect("No {attr_name} attribute found!");
155 let overrider_attribute = attributes.remove(to_remove_idx);
156
157 let trait_name = overrider_attribute
158 .parse_args::<Path>()
159 .expect("Expected overriden trait identifier")
160 .to_token_stream()
161 .to_string();
162
163 let mut vec = map.get(&trait_name).unwrap_or(&vec![]).clone();
164 vec.push((fn_name, (code, attributes, inputs)));
165 map.insert(trait_name, vec.to_vec());
166 } else {
167 result.push(item);
168 }
169 } else {
170 result.push(item);
171 }
172 });
173
174 (map, result)
175}
176
177pub fn extract_storage_struct_name(items: &[syn::Item]) -> String {
178 let contract_storage_struct = items
179 .iter()
180 .find(|item| {
181 if let Item::Struct(structure) = item {
182 let ink_attr_maybe = structure
183 .attrs
184 .iter()
185 .cloned()
186 .find(|attr| is_attr(&[attr.clone()], "ink"));
187
188 if let Some(ink_attr) = ink_attr_maybe {
189 if let Ok(path) = ink_attr.parse_args::<Path>() {
190 return path.to_token_stream().to_string() == "storage";
191 }
192 }
193 false
194 } else {
195 false
196 }
197 })
198 .expect("Contract storage struct not found!");
199 match contract_storage_struct {
200 Item::Struct(structure) => structure.ident.to_string(),
201 _ => unreachable!("Only Item::Struct allowed here"),
202 }
203}