Mock-builder with generic methods support

libs/mock-builder/Cargo.toml

@@ -10,9 +10,11 @@ version = "0.0.1"
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]
 
+[dependencies]
+frame-support = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
+
 [dev-dependencies]
 codec = { package = "parity-scale-codec", version = "3.0.0", features = ["derive"] }
-frame-support = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
 frame-system = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }
 scale-info = { version = "2.3.0", features = ["derive"] }
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "polkadot-v0.9.37" }

libs/mock-builder/src/lib.rs

@@ -39,177 +39,72 @@
 //!
 //! Take a look to the [pallet tests](`tests/pallet.rs`) to have a user view of how to use this crate.
 
-/// Provide methods for register/execute calls
-pub mod storage {
-	use std::{any::Any, cell::RefCell, collections::HashMap};
+/// Provide functions for register/execute calls
+pub mod storage;
 
-	/// Identify a call in the call storage
-	pub type CallId = u64;
-
-	trait Callable {
-		fn as_any(&self) -> &dyn Any;
-	}
-
-	thread_local! {
-		static CALLS: RefCell<HashMap<CallId, Box<dyn Callable>>>
-			= RefCell::new(HashMap::default());
-	}
-
-	struct CallWrapper<Input, Output>(Box<dyn Fn(Input) -> Output>);
-
-	impl<Input: 'static, Output: 'static> Callable for CallWrapper<Input, Output> {
-		fn as_any(&self) -> &dyn Any {
-			self
-		}
-	}
-
-	/// Register a call into the call storage.
-	/// The registered call can be uniquely identified by the returned `CallId`.
-	pub fn register_call<F: Fn(Args) -> R + 'static, Args: 'static, R: 'static>(f: F) -> CallId {
-		CALLS.with(|state| {
-			let registry = &mut *state.borrow_mut();
-			let call_id = registry.len() as u64;
-			registry.insert(call_id, Box::new(CallWrapper(Box::new(f))));
-			call_id
-		})
-	}
-
-	/// Execute a call from the call storage identified by a `call_id`.
-	pub fn execute_call<Args: 'static, R: 'static>(call_id: CallId, args: Args) -> R {
-		CALLS.with(|state| {
-			let registry = &*state.borrow();
-			let call = registry.get(&call_id).unwrap();
-			call.as_any()
-				.downcast_ref::<CallWrapper<Args, R>>()
-				.expect("Bad mock implementation: expected other function type")
-				.0(args)
-		})
-	}
-}
+/// Provide functions for handle fuction locations
+pub mod location;
 
+use frame_support::{Blake2_128, StorageHasher, StorageMap};
+use location::FunctionLocation;
 pub use storage::CallId;
 
 /// Prefix that the register functions should have.
 pub const MOCK_FN_PREFIX: &str = "mock_";
 
-/// Gives the absolute string identification of a function.
-#[macro_export]
-macro_rules! function_locator {
-	() => {{
-		// Aux function to extract the path
-		fn f() {}
-
-		fn type_name_of<T>(_: T) -> &'static str {
-			std::any::type_name::<T>()
-		}
-		let name = type_name_of(f);
-		&name[..name.len() - "::f".len()]
-	}};
+/// Register a mock function into the mock function storage.
+/// This function should be called with a locator used as a function identification.
+pub fn register<Map, L, F, I, O>(locator: L, f: F)
+where
+	Map: StorageMap<<Blake2_128 as StorageHasher>::Output, CallId>,
+	L: Fn(),
+	F: Fn(I) -> O + 'static,
+	I: 'static,
+	O: 'static,
+{
+	let location = FunctionLocation::from(locator)
+		.normalize()
+		.strip_name_prefix(MOCK_FN_PREFIX)
+		.append_type_signature::<I, O>();
+
+	Map::insert(location.hash::<Blake2_128>(), storage::register_call(f));
 }
 
