Derive for enums.

substrate/codec/derive/src/decode.rs

@@ -16,11 +16,11 @@
 
 use proc_macro2::{Span, TokenStream, Ident};
 use syn::{
-	Data, Fields,
+	Data, Fields, Index,
 	spanned::Spanned,
 };
 
-pub fn quote(data: &Data, name: &Ident, input_: &TokenStream) -> TokenStream {
+pub fn quote(data: &Data, type_name: &Ident, input_: &TokenStream) -> TokenStream {
 	let call_site = Span::call_site();
 	match *data {
 		Data::Struct(ref data) => match data.fields {
@@ -34,7 +34,7 @@ pub fn quote(data: &Data, name: &Ident, input_: &TokenStream) -> TokenStream {
 					}
 				});
 
-				quote! {
+				quote_spanned! {call_site =>
 					Some(Self {
 						#( #recurse, )*
 					})
@@ -47,20 +47,81 @@ pub fn quote(data: &Data, name: &Ident, input_: &TokenStream) -> TokenStream {
 					}
 				});
 
-				quote! {
-					Some(#name(
+				quote_spanned! {call_site =>
+					Some(#type_name (
 						#( #recurse, )*
 					))
 				}
 			},
 			Fields::Unit => {
-				quote! {
+				quote_spanned! {call_site =>
 					drop(#input_);
-					Some(#name)
+					Some(#type_name)
 				}
 			},
 		},
-		Data::Enum(_) | Data::Union(_) => unimplemented!(),
+		Data::Enum(ref data) => {
+			assert!(data.variants.len() < 256, "Currently only enums with less than 255 variants are encodable.");
+
+			let recurse = data.variants.iter().enumerate().map(|(i, v)| {
+				let name = &v.ident;
+				let index = Index { index: i as u32, span: call_site };
+
+				let create = match v.fields {
+					Fields::Named(ref fields) => {
+						let recurse = fields.named.iter().map(|f| {
+							let name = &f.ident;
+							let field = quote_spanned!(call_site => #name);
+
+							quote_spanned! { f.span() =>
+								#field: ::codec::Decode::decode(#input_)?
+							}
+						});
+
+						quote_spanned! {call_site =>
+							Some(#type_name :: #name {
+								#( #recurse, )*
+							})
+						}
+					},
+					Fields::Unnamed(ref fields) => {
+						let recurse = fields.unnamed.iter().map(|f| {
+							quote_spanned! { f.span() =>
+								::codec::Decode::decode(#input_)?
+							}
+						});
+
+						quote_spanned! {call_site =>
+							Some(#type_name :: #name (
+								#( #recurse, )*
+							))
+						}
+					},
+					Fields::Unit => {
+						quote_spanned! {call_site =>
+							Some(#type_name :: #name)
+						}
+					},
+				};
+
+				quote_spanned! { v.span() =>
+					x if x == #index as u8 => {
+						#create
+					},
+				}
+			});
+
+			quote! {
+				match #input_.read_byte()? {
+					#( #recurse )*
+					_ => None,
+				}
+
+			}
+
+		},
+		// NOTE [ToDr] we currently don't use unions at all.
+		Data::Union(_) => unimplemented!(),
 	}
 }
 

substrate/codec/derive/src/encode.rs

@@ -16,11 +16,11 @@
 
 use proc_macro2::{Span, TokenStream};
 use syn::{
-	Data, Fields, Index,
+	Data, Fields, Ident, Index,
 	spanned::Spanned,
 };
 
-pub fn quote(data: &Data, self_: &TokenStream, dest_: &TokenStream) -> TokenStream {
+pub fn quote(data: &Data, type_name: &Ident, self_: &TokenStream, dest_: &TokenStream) -> TokenStream {
 	let call_site = Span::call_site();
 	match *data {
 		Data::Struct(ref data) => match data.fields {
@@ -34,7 +34,7 @@ pub fn quote(data: &Data, self_: &TokenStream, dest_: &TokenStream) -> TokenStre
 					}
 				});
 
-				quote! {
+				quote_spanned! { call_site =>
 					#( #recurse )*
 				}
 			},
@@ -51,13 +51,73 @@ pub fn quote(data: &Data, self_: &TokenStream, dest_: &TokenStream) -> TokenStre
 					#( #recurse )*
 				}
 			},
-			// Do nothing, we don't encode unit types?
+			// Do nothing, we don't encode unit type and assume it's always decodable.
 			Fields::Unit => quote! {
 				drop(#dest_);
 			},
 		},
-		// TODO [ToDr] Encode enums and unions
-		Data::Enum(_) | Data::Union(_) => unimplemented!(),
+		Data::Enum(ref data) => {
+			assert!(data.variants.len() < 256, "Currently only enums with less than 255 variants are encodable.");
+
+			let recurse = data.variants.iter().enumerate().map(|(i, f)| {
+				let name = &f.ident;
+
+				match f.fields {
+					Fields::Named(ref fields) => {
+						let names = fields.named.iter().map(|f| {
+							let field_name = &f.ident;
+							quote_spanned!(call_site => #field_name)
+						}).collect::<Vec<_>>();
+
+						let fields = names.iter();
+						let encode_fields = names.iter().map(|f| {
+							quote_spanned! { call_site => #dest_.push(#f); }
+						});
+
+						quote_spanned! { f.span() =>
+							#type_name :: #name { #( ref #fields, )* } => {
+								#dest_.push_byte(#i as u8);
+								#( #encode_fields )*
+							}
+						}
+					},
+					Fields::Unnamed(ref fields) => {
+						let names = fields.unnamed.iter().enumerate().map(|(i, _f)| {
+							let ident = Ident::new(&format!("f_{}", i), call_site);
+							quote_spanned!(call_site => #ident)
+						}).collect::<Vec<_>>();
+
+						let fields = names.iter();
+						let encode_fields = names.iter().map(|f| {
+							quote_spanned! { call_site => #dest_.push(#f); }
+						});
+
+						quote_spanned! { f.span() =>
+							#type_name :: #name ( #( ref #fields, )* ) => {
+								#dest_.push_byte(#i as u8);
+								#( #encode_fields )*
+							}
+						}
+					},
+					Fields::Unit => {
+						quote_spanned! { f.span() =>
+							#type_name :: #name => {
+								#dest_.push_byte(#i as u8);
+							}
+						}
+					},
+				}
+			});
+
+			quote! {
+				match *#self_ {
+					#( #recurse )*,
+					_ => unreachable!(),
+				}
+			}
+		},
+		// NOTE [ToDr] we currently don't use unions at all.
+		Data::Union(_) => unimplemented!(),
 	}
 }
 

