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
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
//! A collection of attribute macros to reduce boilerplate in the
//! [solders](https://github.com/kevinheavey/solders) project.
//!
//! These macros make some very specific assumptions about the structs
//! they're applied to, so they're unlikely to be useful for other projects.
use proc_macro::TokenStream;
use proc_macro2::Ident;
use quote::{quote, ToTokens};
use syn::{parse_macro_input, ImplItem, ItemEnum, ItemImpl};

/// Add a `__hash__` to the impl using the `PyHash` trait.
///
/// # Example
///
/// ```rust
/// use solders_macros::pyhash;
///
/// #[derive(Debug)]
/// struct Foo(u8);
///
/// #[pyhash]
/// impl Foo {
///   pub fn pyhash(&self) -> u64 {  // Fake implementation in place of `PyHash`.
///      self.0.into()
///   }
/// }
///
/// let foo = Foo(3);
/// assert_eq!(3, foo.__hash__());
///
/// ```
#[proc_macro_attribute]
pub fn pyhash(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let to_add = quote! {pub fn __hash__(&self) -> u64 {
        solders_traits_core::PyHash::pyhash(self)
    }};
    ast.items.push(ImplItem::Verbatim(to_add));
    TokenStream::from(ast.to_token_stream())
}

/// Add a `__richcmp__` to the impl using the `RichcmpFull` trait.
///
/// # Example
///
/// ```rust
/// use solders_macros::richcmp_full;
///
///
/// #[derive(Debug)]
/// struct Foo(u8);
///
/// #[richcmp_full]
/// impl Foo {
///   pub fn richcmp(&self, other: &Self, op: CompareOp) -> bool {  // Fake implementation in place of `RichcmpFull`.
///      true
///   }
/// }
///
/// let foo = Foo(3);
/// assert_eq!(true, foo.__richcmp__(&foo, CompareOp::Eq));
///
/// ```
#[proc_macro_attribute]
pub fn richcmp_full(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let to_add = quote! {pub fn __richcmp__(&self, other: &Self, op: pyo3::basic::CompareOp) -> bool {
        solders_traits_core::RichcmpFull::richcmp(self, other, op)
    }};
    ast.items.push(ImplItem::Verbatim(to_add));
    TokenStream::from(ast.to_token_stream())
}

/// Add a `__richcmp__` to the impl using the `RichcmpEqualityOnly` trait.
#[proc_macro_attribute]
pub fn richcmp_eq_only(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let to_add = quote! {pub fn __richcmp__(&self, other: &Self, op: pyo3::basic::CompareOp) -> pyo3::prelude::PyResult<bool> {
        solders_traits_core::RichcmpEqualityOnly::richcmp(self, other, op)
    }};
    ast.items.push(ImplItem::Verbatim(to_add));
    TokenStream::from(ast.to_token_stream())
}

/// Add a `__richcmp__` to the impl using the `RichcmpSigner` trait.
#[proc_macro_attribute]
pub fn richcmp_signer(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let to_add = quote! {pub fn __richcmp__(&self, other: crate::signer::Signer, op: pyo3::basic::CompareOp) -> pyo3::prelude::PyResult<bool> {
        solders_traits::RichcmpSigner::richcmp(self, other, op)
    }};
    ast.items.push(ImplItem::Verbatim(to_add));
    TokenStream::from(ast.to_token_stream())
}

