test: add unit test to catch missing distribution to subsystems faster (#2495)

* test: add unit test to catch missing distribution to subsystems faster

* add a simple count

* introduce proc macro to generate dispatch type

* refactor

* refactor

* chore: add license

* fixup unit test

* fixup merge

* better errors

* better fmt

* fix error spans

* better docs

* better error messages

* ui test foo

* Update node/subsystem/dispatch-gen/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update node/network/bridge/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update node/subsystem/Cargo.toml

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update node/subsystem/dispatch-gen/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update node/subsystem/dispatch-gen/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Update node/network/bridge/src/lib.rs

Co-authored-by: Andronik Ordian <write@reusable.software>

* fix compilation

* use find_map

* drop the silly 2, use _inner instead

* Update node/network/bridge/src/lib.rs

Co-authored-by: Andronik Ordian <write@reusable.software>

* Update node/subsystem/dispatch-gen/src/lib.rs

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* nail deps down

* more into()

* flatten

* missing use statement

* fix messages order

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>
Co-authored-by: Andronik Ordian <write@reusable.software>

polkadot/node/subsystem/Cargo.toml

@@ -23,6 +23,7 @@ polkadot-node-network-protocol = { path = "../network/protocol" }
 polkadot-primitives = { path = "../../primitives" }
 polkadot-statement-table = { path = "../../statement-table" }
 polkadot-node-jaeger = { path = "../jaeger" }
+polkadot-procmacro-subsystem-dispatch-gen = { path = "dispatch-gen" }
 sc-network = { git = "https://github.com/paritytech/substrate", branch = "master" }
 smallvec = "1.6.1"
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }

polkadot/node/subsystem/dispatch-gen/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "polkadot-procmacro-subsystem-dispatch-gen"
+version = "0.1.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2018"
+description = "Small proc macro to create the distribution code for network events"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { version = "1.0.60", features = ["full"] }
+quote = "1.0.9"
+proc-macro2 = "1.0.24"
+assert_matches = "1.5.0"
+
+[dev-dependencies]
+trybuild = "1.0.41"

polkadot/node/subsystem/dispatch-gen/src/lib.rs

