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

// 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::utils::{self, codec_crate_path, custom_mel_trait_bound};
use quote::{quote, quote_spanned};
use syn::{
	Data, DeriveInput, Fields, GenericParam, Generics, TraitBound, Type, TypeParamBound,
	parse_quote, spanned::Spanned,
};

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

	let mel_trait = match max_encoded_len_trait(&input) {
		Ok(mel_trait) => mel_trait,
		Err(e) => return e.to_compile_error().into(),
	};

	let name = &input.ident;
	let generics = if let Some(custom_bound) = custom_mel_trait_bound(&input.attrs) {
		add_custom_trait_bounds(input.generics, custom_bound)
	} else {
		add_trait_bounds(input.generics, mel_trait.clone())
	};
	let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

	let data_expr = data_length_expr(&input.data);

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

fn max_encoded_len_trait(input: &DeriveInput) -> syn::Result<TraitBound> {
	let mel = codec_crate_path(&input.attrs)?;
	Ok(parse_quote!(#mel::MaxEncodedLen))
}

// Add a bound `T: MaxEncodedLen` to every type parameter T.
fn add_trait_bounds(mut generics: Generics, mel_trait: TraitBound) -> Generics {
	for param in &mut generics.params {
		if let GenericParam::Type(ref mut type_param) = *param {
			type_param.bounds.push(TypeParamBound::Trait(mel_trait.clone()));
		}
	}
	generics
}

// Add custom trait bounds to the type parameters as specified by the user.
fn add_custom_trait_bounds(mut generics: Generics, custom_bound: utils::TraitBounds) -> Generics {
	generics.make_where_clause().predicates.extend(custom_bound);
	generics
}

/// 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()
		}
	}
}