fn add_core_methods(ast: &mut ItemImpl) {
    let mut methods = vec![
        ImplItem::Verbatim(
            quote! {pub fn __bytes__<'a>(&self, py: pyo3::prelude::Python<'a>) -> &'a pyo3::types::PyBytes  {
                solders_traits_core::CommonMethodsCore::pybytes(self, py)
            }},
        ),
        ImplItem::Verbatim(quote! { pub fn __str__(&self) -> String {
            solders_traits_core::CommonMethodsCore::pystr(self)
        } }),
        ImplItem::Verbatim(quote! { pub fn __repr__(&self) -> String {
            solders_traits_core::CommonMethodsCore::pyrepr(self)
        } }),
        ImplItem::Verbatim(
            quote! { pub fn __reduce__(&self) -> pyo3::prelude::PyResult<(pyo3::prelude::PyObject, pyo3::prelude::PyObject)> {
                solders_traits_core::CommonMethodsCore::pyreduce(self)
            } },
        ),
    ];
    if !ast.items.iter().any(|item| match item {
        ImplItem::Method(m) => m.sig.ident == "from_bytes",
        _ => false,
    }) {
        let from_bytes = ImplItem::Verbatim(quote! {
            /// Deserialize from bytes.
            ///
            /// Args:
            ///     data (bytes): the serialized object.
            ///
            /// Returns: the deserialized object.
            ///
            #[staticmethod]
            pub fn from_bytes(data: &[u8]) -> PyResult<Self> {
                <Self as solders_traits_core::CommonMethodsCore>::py_from_bytes(data)
            }
        });
        methods.push(from_bytes);
    };
    ast.items.extend_from_slice(&methods);
}

/// Add `__bytes__`, `__str__`, `__repr__` and `__reduce__` using the `CommonMethodsCore` trait.
///
/// Also add `from_bytes` if not already defined.
#[proc_macro_attribute]
pub fn common_methods_core(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    add_core_methods(&mut ast);
    TokenStream::from(ast.to_token_stream())
}

/// Add `__bytes__`, `__str__`, `__repr__` and `__reduce__`, `to_json` and `from_json` using the `CommonMethods` trait.
///
/// Also add `from_bytes` if not already defined.
#[proc_macro_attribute]
pub fn common_methods(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    add_core_methods(&mut ast);
    let methods = vec![
        ImplItem::Verbatim(quote! {
        /// Convert to a JSON string.
        pub fn to_json(&self) -> String {
            solders_traits_core::CommonMethods::py_to_json(self)
        } }),
        ImplItem::Verbatim(quote! {
        /// Build from a JSON string.
        #[staticmethod] pub fn from_json(raw: &str) -> PyResult<Self> {
            <Self as solders_traits_core::CommonMethods>::py_from_json(raw)
        } }),
    ];
    ast.items.extend_from_slice(&methods);
    TokenStream::from(ast.to_token_stream())
}

/// Add `__bytes__`, `__str__`, `__repr__`, `__reduce__`, `to_json`, `from_json`, `from_bytes` and `__richcmp__` using the `CommonMethodsRpcResp` trait.
#[proc_macro_attribute]
pub fn common_methods_rpc_resp(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let methods = vec![
        ImplItem::Verbatim(
            quote! {pub fn __bytes__<'a>(&self, py: pyo3::prelude::Python<'a>) -> &'a pyo3::types::PyBytes  {
                CommonMethodsRpcResp::pybytes(self, py)
            }},
        ),
        ImplItem::Verbatim(quote! { pub fn __str__(&self) -> String {
            CommonMethodsRpcResp::pystr(self)
        } }),
        ImplItem::Verbatim(quote! { pub fn __repr__(&self) -> String {
            CommonMethodsRpcResp::pyrepr(self)
        } }),
        ImplItem::Verbatim(
            quote! { pub fn __reduce__(&self) -> pyo3::prelude::PyResult<(pyo3::prelude::PyObject, pyo3::prelude::PyObject)> {
                CommonMethodsRpcResp::pyreduce(self)
            } },
        ),
        ImplItem::Verbatim(quote! {
        /// Convert to a JSON string.
        pub fn to_json(&self) -> String {
            CommonMethodsRpcResp::py_to_json(self)
        } }),
        ImplItem::Verbatim(quote! {
        /// Build from a JSON string.
        ///
        /// Args:
        ///     raw (str): The RPC JSON response (can be an error response).
        ///
        /// Returns:
        ///     Either the deserialized object or ``RPCError``.
        ///
        #[staticmethod]
        pub fn from_json(raw: &str) -> PyResult<Resp<Self>> {
            <Self as CommonMethodsRpcResp>::py_from_json(raw)
        } }),
        ImplItem::Verbatim(quote! {
            /// Deserialize from bytes.
            ///
            /// Args:
            ///     data (bytes): the serialized object.
            ///
            /// Returns: the deserialized object.
            ///
            #[staticmethod]
            pub fn from_bytes(data: &[u8]) -> PyResult<Self> {
                <Self as CommonMethodsRpcResp>::py_from_bytes(data)
            }
        }),
        ImplItem::Verbatim(
            quote! {pub fn __richcmp__(&self, other: &Self, op: pyo3::basic::CompareOp) -> pyo3::prelude::PyResult<bool> {
                solders_traits_core::RichcmpEqualityOnly::richcmp(self, other, op)
            }},
        ),
    ];
    ast.items.extend_from_slice(&methods);
    TokenStream::from(ast.to_token_stream())
}

