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

// Copyright (C) 2021 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![cfg(feature = "max-encoded-len")]

use crate::{
	trait_bounds,
	utils::{codec_crate_path, custom_mel_trait_bound, has_dumb_trait_bound},
};
use quote::{quote, quote_spanned};
use syn::{parse_quote, spanned::Spanned, Data, DeriveInput, Fields, Type};

/// impl for `#[derive(MaxEncodedLen)]`
pub fn derive_max_encoded_len(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	let mut input: DeriveInput = match syn::parse(input) {
		Ok(input) => input,
		Err(e) => return e.to_compile_error().into(),
	};

	let crate_path = match codec_crate_path(&input.attrs) {
		Ok(crate_path) => crate_path,
		Err(error) => return error.into_compile_error().into(),
	};

	let name = &input.ident;
	if let Err(e) = trait_bounds::add(
		&input.ident,
		&mut input.generics,
		&input.data,
		custom_mel_trait_bound(&input.attrs),
		parse_quote!(#crate_path::MaxEncodedLen),
		None,
		has_dumb_trait_bound(&input.attrs),
		&crate_path
	) {
		return e.to_compile_error().into()
	}
	let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();

	let data_expr = data_length_expr(&input.data);

	quote::quote!(
		const _: () = {
			impl #impl_generics #crate_path::MaxEncodedLen for #name #ty_generics #where_clause {
				fn max_encoded_len() -> ::core::primitive::usize {
					#data_expr
				}
			}
		};
	)
	.into()
}

/// generate an expression to sum up the max encoded length from several fields
fn fields_length_expr(fields: &Fields) -> proc_macro2::TokenStream {
	let type_iter: Box<dyn Iterator<Item = &Type>> = match fields {
		Fields::Named(ref fields) => Box::new(fields.named.iter().map(|field| &field.ty)),
		Fields::Unnamed(ref fields) => Box::new(fields.unnamed.iter().map(|field| &field.ty)),
		Fields::Unit => Box::new(std::iter::empty()),
	};
	// expands to an expression like
	//
	//   0
	//     .saturating_add(<type of first field>::max_encoded_len())
	//     .saturating_add(<type of second field>::max_encoded_len())
	//
	// We match the span of each field to the span of the corresponding
	// `max_encoded_len` call. This way, if one field's type doesn't implement
	// `MaxEncodedLen`, the compiler's error message will underline which field
	// caused the issue.
	let expansion = type_iter.map(|ty| {
		quote_spanned! {
			ty.span() => .saturating_add(<#ty>::max_encoded_len())
		}
	});
	quote! {
		0_usize #( #expansion )*
	}
}

// generate an expression to sum up the max encoded length of each field
fn data_length_expr(data: &Data) -> proc_macro2::TokenStream {
	match *data {
		Data::Struct(ref data) => fields_length_expr(&data.fields),
		Data::Enum(ref data) => {
			// We need an expression expanded for each variant like
			//
			//   0
			//     .max(<variant expression>)
			//     .max(<variant expression>)
			//     .saturating_add(1)
			//
			// The 1 derives from the discriminant; see
			// https://github.com/paritytech/parity-scale-codec/
			//   blob/f0341dabb01aa9ff0548558abb6dcc5c31c669a1/derive/src/encode.rs#L211-L216
			//
			// Each variant expression's sum is computed the way an equivalent struct's would be.

			let expansion = data.variants.iter().map(|variant| {
				let variant_expression = fields_length_expr(&variant.fields);
				quote! {
					.max(#variant_expression)
				}
			});

			quote! {
				0_usize #( #expansion )* .saturating_add(1)
			}
		},
		Data::Union(ref data) => {
			// https://github.com/paritytech/parity-scale-codec/
			//   blob/f0341dabb01aa9ff0548558abb6dcc5c31c669a1/derive/src/encode.rs#L290-L293
			syn::Error::new(data.union_token.span(), "Union types are not supported.")
				.to_compile_error()
		},
	}
}