@@ -0,0 +1,208 @@
+// Copyright 2021 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+use proc_macro2::TokenStream;
+use quote::{quote, ToTokens};
+use std::fmt;
+use syn::{parse2, Error, Fields, FieldsNamed, FieldsUnnamed, Ident, ItemEnum, Path, Result, Type, Variant};
+
+#[proc_macro_attribute]
+pub fn subsystem_dispatch_gen(attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
+	let attr: TokenStream = attr.into();
+	let item: TokenStream = item.into();
+	let mut backup = item.clone();
+	impl_subsystem_dispatch_gen(attr.into(), item).unwrap_or_else(|err| {
+		backup.extend(err.to_compile_error());
+		backup
+	}).into()
+}
+
+/// An enum variant without base type.
+#[derive(Clone)]
+struct EnumVariantDispatchWithTy {
+	// enum ty name
+	ty: Ident,
+	// variant
+	variant: EnumVariantDispatch,
+}
+
+impl fmt::Debug for EnumVariantDispatchWithTy {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		write!(f, "{}::{:?}", self.ty, self.variant)
+	}
+}
+
+impl ToTokens for EnumVariantDispatchWithTy {
+	fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+		if let Some(inner) = &self.variant.inner {
+			let enum_name = &self.ty;
+			let variant_name = &self.variant.name;
+
+			let quoted = quote! {
+				#enum_name::#variant_name(#inner::from(event))
+			};
+			quoted.to_tokens(tokens);
+		}
+	}
+}
+
+/// An enum variant without the base type, contains the relevant inner type.
+#[derive(Clone)]
+struct EnumVariantDispatch {
+	/// variant name
+	name: Ident,
+	/// The inner type for which a `From::from` impl is anticipated from the input type.
+	/// No code will be generated for this enum variant if `inner` is `None`.
+	inner: Option<Type>,
+}
+
+impl fmt::Debug for EnumVariantDispatch {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		write!(f, "{}(..)", self.name)
+	}
+}
+
+fn prepare_enum_variant(variant: &mut Variant) -> Result<EnumVariantDispatch> {
+	let skip = variant.attrs.iter().find(|attr| attr.path.is_ident("skip")).is_some();
+	variant.attrs = variant.attrs.iter().filter(|attr| !attr.path.is_ident("skip")).cloned().collect::<Vec<_>>();
+
+	let variant = variant.clone();
+	let span = variant.ident.span();
+	let inner = match variant.fields.clone() {
+		// look for one called inner
+		Fields::Named(FieldsNamed { brace_token: _, named }) if !skip => named
+			.iter()
+			.find_map(
+				|field| {
+					if let Some(ident) = &field.ident {
+						if ident == "inner" {
+							return Some(Some(field.ty.clone()))
+						}
+					}
+					None
+				},
+			)
+			.ok_or_else(|| {
+				Error::new(span, "To dispatch with struct enum variant, one element must named `inner`")
+			})?,
+
+		// technically, if it has no inner types we cound not require the #[skip] annotation, but better make it consistent
+		Fields::Unnamed(FieldsUnnamed { paren_token: _, unnamed }) if !skip => unnamed
+			.first()
+			.map(|field| Some(field.ty.clone()))
+			.ok_or_else(|| Error::new(span, "Must be annotated with skip, even if no inner types exist."))?,
+		_ if skip => None,
+		Fields::Unit => {
+			return Err(Error::new(
+				span,
+				"Must be annotated with #[skip].",
+			))
+		}
+		Fields::Unnamed(_) => {
+			return Err(Error::new(
+				span,
+				"Must be annotated with #[skip] or have in `inner` element which impls `From<_>`.",
+			))
+		}
+		Fields::Named(_) => {
+			return Err(Error::new(
+				span,
+				"Must be annotated with #[skip] or the first wrapped type must impl `From<_>`.",
+			))
+		}
+	};
+
+	Ok(EnumVariantDispatch { name: variant.ident, inner })
+}
+
+fn impl_subsystem_dispatch_gen(attr: TokenStream, item: TokenStream) -> Result<proc_macro2::TokenStream> {
+	let event_ty = parse2::<Path>(attr)?;
+
+	let mut ie = parse2::<ItemEnum>(item)?;
+
+	let message_enum = ie.ident.clone();
+	let variants = ie.variants.iter_mut().try_fold(Vec::<EnumVariantDispatchWithTy>::new(), |mut acc, variant| {
+		let variant = prepare_enum_variant(variant)?;
+		if variant.inner.is_some() {
+			acc.push(EnumVariantDispatchWithTy { ty: message_enum.clone(), variant })
+		}
+		Ok::<_, syn::Error>(acc)
+	})?;
+
+	let mut orig = ie.to_token_stream();
+
+	let msg = "Generated by #[subsystem_dispatch_gen] proc-macro.";
+
+	orig.extend(quote! {
+		impl #message_enum {
+			#[doc = #msg]
+			pub fn dispatch_iter(event: #event_ty) -> impl Iterator<Item=Self> + Send {
+				let mut iter = None.into_iter();
+
+				#(
+					let mut iter = iter.chain(std::iter::once(event.focus().ok().map(|event| {
+						#variants
+					})));
+				)*
+				iter.filter_map(|x| x)
+			}
+		}
+	});
+	Ok(orig)
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	#[test]
+	fn basic() {
+		let attr = quote! {
+			NetEvent<foo::Bar>
+		};
+
+		let item = quote! {
+			/// Documentation.
+			#[derive(Clone)]
+			enum AllMessages {
+
+				Sub1(Inner1),
+
+				#[skip]
+				/// D3
+				Sub3,
+
+				/// D4
+				#[skip]
+				Sub4(Inner2),
+
+				/// D2
+				Sub2(Inner2),
+			}
+		};
+
+		let output = impl_subsystem_dispatch_gen(attr, item).expect("Simple example always works. qed");
+		println!("//generated:");
+		println!("{}", output);
+	}
+
+	#[test]
+	fn ui() {
+		let t = trybuild::TestCases::new();
+		t.compile_fail("tests/ui/err-*.rs");
+		t.pass("tests/ui/ok-*.rs");
+	}
+}

