pendzl_lang_codegen/
implementation.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2012-2022 Supercolony
3//
4// Permission is hereby granted, free of charge, to any person obtaining
5// a copy of this software and associated documentation files (the"Software"),
6// to deal in the Software without restriction, including
7// without limitation the rights to use, copy, modify, merge, publish,
8// distribute, sublicense, and/or sell copies of the Software, and to
9// permit persons to whom the Software is furnished to do so, subject to
10// the following conditions:
11//
12// The above copyright notice and this permission notice shall be
13// included in all copies or substantial portions of the Software.
14//
15// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
16// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
17// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
18// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
19// LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
20// OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
21// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
22
23use 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    // map attribute args to default contract names
36    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    // name of struct for which we will implement the traits
57    let ident = extract_storage_struct_name(&items);
58    // we will look for overriden functions and remove them from the mod
59    let (map, mut items) = consume_overriders(items);
60
61    // to save importing of stuff by users
62    let mut imports = HashMap::<&str, syn::ItemUse>::default();
63    // if multiple contracts are using the same trait implemented differently we override it this way
64    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    // add the imports
96    impl_args.items.append(
97        &mut impl_args
98            .imports
99            .values()
100            .cloned()
101            .map(syn::Item::Use)
102            .collect(),
103    );
104
105    // add overriden traits
106    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
117//TODO verify
118fn cleanup_imports(imports: &mut HashMap<&str, syn::ItemUse>) {
119    // we will remove unnecessary imports
120    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
137// this method consumes override annotated methods and returns then mapped to code and the mod without them
138// we will later override the methods
139fn 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                // remove the overrider attribute
151                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}