1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
#![no_std]

extern crate alloc;

use alloc::format;
use alloc::string::ToString;

use proc_macro2::{Ident, Span, TokenStream};
use quote::quote;
use syn::{
    parse_macro_input, parse_quote, Data, DeriveInput, Expr, Fields, GenericParam, Generics, Lit,
};

#[proc_macro_derive(EpeeObject, attributes(epee_default, epee_alt_name, epee_flatten))]
pub fn derive_epee_object(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    // Parse the input tokens into a syntax tree.
    let input = parse_macro_input!(input as DeriveInput);

    let struct_name = input.ident;

    let generics = add_trait_bounds(input.generics);
    let (_impl_generics, _ty_generics, _where_clause) = generics.split_for_impl();

    let output = match input.data {
        Data::Struct(data) => build(&data.fields, &struct_name),
        _ => panic!("Only structs can be epee objects"),
    };

    output.into()
}

fn add_trait_bounds(mut generics: Generics) -> Generics {
    for param in &mut generics.params {
        if let GenericParam::Type(ref mut type_param) = *param {
            type_param
                .bounds
                .push(parse_quote!(epee_encoding::EpeeValue));
        }
    }
    generics
}

fn build(fields: &Fields, struct_name: &Ident) -> TokenStream {
    let mut struct_fields = TokenStream::new();
    let mut default_values = TokenStream::new();
    let mut count_fields = TokenStream::new();
    let mut write_fields = TokenStream::new();

    let mut read_match_body = TokenStream::new();
    let mut read_catch_all = TokenStream::new();

    let mut object_finish = TokenStream::new();

    let numb_o_fields: u64 = fields.len().try_into().unwrap();

    for field in fields {
        let field_name = field.ident.clone().expect("Epee only accepts named fields");
        let field_type = &field.ty;
        // If this field has a default value find it
        let default_val: Option<Expr> = field
            .attrs
            .iter()
            .find(|f| f.path().is_ident("epee_default"))
            .map(|f| f.parse_args().unwrap());
        // If this field has a different name when encoded find it
        let alt_name: Option<Lit> = field
            .attrs
            .iter()
            .find(|f| f.path().is_ident("epee_alt_name"))
            .map(|f| f.parse_args().unwrap());

        let is_flattened = field
            .attrs
            .iter()
            .any(|f| f.path().is_ident("epee_flatten"));

        // Gets this objects epee name, the name its encoded with
        let epee_name = if let Some(alt) = alt_name {
            if is_flattened {
                panic!("Cant rename a flattened field")
            }
            match alt {
                Lit::Str(name) => name.value(),
                _ => panic!("Alt name was not a string"),
            }
        } else {
            field_name.to_string()
        };

        // This is fields part of a struct:
        // struct T {
        //  #struct_fields
        // }
        if is_flattened {
            struct_fields = quote! {
                #struct_fields
                #field_name: <#field_type as epee_encoding::EpeeObject>::Builder,
            };

            count_fields = quote! {
                #count_fields
                // This filed has been flattened so dont count it.
                numb_o_fields -= 1;
                // Add the flattend fields to this one.
                numb_o_fields += self.#field_name.number_of_fields();

            };
        } else {
            struct_fields = quote! {
                #struct_fields
                #field_name: Option<#field_type>,
            };
        }

        // `default_val`: this is the body of a default function:
        // fn default() -> Self {
        //    Self {
        //       #default_values
        //    }
        // }

        // `count_fields`: this is the part of the write function that takes
        // away from the number of fields if the field is the default value.

        // `write_fields`: this is the part of the write function that writes
        // this specific epee field.
        if let Some(default_val) = default_val {
            if is_flattened {
                panic!("Cant have a default on a flattened field");
            };

            default_values = quote! {
                #default_values
                #field_name: Some(#default_val),
            };

            count_fields = quote! {
                #count_fields
                if self.#field_name == #default_val {
                    numb_o_fields -= 1;
                };
            };

            write_fields = quote! {
                #write_fields
                if self.#field_name != #default_val {
                    epee_encoding::write_field(&self.#field_name, &#epee_name, w)?;
                }
            }
        } else if !is_flattened {
            default_values = quote! {
                #default_values
                #field_name: None,
            };

            write_fields = quote! {
                #write_fields
                epee_encoding::write_field(&self.#field_name, #epee_name, w)?;
            };
        } else {
            default_values = quote! {
                #default_values
                #field_name: Default::default(),
            };

            write_fields = quote! {
                #write_fields
                self.#field_name.write_fields(w)?;
            };
        };

        // This is what these values do:
        // fn add_field(name: &str, r: &mut r) -> Result<bool> {
        //    match name {
        //        #read_match_body
        //        _ => {
        //           #read_catch_all
        //           return Ok(false);
        //         }
        //    }
        //    Ok(true)
        // }
        if is_flattened {
            read_catch_all = quote! {
                #read_catch_all
                if self.#field_name.add_field(name, r)? {
                    return Ok(true);
                };
            };

            object_finish = quote! {
                #object_finish
                #field_name: self.#field_name.finish()?,
            };
        } else {
            read_match_body = quote! {
                #read_match_body
                #epee_name => {self.#field_name = Some(epee_encoding::read_epee_value(r)?);},
            };

            object_finish = quote! {
                #object_finish
                #field_name: self.#field_name.ok_or_else(|| epee_encoding::error::Error::Format("Required field was not found!"))?,
            };
        }
    }

    let builder_name = Ident::new(&format!("__{}EpeeBuilder", struct_name), Span::call_site());
    let mod_name = Ident::new(&format!("__{}_epee_module", struct_name), Span::call_site());

    let builder_impl = quote! {
        pub struct #builder_name {
            #struct_fields
        }

        impl Default for #builder_name {
            fn default() -> Self {
                Self {
                    #default_values
                }
            }
        }

        impl epee_encoding::EpeeObjectBuilder<#struct_name> for #builder_name {
            fn add_field<R: epee_encoding::io::Read>(&mut self, name: &str, r: &mut R) -> epee_encoding::error::Result<bool> {
                match name {
                    #read_match_body
                    _ => {
                        #read_catch_all
                        return Ok(false);
                    }
                };

                Ok(true)
            }

            fn finish(self) -> epee_encoding::error::Result<#struct_name> {
                Ok(#struct_name {
                    #object_finish
                })
            }
        }
    };

    let object_impl = quote! {
        impl EpeeObject for #struct_name {
            type Builder = #mod_name::#builder_name;

            fn number_of_fields(&self) -> u64 {
                let mut numb_o_fields: u64 = #numb_o_fields;
                #count_fields
                numb_o_fields
            }


            fn write_fields<W: epee_encoding::io::Write>(&self, w: &mut W) -> epee_encoding::error::Result<()> {

                #write_fields

                Ok(())
            }
        }
    };

    quote! {
        mod #mod_name {
            use super::*;
            #builder_impl
        }

        #object_impl
    }
}