custom weight function wrapper

frame/staking/src/lib.rs

@@ -1607,7 +1607,6 @@ impl<T: Trait> Module<T> {
 	/// For each element in the iterator the given number of points in u32 is added to the
 	/// validator, thus duplicates are handled.
 	pub fn reward_by_indices(validators_points: impl IntoIterator<Item = (u32, u32)>) {
-		// TODO: This can be optimised once #3302 is implemented.
 		let current_elected_len = <Module<T>>::current_elected().len() as u32;
 
 		CurrentEraPointsEarned::mutate(|rewards| {

frame/staking/src/slashing.rs

@@ -649,9 +649,6 @@ fn pay_reporters<T: Trait>(
 	T::Slash::on_unbalanced(value_slashed);
 }
 
-// TODO: function for undoing a slash.
-//
-
 #[cfg(test)]
 mod tests {
 	use super::*;

frame/support/src/weights.rs

@@ -122,6 +122,17 @@ impl Default for DispatchClass {
 	}
 }
 
+// Utility implementation to provide numbers instead of the full enum.
+impl From<u8> for DispatchClass {
+	fn from(x: u8) -> Self {
+		match x {
+			0 => DispatchClass::Normal,
+			1 => DispatchClass::Operational,
+			_ => panic!("Reserved for other dispatch classes."),
+		}
+	}
+}
+
 impl From<SimpleDispatchInfo> for DispatchClass {
 	fn from(tx: SimpleDispatchInfo) -> Self {
 		match tx {
@@ -236,6 +247,35 @@ impl SimpleDispatchInfo {
 	}
 }
 
+/// A struct to represent a weight which is a function of the input arguments. The given closure's
+/// arguments must match the dispatch's argument types, given as a tuple.
+pub struct FunctionOf<F, Class>(pub F, pub Class, pub bool);
+
+impl<Args, F, Class> WeighData<Args> for FunctionOf<F, Class>
+where
+	F : Fn(Args) -> Weight
+{
+	fn weigh_data(&self, args: Args) -> Weight {
+		(self.0)(args)
+	}
+}
+
+impl<Args, F, Class> ClassifyDispatch<Args> for FunctionOf<F, Class>
+where
+	Class: Into<DispatchClass> + Clone,
+{
+	fn classify_dispatch(&self, _: Args) -> DispatchClass {
+		self.1.clone().into()
+	}
+}
+
+impl<F, Class> PaysFee for FunctionOf<F, Class> {
+	fn pays_fee(&self) -> bool {
+		self.2
+	}
+}
+
+
 /// Implementation for unchecked extrinsic.
 impl<Address, Call, Signature, Extra> GetDispatchInfo
 	for UncheckedExtrinsic<Address, Call, Signature, Extra>
@@ -271,3 +311,58 @@ impl<Call: Encode, Extra: Encode> GetDispatchInfo for sp_runtime::testing::TestX
 		}
 	}
 }
+
+#[cfg(test)]
+#[allow(dead_code)]
+mod tests {
+	use crate::decl_module;
+	use super::*;
+
+	pub trait Trait {
+		type Origin;
+		type Balance;
+		type BlockNumber;
+	}
+
+	pub struct TraitImpl {}
+
+	impl Trait for TraitImpl {
+		type Origin = u32;
+		type BlockNumber = u32;
+		type Balance = u32;
+	}
+
+	decl_module! {
+		pub struct Module<T: Trait> for enum Call where origin: T::Origin {
+			// no arguments, fixed weight
+			#[weight = SimpleDispatchInfo::FixedNormal(1000)]
+			fn f0(_origin) { unimplemented!(); }
+
+			// weight = a x 10 + b
+			#[weight = FunctionOf(|args: (&u32, &u32)| args.0 * 10 + args.1, DispatchClass::Normal, true)]
+			fn f11(_origin, _a: u32, _eb: u32) { unimplemented!(); }
+
+			#[weight = FunctionOf(|_: (&u32, &u32)| 0, 1u8, true)]
+			fn f12(_origin, _a: u32, _eb: u32) { unimplemented!(); }
+		}
+	}
+
+	#[test]
+	fn weights_are_correct() {
+		assert_eq!(Call::<TraitImpl>::f11(10, 20).get_dispatch_info().weight, 120);
+		assert_eq!(Call::<TraitImpl>::f11(10, 20).get_dispatch_info().class, DispatchClass::Normal);
+		assert_eq!(Call::<TraitImpl>::f0().get_dispatch_info().weight, 1000);
+	}
+
+	#[test]
+	fn can_use_number_as_class() {
+		assert_eq!(
+			Call::<TraitImpl>::f11(0, 0).get_dispatch_info().class,
+			DispatchClass::Normal,
+		);
+		assert_eq!(
+			Call::<TraitImpl>::f12(0, 0).get_dispatch_info().class,
+			DispatchClass::Operational,
+		);
+	}
+}