polkadot/node/subsystem/dispatch-gen/tests/ui/err-01-missing-skip.rs

@@ -0,0 +1,37 @@
+#![allow(dead_code)]
+
+use polkadot_procmacro_subsystem_dispatch_gen::subsystem_dispatch_gen;
+
+/// The event type in question.
+#[derive(Clone, Copy)]
+enum Event {
+	Smth,
+	Else,
+}
+
+impl Event {
+	fn focus(&self) -> std::result::Result<Inner, ()> {
+		unimplemented!("foo")
+	}
+}
+
+/// This should have a `From<Event>` impl but does not.
+#[derive(Clone)]
+enum Inner {
+	Foo,
+	Bar(Event),
+}
+
+#[subsystem_dispatch_gen(Event)]
+#[derive(Clone)]
+enum AllMessages {
+	/// Foo
+	Vvvvvv(Inner),
+
+    /// Missing a `#[skip]` annotation
+    Uuuuu,
+}
+
+fn main() {
+    let _x = AllMessages::dispatch_iter(Event::Else);
+}

polkadot/node/subsystem/dispatch-gen/tests/ui/err-01-missing-skip.stderr

@@ -0,0 +1,14 @@
+error: Must be annotated with #[skip].
+  --> $DIR/err-01-missing-skip.rs:32:5
+   |
+32 |     Uuuuu,
+   |     ^^^^^
+
+error[E0599]: no variant or associated item named `dispatch_iter` found for enum `AllMessages` in the current scope
+  --> $DIR/err-01-missing-skip.rs:36:27
+   |
+27 | enum AllMessages {
+   | ---------------- variant or associated item `dispatch_iter` not found here
+...
+36 |     let _x = AllMessages::dispatch_iter(Event::Else);
+   |                           ^^^^^^^^^^^^^ variant or associated item not found in `AllMessages`

polkadot/node/subsystem/dispatch-gen/tests/ui/err-02-missing-from.rs

@@ -0,0 +1,41 @@
+#![allow(dead_code)]
+
+use polkadot_procmacro_subsystem_dispatch_gen::subsystem_dispatch_gen;
+
+/// The event type in question.
+#[derive(Clone, Copy, Debug)]
+enum Event {
+	Smth,
+	Else,
+}
+
+impl Event {
+	fn focus(&self) -> std::result::Result<Intermediate, ()> {
+		Ok(Intermediate(self.clone()))
+	}
+}
+
+#[derive(Debug, Clone)]
+struct Intermediate(Event);
+
+
+/// This should have a `From<Event>` impl but does not.
+#[derive(Debug, Clone)]
+enum Inner {
+	Foo,
+	Bar(Intermediate),
+}
+
+#[subsystem_dispatch_gen(Event)]
+#[derive(Clone)]
+enum AllMessages {
+	/// Foo
+	Vvvvvv(Inner),
+
+    #[skip]
+    Uuuuu,
+}
+
+fn main() {
+    let _x = AllMessages::dispatch_iter(Event::Else);
+}

polkadot/node/subsystem/dispatch-gen/tests/ui/err-02-missing-from.stderr

@@ -0,0 +1,10 @@
+error[E0308]: mismatched types
+  --> $DIR/err-02-missing-from.rs:29:1
+   |
+29 | #[subsystem_dispatch_gen(Event)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   | |
+   | expected enum `Inner`, found struct `Intermediate`
+   | help: try using a variant of the expected enum: `Inner::Bar(#[subsystem_dispatch_gen(Event)])`
+   |
+   = note: this error originates in an attribute macro (in Nightly builds, run with -Z macro-backtrace for more info)

polkadot/node/subsystem/dispatch-gen/tests/ui/ok-01-with-intermediate.rs

@@ -0,0 +1,48 @@
+#![allow(dead_code)]
+
+use polkadot_procmacro_subsystem_dispatch_gen::subsystem_dispatch_gen;
+
+/// The event type in question.
+#[derive(Clone, Copy, Debug)]
+enum Event {
+	Smth,
+	Else,
+}
+
+impl Event {
+	fn focus(&self) -> std::result::Result<Intermediate, ()> {
+		Ok(Intermediate(self.clone()))
+	}
+}
+
+
+#[derive(Debug, Clone)]
+struct Intermediate(Event);
+
+
+/// This should have a `From<Event>` impl but does not.
+#[derive(Clone, Debug)]
+enum Inner {
+	Foo,
+	Bar(Intermediate),
+}
+
+impl From<Intermediate> for Inner {
+	fn from(src: Intermediate) -> Self {
+		Inner::Bar(src)
+	}
+}
+
+#[subsystem_dispatch_gen(Event)]
+#[derive(Clone)]
+enum AllMessages {
+	/// Foo
+	Vvvvvv(Inner),
+
+    #[skip]
+    Uuuuu,
+}
+
+fn main() {
+    let _x = AllMessages::dispatch_iter(Event::Else);
+}

polkadot/node/subsystem/src/lib.rs

@@ -32,14 +32,14 @@ use polkadot_primitives::v1::{Hash, BlockNumber};
 use async_trait::async_trait;
 use smallvec::SmallVec;
 
-use crate::messages::AllMessages;
-
 pub mod errors;
 pub mod messages;
 
 pub use polkadot_node_jaeger as jaeger;
 pub use jaeger::*;
 
+use self::messages::AllMessages;
+
 /// How many slots are stack-reserved for active leaves updates
 ///
 /// If there are fewer than this number of slots, then we've wasted some stack space.

polkadot/node/subsystem/src/messages.rs

@@ -44,6 +44,7 @@ use polkadot_primitives::v1::{
 	CandidateIndex, GroupIndex,
 };
 use polkadot_statement_table::v1::Misbehavior;
+use polkadot_procmacro_subsystem_dispatch_gen::subsystem_dispatch_gen;
 use std::{sync::Arc, collections::btree_map::BTreeMap};
 
 
@@ -171,7 +172,7 @@ impl CandidateValidationMessage {
 
 
 /// Messages received by the Collator Protocol subsystem.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum CollatorProtocolMessage {
 	/// Signal to the collator protocol that it should connect to validators with the expectation
 	/// of collating on the given para. This is only expected to be called once, early on, if at all,
@@ -195,6 +196,7 @@ pub enum CollatorProtocolMessage {
 	/// Notify a collator that its collation was seconded.
 	NotifyCollationSeconded(CollatorId, SignedFullStatement),
 	/// Get a network bridge update.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::CollatorProtocolMessage>),
 }
 
@@ -270,13 +272,15 @@ impl NetworkBridgeMessage {
 #[derive(Debug, derive_more::From)]
 pub enum AvailabilityDistributionMessage {
 	/// Event from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::AvailabilityDistributionMessage>),
 	/// Incoming request for an availability chunk.
+	#[from]
 	AvailabilityFetchingRequest(IncomingRequest<req_res_v1::AvailabilityFetchingRequest>)
 }
 
 /// Availability Recovery Message.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum AvailabilityRecoveryMessage {
 	/// Recover available data from validators on the network.
 	RecoverAvailableData(
@@ -286,6 +290,7 @@ pub enum AvailabilityRecoveryMessage {
 		oneshot::Sender<Result<AvailableData, crate::errors::RecoveryError>>,
 	),
 	/// Event from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::AvailabilityRecoveryMessage>),
 }
 
@@ -300,12 +305,13 @@ impl AvailabilityDistributionMessage {
 }
 
 /// Bitfield distribution message.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum BitfieldDistributionMessage {
 	/// Distribute a bitfield via gossip to other validators.
 	DistributeBitfield(Hash, SignedAvailabilityBitfield),
 
 	/// Event from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::BitfieldDistributionMessage>),
 }
 
@@ -509,12 +515,13 @@ impl RuntimeApiMessage {
 }
 
 /// Statement distribution message.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum StatementDistributionMessage {
 	/// We have originated a signed statement in the context of
 	/// given relay-parent hash and it should be distributed to other validators.
 	Share(Hash, SignedFullStatement),
 	/// Event from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::StatementDistributionMessage>),
 }
 
@@ -572,7 +579,7 @@ impl BoundToRelayParent for ProvisionerMessage {
 }
 
 /// Message to the PoV Distribution subsystem.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum PoVDistributionMessage {
 	/// Fetch a PoV from the network.
 	///
@@ -583,6 +590,7 @@ pub enum PoVDistributionMessage {
 	/// The PoV should correctly hash to the PoV hash mentioned in the CandidateDescriptor
 	DistributePoV(Hash, CandidateDescriptor, Arc<PoV>),
 	/// An update from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::PoVDistributionMessage>),
 }
 
@@ -662,7 +670,7 @@ pub enum ApprovalVotingMessage {
 }
 
 /// Message to the Approval Distribution subsystem.
-#[derive(Debug)]
+#[derive(Debug, derive_more::From)]
 pub enum ApprovalDistributionMessage {
 	/// Notify the `ApprovalDistribution` subsystem about new blocks
 	/// and the candidates contained within them.
@@ -675,21 +683,28 @@ pub enum ApprovalDistributionMessage {
 	/// If not, the subsystem is free to drop the message.
 	DistributeApproval(IndirectSignedApprovalVote),
 	/// An update from the network bridge.
+	#[from]
 	NetworkBridgeUpdateV1(NetworkBridgeEvent<protocol_v1::ApprovalDistributionMessage>),
 }
 
 /// A message type tying together all message types that are used across Subsystems.
+#[subsystem_dispatch_gen(NetworkBridgeEvent<protocol_v1::ValidationProtocol>)]
 #[derive(Debug, derive_more::From)]
 pub enum AllMessages {
 	/// Message for the validation subsystem.
+	#[skip]
 	CandidateValidation(CandidateValidationMessage),
 	/// Message for the candidate backing subsystem.
+	#[skip]
 	CandidateBacking(CandidateBackingMessage),
 	/// Message for the candidate selection subsystem.
+	#[skip]
 	CandidateSelection(CandidateSelectionMessage),
 	/// Message for the Chain API subsystem.
+	#[skip]
 	ChainApi(ChainApiMessage),
 	/// Message for the Collator Protocol subsystem.
+	#[skip]
 	CollatorProtocol(CollatorProtocolMessage),
 	/// Message for the statement distribution subsystem.
 	StatementDistribution(StatementDistributionMessage),
@@ -700,20 +715,27 @@ pub enum AllMessages {
 	/// Message for the bitfield distribution subsystem.
 	BitfieldDistribution(BitfieldDistributionMessage),
 	/// Message for the bitfield signing subsystem.
+	#[skip]
 	BitfieldSigning(BitfieldSigningMessage),
 	/// Message for the Provisioner subsystem.
+	#[skip]
 	Provisioner(ProvisionerMessage),
 	/// Message for the PoV Distribution subsystem.
 	PoVDistribution(PoVDistributionMessage),
 	/// Message for the Runtime API subsystem.
+	#[skip]
 	RuntimeApi(RuntimeApiMessage),
 	/// Message for the availability store subsystem.
+	#[skip]
 	AvailabilityStore(AvailabilityStoreMessage),
 	/// Message for the network bridge subsystem.
+	#[skip]
 	NetworkBridge(NetworkBridgeMessage),
 	/// Message for the Collation Generation subsystem.
+	#[skip]
 	CollationGeneration(CollationGenerationMessage),
 	/// Message for the Approval Voting subsystem.
+	#[skip]
 	ApprovalVoting(ApprovalVotingMessage),
 	/// Message for the Approval Distribution subsystem.
 	ApprovalDistribution(ApprovalDistributionMessage),