substrate/codec/derive/src/lib.rs

@@ -46,7 +46,7 @@ pub fn encode_derive(input: TokenStream) -> TokenStream {
 
 	let self_ = quote!(self);
 	let dest_ = quote!(dest);
-	let encoding = encode::quote(&input.data, &self_, &dest_);
+	let encoding = encode::quote(&input.data, name, &self_, &dest_);
 
 	let expanded = quote! {
 		impl #impl_generics ::codec::Encode for #name #ty_generics #where_clause {

substrate/codec/derive/tests/mod.rs

@@ -24,7 +24,6 @@ use codec::{Encode, Decode};
 #[derive(Debug, PartialEq, Encode, Decode)]
 struct Unit;
 
-
 #[derive(Debug, PartialEq, Encode, Decode)]
 struct Indexed(u32, u64);
 
@@ -35,13 +34,78 @@ struct Struct<A, B, C> {
 	pub c: C,
 }
 
+type TestType = Struct<u32, u64, Vec<u8>>;
+
 impl <A, B, C> Struct<A, B, C> {
 	fn new(a: A, b: B, c: C) -> Self {
 		Self { a, b, c }
 	}
 }
 
-type TestType = Struct<u32, u64, Vec<u8>>;
+#[derive(Debug, PartialEq, Encode, Decode)]
+enum EnumType {
+	A,
+	B(u32, u64),
+	C {
+		a: u32,
+		b: u64,
+	},
+}
+
+#[derive(Debug, PartialEq, Encode, Decode)]
+enum EnumWithDiscriminant {
+	A = 1,
+	B = 15,
+	C = 255,
+}
+
+#[test]
+fn should_work_for_simple_enum() {
+	let a = EnumType::A;
+	let b = EnumType::B(1, 2);
+	let c = EnumType::C { a: 1, b: 2 };
+
+	a.using_encoded(|ref slice| {
+		assert_eq!(slice, &b"\0");
+	});
+	b.using_encoded(|ref slice| {
+		assert_eq!(slice, &b"\x01\x01\0\0\0\x02\0\0\0\0\0\0\0");
+	});
+	c.using_encoded(|ref slice| {
+		assert_eq!(slice, &b"\x02\x01\0\0\0\x02\0\0\0\0\0\0\0");
+	});
+
+	let mut da: &[u8] = b"\0";
+	assert_eq!(EnumType::decode(&mut da), Some(a));
+	let mut db: &[u8] = b"\x01\x01\0\0\0\x02\0\0\0\0\0\0\0";
+	assert_eq!(EnumType::decode(&mut db), Some(b));
+	let mut dc: &[u8] = b"\x02\x01\0\0\0\x02\0\0\0\0\0\0\0";
+	assert_eq!(EnumType::decode(&mut dc), Some(c));
+	let mut dz: &[u8] = &[4];
+	assert_eq!(EnumType::decode(&mut dz), None);
+}
+
+#[test]
+fn should_work_for_enum_with_discriminant() {
+	EnumWithDiscriminant::A.using_encoded(|ref slice| {
+		assert_eq!(slice, &[1]);
+	});
+	EnumWithDiscriminant::B.using_encoded(|ref slice| {
+		assert_eq!(slice, &[15]);
+	});
+	EnumWithDiscriminant::C.using_encoded(|ref slice| {
+		assert_eq!(slice, &[255]);
+	});
+
+	let mut da: &[u8] = &[1];
+	assert_eq!(EnumWithDiscriminant::decode(&mut da), Some(EnumWithDiscriminant::A));
+	let mut db: &[u8] = &[15];
+	assert_eq!(EnumWithDiscriminant::decode(&mut db), Some(EnumWithDiscriminant::B));
+	let mut dc: &[u8] = &[255];
+	assert_eq!(EnumWithDiscriminant::decode(&mut dc), Some(EnumWithDiscriminant::C));
+	let mut dz: &[u8] = &[2];
+	assert_eq!(EnumWithDiscriminant::decode(&mut dz), None);
+}
 
 #[test]
 fn should_derive_encode() {

substrate/network/src/message.rs

@@ -114,8 +114,7 @@ impl Decode for BlockAttributes {
 	}
 }
 
-// TODO [ToDr] as u8?
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Encode, Decode)]
 /// Block enumeration direction.
 pub enum Direction {
 	/// Enumerate in ascending order (from child to parent).
@@ -124,22 +123,6 @@ pub enum Direction {
 	Descending = 1,
 }
 
-impl Encode for Direction {
-	fn encode_to<T: Output>(&self, dest: &mut T) {
-		dest.push_byte(*self as u8)
-	}
-}
-
-impl Decode for Direction {
-	fn decode<I: Input>(input: &mut I) -> Option<Self> {
-		match input.read_byte()? {
-			x if x == Direction::Ascending as u8 => Some(Direction::Ascending),
-			x if x == Direction::Descending as u8 => Some(Direction::Descending),
-			_ => None,
-		}
-	}
-}
-
 /// Remote call response.
 #[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
 pub struct RemoteCallResponse {