Update MaxEncodedLen derive macro (#512)

pgherveou · bkchr · web-flow · commit 1e3f80c1e047 · 2023-09-08T11:02:20.000+02:00
Fix the generated code for MaxEncoded derive trait

---------

Co-authored-by: Bastian Köcher &lt;info@kchr.de&gt;
diff --git a/derive/src/max_encoded_len.rs b/derive/src/max_encoded_len.rs
@@ -96,7 +96,9 @@ fn fields_length_expr(fields: &Fields, crate_path: &syn::Path) -> proc_macro2::T
 		let ty = &field.ty;
 		if utils::is_compact(&field) {
 			quote_spanned! {
-				ty.span() => .saturating_add(<#crate_path::Compact::<#ty> as #crate_path::MaxEncodedLen>::max_encoded_len())
+				ty.span() => .saturating_add(
+					<<#ty as #crate_path::HasCompact>::Type as #crate_path::MaxEncodedLen>::max_encoded_len()
+				)
 			}
 		} else {
 			quote_spanned! {
diff --git a/src/compact.rs b/src/compact.rs
@@ -20,6 +20,8 @@ use crate::alloc::vec::Vec;
 use crate::codec::{Encode, Decode, Input, Output, EncodeAsRef};
 use crate::encode_like::EncodeLike;
 use crate::Error;
+#[cfg(feature = "max-encoded-len")]
+use crate::MaxEncodedLen;
 #[cfg(feature = "fuzz")]
 use arbitrary::Arbitrary;
 
@@ -206,18 +208,39 @@ impl<'de, T> serde::Deserialize<'de> for Compact<T> where T: serde::Deserialize<
 	}
 }
 
+/// Requires the presence of `MaxEncodedLen` when the `max-encoded-len` feature is active.
+// Remove this trait when the feature is removed.
+#[cfg(feature = "max-encoded-len")]
+pub trait MaybeMaxEncodedLen: MaxEncodedLen {}
+#[cfg(feature = "max-encoded-len")]
+impl<T: MaxEncodedLen> MaybeMaxEncodedLen for T {}
+
+/// Requires the presence of `MaxEncodedLen` when the `max-encoded-len` feature is active.
+// Remove this trait when the feature is removed.
+#[cfg(not(feature = "max-encoded-len"))]
+pub trait MaybeMaxEncodedLen {}
+#[cfg(not(feature = "max-encoded-len"))]
+impl<T> MaybeMaxEncodedLen for T {}
+
 /// Trait that tells you if a given type can be encoded/decoded in a compact way.
 pub trait HasCompact: Sized {
 	/// The compact type; this can be
-	type Type: for<'a> EncodeAsRef<'a, Self> + Decode + From<Self> + Into<Self>;
+	type Type: for<'a> EncodeAsRef<'a, Self> + Decode + From<Self> + Into<Self> + MaybeMaxEncodedLen;
 }
 
 impl<'a, T: 'a> EncodeAsRef<'a, T> for Compact<T> where CompactRef<'a, T>: Encode + From<&'a T> {
 	type RefType = CompactRef<'a, T>;
 }
 
+#[cfg(feature = "max-encoded-len")]
+impl<T> MaxEncodedLen for Compact<T> where T: CompactAs, Compact<T::As>: MaxEncodedLen, Compact<T>: Encode {
+	fn max_encoded_len() -> usize {
+		Compact::<T::As>::max_encoded_len()
+	}
+}
+
 impl<T: 'static> HasCompact for T where
-	Compact<T>: for<'a> EncodeAsRef<'a, T> + Decode + From<Self> + Into<Self>
+	Compact<T>: for<'a> EncodeAsRef<'a, T> + Decode + From<Self> + Into<Self> + MaybeMaxEncodedLen
 {
 	type Type = Compact<T>;
 }
diff --git a/src/max_encoded_len.rs b/src/max_encoded_len.rs
@@ -46,9 +46,8 @@ macro_rules! impl_primitives {
 	};
 }
 
-impl_primitives!(u8, u16, u32, u64, u128, i8, i16, i32, i64, i128, bool);
-
 impl_primitives!(
+	u8, i8, u16, i16, u32, i32, u64, i64, u128, i128, bool,
 	NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU128, NonZeroI8, NonZeroI16, NonZeroI32,
 	NonZeroI64, NonZeroI128
 );
diff --git a/tests/max_encoded_len.rs b/tests/max_encoded_len.rs
@@ -73,12 +73,29 @@ struct CompactField {
 
 #[test]
 fn compact_field_max_length() {
+	assert_eq!(CompactField::max_encoded_len(), 17);
 	assert_eq!(
 		CompactField::max_encoded_len(),
 		Compact::<u64>::max_encoded_len() + u64::max_encoded_len()
 	);
 }
 
+
+#[derive(Encode, MaxEncodedLen)]
+struct CompactFieldGenerics<T: MaxEncodedLen> {
+	#[codec(compact)]
+	t: T,
+	v: u64,
+}
+
+#[test]
+fn compact_field_generics_max_length() {
+	assert_eq!(
+		CompactFieldGenerics::<u64>::max_encoded_len(),
+		CompactField::max_encoded_len()
+	);
+}
+
 #[derive(Encode, MaxEncodedLen)]
 struct CompactStruct(#[codec(compact)] u64);
 