/// Add an `id` getter to an RPC request object.
///
/// By convention, assumes the `id` lives at `self.base.id`.
#[proc_macro_attribute]
pub fn rpc_id_getter(_: TokenStream, item: TokenStream) -> TokenStream {
    let mut ast = parse_macro_input!(item as ItemImpl);
    let to_add = quote! {
    /// int: The ID of the RPC request.
    #[getter]
    pub fn id(&self) -> u64 {
        self.base.id
    }};
    ast.items.push(ImplItem::Verbatim(to_add));
    TokenStream::from(ast.to_token_stream())
}

/// Add mappings to and from another enum that has the exact same fields.
///
/// # Example
///
/// ```rust
/// use solders_macros::enum_original_mapping;
///
/// #[derive(PartialEq, Debug)]
/// pub enum Foo {
///   A,
///   B
/// }
/// #[enum_original_mapping(Foo)]
/// #[derive(PartialEq, Debug)]
/// pub enum Bar {
///   A,
///   B,
/// }
///
/// let a = Bar::A;
/// let b = Foo::B;
/// assert_eq!(Foo::from(a), Foo::A);
/// assert_eq!(Bar::from(b), Bar::B);
///
#[proc_macro_attribute]
pub fn enum_original_mapping(original: TokenStream, item: TokenStream) -> TokenStream {
    let mut new_stream = proc_macro2::TokenStream::from(item.clone());
    let ast = parse_macro_input!(item as ItemEnum);
    let enum_name = ast.ident;
    let orig = parse_macro_input!(original as Ident);
    let variant_names: Vec<Ident> = ast.variants.into_iter().map(|v| v.ident).collect();
    let from_impl = quote! {
        impl From<#orig> for #enum_name {
            fn from(left: #orig) -> Self {
                match left {
                    #(#orig::#variant_names => Self::#variant_names),*,
                    _ => panic!("Unrecognized variant: {:?}", left)
                }
            }
        }

        impl From<#enum_name> for #orig {
            fn from(left: #enum_name) -> Self {
                match left {
                    #(#enum_name::#variant_names => Self::#variant_names),*
                }
            }
        }
    };
    new_stream.extend(from_impl);
    TokenStream::from(new_stream)
}

/// Impl IntoPy<PyObject> for an ADT where each variant is a newtype.
///
/// # Example
///
/// ```rust
/// use solders_macros::EnumIntoPy;
///
/// #[derive(PartialEq, Debug, EnumIntoPy)]
/// pub enum Foo {
///   A(u8),
///   B(u8)
/// }
///
#[proc_macro_derive(EnumIntoPy)]
pub fn enum_into_py(item: TokenStream) -> TokenStream {
    let ast = parse_macro_input!(item as ItemEnum);
    let enum_name = ast.ident;
    let variant_names: Vec<Ident> = ast.variants.into_iter().map(|v| v.ident).collect();
    let into_py_impl = quote! {
        impl IntoPy<PyObject> for #enum_name {
            fn into_py(self, py: Python<'_>) -> PyObject {
                match self {
                    #(Self::#variant_names(x) => x.into_py(py)),*,
                }
            }
        }
    };
    into_py_impl.to_token_stream().into()
}