-/// Gives the string identification of a function.
-/// The identification will be the same no matter if it belongs to a trait or has an `except_`
-/// prefix name.
-#[macro_export]
-macro_rules! call_locator {
-	() => {{
-		let path_name = $crate::function_locator!();
-		let (path, name) = path_name.rsplit_once("::").expect("always ::");
-
-		let base_name = name.strip_prefix($crate::MOCK_FN_PREFIX).unwrap_or(name);
-		let correct_path = path
-			.strip_prefix("<")
-			.map(|trait_path| trait_path.split_once(" as").expect("always ' as'").0)
-			.unwrap_or(path);
-
-		format!("{}::{}", correct_path, base_name)
-	}};
+/// Execute a function from the function storage.
+/// This function should be called with a locator used as a function identification.
+pub fn execute<Map, L, I, O>(locator: L, input: I) -> O
+where
+	Map: StorageMap<<Blake2_128 as StorageHasher>::Output, CallId>,
+	L: Fn(),
+	I: 'static,
+	O: 'static,
+{
+	let location = FunctionLocation::from(locator)
+		.normalize()
+		.append_type_signature::<I, O>();
+
+	let call_id = Map::try_get(location.hash::<Blake2_128>())
+		.unwrap_or_else(|_| panic!("Mock was not found. Location: {location:?}"));
+
+	storage::execute_call(call_id, input).unwrap_or_else(|err| {
+		panic!("{err}. Location: {location:?}");
+	})
 }
 
-/// Register a call into the call storage.
-/// This macro should be called from the method that wants to register `f`.
-/// This macro must be called from a pallet with the `CallIds` storage.
-/// Check the main documentation.
+/// Register a mock function into the mock function storage.
+/// Same as `register()` but with using the locator who calls this macro.
 #[macro_export]
 macro_rules! register_call {
 	($f:expr) => {{
-		use frame_support::StorageHasher;
-
-		let call_id = frame_support::Blake2_128::hash($crate::call_locator!().as_bytes());
-
-		CallIds::<T>::insert(call_id, $crate::storage::register_call($f));
+		$crate::register::<CallIds<T>, _, _, _, _>(|| (), $f)
 	}};
 }
 
-/// Execute a call from the call storage.
-/// This macro should be called from the method that wants to execute `f`.
-/// This macro must be called from a pallet with the `CallIds` storage.
-/// Check the main documentation.
+/// Execute a function from the function storage.
+/// Same as `execute()` but with using the locator who calls this macro.
 #[macro_export]
 macro_rules! execute_call {
-	($params:expr) => {{
-		use frame_support::StorageHasher;
-
-		let hash = frame_support::Blake2_128::hash($crate::call_locator!().as_bytes());
-		let call_id = CallIds::<T>::get(hash).expect(&format!(
-			"Called to {}, but mock was not found",
-			$crate::call_locator!()
-		));
-
-		$crate::storage::execute_call(call_id, $params)
+	($input:expr) => {{
+		$crate::execute::<CallIds<T>, _, _, _>(|| (), $input)
 	}};
 }
-
-#[cfg(test)]
-mod tests {
-	trait TraitExample {
-		fn function_locator() -> String;
-		fn call_locator() -> String;
-	}
-
-	struct Example;
-
-	impl Example {
-		fn mock_function_locator() -> String {
-			function_locator!().into()
-		}
-
-		fn mock_call_locator() -> String {
-			call_locator!().into()
-		}
-	}
-
-	impl TraitExample for Example {
-		fn function_locator() -> String {
-			function_locator!().into()
-		}
-
-		fn call_locator() -> String {
-			call_locator!().into()
-		}
-	}
-
-	#[test]
-	fn function_locator() {
-		assert_eq!(
-			Example::mock_function_locator(),
-			"mock_builder::tests::Example::mock_function_locator"
-		);
-
-		assert_eq!(
-			Example::function_locator(),
-			"<mock_builder::tests::Example as \
-            mock_builder::tests::TraitExample>::function_locator"
-		);
-	}
-
-	#[test]
-	fn call_locator() {
-		assert_eq!(
-			Example::call_locator(),
-			"mock_builder::tests::Example::call_locator"
-		);
-
-		assert_eq!(Example::call_locator(), Example::mock_call_locator());
-	}
-}

libs/mock-builder/src/location.rs

@@ -0,0 +1,160 @@
+use frame_support::StorageHasher;
+
+/// Absolute string identification of function.
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct FunctionLocation(String);
+
+impl FunctionLocation {
+	/// Creates a location for the function which created the given closure used as a locator
+	pub fn from<F: Fn()>(_: F) -> Self {
+		let location = std::any::type_name::<F>();
+		let location = &location[..location.len() - "::{{closure}}".len()];
+
+		// Remove generic attributes from signature if it has any
+		let location = location
+			.ends_with('>')
+			.then(|| {
+				let mut count = 0;
+				for (i, c) in location.chars().rev().enumerate() {
+					if c == '>' {
+						count += 1;
+					} else if c == '<' {
+						count -= 1;
+						if count == 0 {
+							return location.split_at(location.len() - i - 1).0;
+						}
+					}
+				}
+				panic!("Expected '<' symbol to close '>'");
+			})
+			.unwrap_or(location);
+
+		Self(location.into())
+	}
+
+	/// Normalize the location, allowing to identify the function
+	/// no matter if it belongs to a trait or not.
+	pub fn normalize(self) -> Self {
+		let (path, name) = self.0.rsplit_once("::").expect("always ::");
+		let path = path
+			.strip_prefix('<')
+			.map(|trait_path| trait_path.split_once(" as").expect("always ' as'").0)
+			.unwrap_or(path);
+
+		Self(format!("{}::{}", path, name))
+	}
+
+	/// Remove the prefix from the function name.
+	pub fn strip_name_prefix(self, prefix: &str) -> Self {
+		let (path, name) = self.0.rsplit_once("::").expect("always ::");
+		let name = name.strip_prefix(prefix).unwrap_or_else(|| {
+			panic!(
+				"Function '{name}' should have a '{prefix}' prefix. Location: {}",
+				self.0
+			)
+		});
+
+		Self(format!("{}::{}", path, name))
+	}
+
+	/// Add a representation of the function input and output types
+	pub fn append_type_signature<I, O>(self) -> Self {
+		Self(format!(
+			"{}:{}->{}",
+			self.0,
+			std::any::type_name::<I>(),
+			std::any::type_name::<O>(),
+		))
+	}
+
+	/// Generate a hash of the location
+	pub fn hash<Hasher: StorageHasher>(&self) -> Hasher::Output {
+		Hasher::hash(self.0.as_bytes())
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	const PREFIX: &str = "mock_builder::location::tests";
+
+	trait TraitExample {
+		fn method() -> FunctionLocation;
+		fn generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation;
+	}
+
+	struct Example;
+
+	impl Example {
+		fn mock_method() -> FunctionLocation {
+			FunctionLocation::from(|| ())
+		}
+
+		fn mock_generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation {
+			FunctionLocation::from(|| ())
+		}
+	}
+
+	impl TraitExample for Example {
+		fn method() -> FunctionLocation {
+			FunctionLocation::from(|| ())
+		}
+
+		fn generic_method<A: Into<i32>>(_: impl Into<u32>) -> FunctionLocation {
+			FunctionLocation::from(|| ())
+		}
+	}
+
+	#[test]
+	fn function_location() {
+		assert_eq!(
+			Example::mock_method().0,
+			format!("{PREFIX}::Example::mock_method")
+		);
+
+		assert_eq!(
+			Example::mock_generic_method::<i8>(0u8).0,
+			format!("{PREFIX}::Example::mock_generic_method")
+		);
+
+		assert_eq!(
+			Example::method().0,
+			format!("<{PREFIX}::Example as {PREFIX}::TraitExample>::method")
+		);
+
+		assert_eq!(
+			Example::generic_method::<i8>(0u8).0,
+			format!("<{PREFIX}::Example as {PREFIX}::TraitExample>::generic_method")
+		);
+	}
+
+	#[test]
+	fn normalized_function_location() {
+		assert_eq!(
+			Example::mock_method().normalize().0,
+			format!("{PREFIX}::Example::mock_method")
+		);
+
+		assert_eq!(
+			Example::method().normalize().0,
+			format!("{PREFIX}::Example::method")
+		);
+	}
+
+	#[test]
+	fn striped_function_location() {
+		assert_eq!(
+			Example::mock_method().strip_name_prefix("mock_").0,
+			format!("{PREFIX}::Example::method")
+		);
+	}
+
+	#[test]
+	fn appended_type_signature() {
+		assert_eq!(
+			Example::mock_method().append_type_signature::<i8, u8>().0,
+			format!("{PREFIX}::Example::mock_method:i8->u8")
+		);
+	}
+}