custom weight function wrapper

frame/support/src/weights.rs

@@ -114,6 +114,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 {
@@ -223,6 +234,28 @@ 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);
+
+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()
+	}
+}
+
 /// Implementation for unchecked extrinsic.
 impl<Address, Call, Signature, Extra> GetDispatchInfo
 	for UncheckedExtrinsic<Address, Call, Signature, Extra>
@@ -258,3 +291,57 @@ impl<Call: Encode, Extra: Encode> GetDispatchInfo for sr_primitives::testing::Te
 	}
 }
 
+#[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)]
+			fn f11(_origin, _a: u32, _eb: u32) { unimplemented!(); }
+
+			#[weight = FunctionOf(|_: (&u32, &u32)| 0, 1u8)]
+			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,
+		);
+	}
+}