refactor dynamic dispatch

Author: NikVolf

client/executor/common/src/error.rs

@@ -81,6 +81,14 @@ pub enum Error {
 	/// Execution of a host function failed.
 	#[display(fmt="Host function {} execution failed with: {}", _0, _1)]
 	FunctionExecution(String, String),
+	/// No table is present.
+	///
+	/// Call was requested that requires table but none was present in the instance.
+	NoTable,
+	/// No table is present.
+	///
+	/// Call was requested that requires specific entry in the table to be present.
+	NoTableEntryWithIndex(u32),
 }
 
 impl std::error::Error for Error {

client/executor/common/src/wasm_runtime.rs

@@ -19,6 +19,35 @@
 use crate::error::Error;
 use sp_wasm_interface::Value;
 
+/// How to locate method to call.
+pub enum InvokeLocation<'a> {
+	/// Call function exported with this name.
+	///
+	/// Located function should have (u32, u32) -> u64 signature.
+	Export(&'a str),
+	/// Call function by reference from table.
+	///
+	/// Located function should have (u32, u32) -> u64 signature.
+	Table(u32),
+	/// Call function by reference from table through a wrapper.
+	///
+	/// Located function should have (u32, u32, u32) -> u64 signature.
+	///
+	/// func will be passed to the
+	TableWithWrapper {
+		/// Wrapper for call.
+		dispatcher_ref: u32,
+		/// Actual function index that should be invoked.
+		func: u32,
+	},
+}
+
+impl<'a> From<&'a str> for InvokeLocation<'a> {
+	fn from(val: &'a str) -> InvokeLocation<'a> {
+		InvokeLocation::Export(val)
+	}
+}
+
 /// A trait that defines an abstract WASM runtime module.
 ///
 /// This can be implemented by an execution engine.
@@ -31,11 +60,24 @@ pub trait WasmModule: Sync + Send {
 ///
 /// This can be implemented by an execution engine.
 pub trait WasmInstance: Send {
-	/// Call a method on this WASM instance and reset it afterwards.
+	/// Call a method on this WASM instance.
+	///
+	/// Before execution, instance is reset.
+	///
+	/// Returns the encoded result on success.
+	fn call(&self, method: InvokeLocation, data: &[u8]) -> Result<Vec<u8>, Error>;
+
+	/// Call an exported method on this WASM instance.
+	///
+	/// Before execution, instance is reset.
+	///
 	/// Returns the encoded result on success.
-	fn call(&self, method: &str, data: &[u8]) -> Result<Vec<u8>, Error>;
+	fn call_export(&self, method: &str, data: &[u8]) -> Result<Vec<u8>, Error> {
+		self.call(method.into(), data)
+	}
 
 	/// Get the value from a global with the given `name`.
+	///
 	/// This method is only suitable for getting immutable globals.
 	fn get_global_const(&self, name: &str) -> Result<Option<Value>, Error>;
 }

client/executor/src/integration_tests/mod.rs

@@ -554,13 +554,13 @@ fn returns_mutable_static(wasm_method: WasmExecutionMethod) {
 	).expect("Creates runtime");
 
 	let instance = runtime.new_instance().unwrap();
-	let res = instance.call("returns_mutable_static", &[0]).unwrap();
+	let res = instance.call_export("returns_mutable_static", &[0]).unwrap();
 	assert_eq!(33, u64::decode(&mut &res[..]).unwrap());
 
 	// We expect that every invocation will need to return the initial
 	// value plus one. If the value increases more than that then it is
 	// a sign that the wasm runtime preserves the memory content.
-	let res = instance.call("returns_mutable_static", &[0]).unwrap();
+	let res = instance.call_export("returns_mutable_static", &[0]).unwrap();
 	assert_eq!(33, u64::decode(&mut &res[..]).unwrap());
 }
 
@@ -589,11 +589,11 @@ fn restoration_of_globals(wasm_method: WasmExecutionMethod) {
 	let instance = runtime.new_instance().unwrap();
 
 	// On the first invocation we allocate approx. 768KB (75%) of stack and then trap.
-	let res = instance.call("allocates_huge_stack_array", &true.encode());
+	let res = instance.call_export("allocates_huge_stack_array", &true.encode());
 	assert!(res.is_err());
 
 	// On the second invocation we allocate yet another 768KB (75%) of stack
-	let res = instance.call("allocates_huge_stack_array", &false.encode());
+	let res = instance.call_export("allocates_huge_stack_array", &false.encode());
 	assert!(res.is_ok());
 }
 
@@ -615,10 +615,10 @@ fn heap_is_reset_between_calls(wasm_method: WasmExecutionMethod) {
 		.expect("`__heap_base` is an `i32`");
 
 	let params = (heap_base as u32, 512u32 * 64 * 1024).encode();
-	instance.call("check_and_set_in_heap", &params).unwrap();
+	instance.call_export("check_and_set_in_heap", &params).unwrap();
 
 	// Cal it a second time to check that the heap was freed.
-	instance.call("check_and_set_in_heap", &params).unwrap();
+	instance.call_export("check_and_set_in_heap", &params).unwrap();
 }
 
 #[test_case(WasmExecutionMethod::Interpreted)]

client/executor/src/native_executor.rs

@@ -29,7 +29,7 @@ use sp_core::{
 use log::trace;
 use std::{result, panic::{UnwindSafe, AssertUnwindSafe}, sync::Arc, collections::HashMap};
 use sp_wasm_interface::{HostFunctions, Function};
-use sc_executor_common::wasm_runtime::{WasmInstance, WasmModule};
+use sc_executor_common::wasm_runtime::{WasmInstance, WasmModule, InvokeLocation};
 use sp_externalities::ExternalitiesExt as _;
 use sp_io::RuntimeSpawnExt;
 
@@ -189,7 +189,7 @@ impl sp_core::traits::CallInWasm for WasmExecutor {
 					&mut **ext,
 					move || {
 						RuntimeInstanceSpawn::register_on_externalities(module.clone());
-						instance.call(method, call_data)
+						instance.call(InvokeLocation::Export(method), call_data)
 					}
 				)
 			}).map_err(|e| e.to_string())
@@ -214,7 +214,7 @@ impl sp_core::traits::CallInWasm for WasmExecutor {
 				&mut **ext,
 				move || {
 					RuntimeInstanceSpawn::register_on_externalities(module.clone());
-					instance.call(method, call_data)
+					instance.call(InvokeLocation::Export(method), call_data)
 				}
 			)
 			.and_then(|r| r)
@@ -295,9 +295,7 @@ pub struct RuntimeInstanceSpawn {
 }
 
 impl sp_io::RuntimeSpawn for RuntimeInstanceSpawn {
-	fn dyn_dispatch(&self, func: u32, data: Vec<u8>) -> u32 {
-		use codec::Encode as _;
-
+	fn dyn_dispatch(&self, dispatcher_ref: u32, func: u32, data: Vec<u8>) -> u32 {
 		let new_handle = self.counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
 
 		let (sender, receiver) = std::sync::mpsc::channel();
@@ -317,10 +315,10 @@ impl sp_io::RuntimeSpawn for RuntimeInstanceSpawn {
 					// pool of istances should be used.
 					let instance = module.new_instance().expect("Failed to create new instance for fork");
 
-					let mut dispatch_data = Vec::new();
-					func.encode_to(&mut dispatch_data);
-					data.encode_to(&mut dispatch_data);
-					instance.call("dyn_dispatch", &dispatch_data[..]).expect("Failed to invoke instance.")
+					instance.call(
+						InvokeLocation::TableWithWrapper { dispatcher_ref, func },
+						&data[..],
+					).expect("Failed to invoke instance.")
 				}
 			);
 
@@ -430,7 +428,7 @@ impl<D: NativeExecutionDispatch + 'static> CodeExecutor for NativeExecutor<D> {
 							&mut **ext,
 							move || {
 								RuntimeInstanceSpawn::register_on_externalities(module.clone());
-								instance.call(method, data).map(NativeOrEncoded::Encoded)
+								instance.call(InvokeLocation::Export(method), data).map(NativeOrEncoded::Encoded)
 							}
 						)
 					},

client/executor/src/wasm_runtime.rs

@@ -338,7 +338,7 @@ fn create_versioned_wasm_runtime(
 		let runtime = AssertUnwindSafe(runtime.as_ref());
 		crate::native_executor::with_externalities_safe(
 			&mut **ext,
-			move || runtime.new_instance()?.call("Core_version", &[])
+			move || runtime.new_instance()?.call("Core_version".into(), &[])
 		).map_err(|_| WasmError::Instantiation("panic in call to get runtime version".into()))?
 	};
 	let version = match version_result {

client/executor/wasmi/src/lib.rs

@@ -19,7 +19,7 @@
 use std::{str, cell::RefCell, sync::Arc};
 use wasmi::{
 	Module, ModuleInstance, MemoryInstance, MemoryRef, TableRef, ImportsBuilder, ModuleRef,
-	memory_units::Pages,
+	FuncInstance, memory_units::Pages,
 	RuntimeValue::{I32, I64, self},
 };
 use codec::{Encode, Decode};
@@ -29,7 +29,7 @@ use sp_wasm_interface::{
 	FunctionContext, Pointer, WordSize, Sandbox, MemoryId, Result as WResult, Function,
 };
 use sp_runtime_interface::unpack_ptr_and_len;
-use sc_executor_common::wasm_runtime::{WasmModule, WasmInstance};
+use sc_executor_common::wasm_runtime::{WasmModule, WasmInstance, InvokeLocation};
 use sc_executor_common::{
 	error::{Error, WasmError},
 	sandbox,
@@ -434,7 +434,7 @@ fn get_heap_base(module: &ModuleRef) -> Result<u32, Error> {
 fn call_in_wasm_module(
 	module_instance: &ModuleRef,
 	memory: &MemoryRef,
-	method: &str,
+	method: InvokeLocation,
 	data: &[u8],
 	host_functions: &[&'static dyn Function],
 	allow_missing_func_imports: bool,
@@ -449,7 +449,7 @@ fn call_in_wasm_module(
 	let mut fec = FunctionExecutor::new(
 		memory.clone(),
 		heap_base,
-		table,
+		table.clone(),
 		host_functions,
 		allow_missing_func_imports,
 		missing_functions,
@@ -459,11 +459,36 @@ fn call_in_wasm_module(
 	let offset = fec.allocate_memory(data.len() as u32)?;
 	fec.write_memory(offset, data)?;
 
-	let result = module_instance.invoke_export(
-		method,
-		&[I32(u32::from(offset) as i32), I32(data.len() as i32)],
-		&mut fec,
-	);
+	let result = match method {
+		InvokeLocation::Export(method) => {
+			module_instance.invoke_export(
+				method,
+				&[I32(u32::from(offset) as i32), I32(data.len() as i32)],
+				&mut fec,
+			)
+		},
+		InvokeLocation::Table(func_ref) => {
+			let func = table.ok_or(Error::NoTable)?
+				.get(func_ref)?
+				.ok_or(Error::NoTableEntryWithIndex(func_ref))?;
+			FuncInstance::invoke(
+				&func,
+				&[I32(u32::from(offset) as i32), I32(data.len() as i32)],
+				&mut fec,
+			).map_err(Into::into)
+		},
+		InvokeLocation::TableWithWrapper { dispatcher_ref, func } => {
+			let dispatcher = table.ok_or(Error::NoTable)?
+				.get(dispatcher_ref)?
+				.ok_or(Error::NoTableEntryWithIndex(dispatcher_ref))?;
+
+			FuncInstance::invoke(
+				&dispatcher,
+				&[I32(func as _), I32(u32::from(offset) as i32), I32(data.len() as i32)],
+				&mut fec,
+			).map_err(Into::into)
+		},
+	};
 
 	match result {
 		Ok(Some(I64(r))) => {
@@ -677,7 +702,7 @@ pub struct WasmiInstance {
 unsafe impl Send for WasmiInstance {}
 
 impl WasmInstance for WasmiInstance {
-	fn call(&self, method: &str, data: &[u8]) -> Result<Vec<u8>, Error> {
+	fn call(&self, method: InvokeLocation, data: &[u8]) -> Result<Vec<u8>, Error> {
 		// We reuse a single wasm instance for multiple calls and a previous call (if any)
 		// altered the state. Therefore, we need to restore the instance to original state.
 

primitives/io/src/lib.rs

@@ -1138,7 +1138,7 @@ pub trait Sandbox {
 /// Runtime spawn extension.
 pub trait RuntimeSpawn : Send {
 	/// Create new runtime instance and use dynamic dispatch to invoke with specified payload.
-	fn dyn_dispatch(&self, func: u32, payload: Vec<u8>) -> u32;
+	fn dyn_dispatch(&self, dispatcher_ref: u32, func: u32, payload: Vec<u8>) -> u32;
 
 	/// Join the result of previously created runtime instance invocation.
 	fn join(&self, handle: u32) -> Vec<u8>;
@@ -1158,10 +1158,10 @@ pub trait RuntimeTasks {
 	/// Wasm host function for spawning task.
 	///
 	/// This should not be used directly. Use `sp_io::tasks::spawn` instead.
-	fn spawn(&mut self, entry: u32, payload: Vec<u8>) -> u32 {
+	fn spawn(&mut self, dispatcher_ref: u32, entry: u32, payload: Vec<u8>) -> u32 {
 		let runtime_spawn = self.extension::<RuntimeSpawnExt>()
 			.expect("Cannot spawn without dynamic runtime dispatcher (RuntimeSpawnExt)");
-		runtime_spawn.dyn_dispatch(entry, payload)
+		runtime_spawn.dyn_dispatch(dispatcher_ref, entry, payload)
 	}
 
 	/// Wasm host function for joining a task.

primitives/io/src/tasks.rs

@@ -17,7 +17,8 @@
 
 //! Runtime tasks.
 //!
-//! Contains runtime-usable interface for spawning parallel purely computational tasks.
+//! Contains runtime-usable functions for spawning parallel purely computational tasks.
+//!
 
 #[cfg(feature = "std")]
 mod inner {
@@ -60,40 +61,41 @@ mod inner {
 #[cfg(not(feature = "std"))]
 mod inner {
 
-	use sp_std::vec::Vec;
+	use sp_std::{vec::Vec, prelude::Box};
 
-	/// Spawn new runtime task (wasm).
-	pub fn spawn(entry_point: fn(Vec<u8>) -> Vec<u8>, payload: Vec<u8>) -> DataJoinHandle {
 
-		/// Dynamic dispatch of wasm blob.
-		///
-		/// Arguments are expected to be scale encoded in vector at address `payload_ptr` with length of
-		/// `payload_len`.
-		///
-		/// Arguments: function pointer (u32), input (Vec<u8>).
-		///
-		/// Function at pointer is expected to have signature of `(Vec<u8>) -> Vec<u8>`. Since this dynamic dispatch
-		/// function and the invoked function are compiled with the same compiler, there should be no problem with
-		/// ABI incompatibility.
-		#[no_mangle]
-		unsafe extern "C" fn dyn_dispatch(payload_ptr: u32, payload_len: u32) -> u64 {
-
-			use codec::Decode as _;
-
-			let mut data: &[u8] = core::slice::from_raw_parts(payload_ptr as usize as *const u8, payload_len as usize);
-			let entry = u32::decode(&mut data).expect("Failed to decode input") as usize;
-			let ptr: fn(Vec<u8>) -> Vec<u8> = core::mem::transmute(entry);
-			let payload = Vec::<u8>::decode(&mut data).expect("Failed to decode input");
-
-			let output = (ptr)(payload);
-
-			let mut output_encode = (output.len() as u64) << 32;
-			output_encode | (output.as_ptr() as usize as u64)
-		}
+	/// Dynamic dispatch of wasm blob.
+	///
+	/// Arguments are expected to be scale encoded in vector at address `payload_ptr` with length of
+	/// `payload_len`.
+	///
+	/// Arguments: function pointer (u32), input (Vec<u8>).
+	///
+	/// Function at pointer is expected to have signature of `(Vec<u8>) -> Vec<u8>`. Since this dynamic dispatch
+	/// function and the invoked function are compiled with the same compiler, there should be no problem with
+	/// ABI incompatibility.
+	extern "C" fn dispatch_wrapper(func_ref: u32, payload_ptr: u32, payload_len: u32) -> u64 {
+		let payload_len = payload_len as usize;
+		let output = unsafe {
+			let payload = Vec::from_raw_parts(payload_ptr as usize as *mut _, payload_len, payload_len);
+			let ptr: fn(Vec<u8>) -> Vec<u8> = core::mem::transmute(func_ref);
+			(ptr)(payload)
+		};
+		let mut output_encode = (output.len() as u64) << 32;
+		output_encode | (output.as_ptr() as usize as u64)
+	}
 
+	/// Spawn new runtime task (wasm).
+	pub fn spawn(entry_point: fn(Vec<u8>) -> Vec<u8>, payload: Vec<u8>) -> DataJoinHandle {
 		let func_ptr: usize = unsafe { core::mem::transmute(entry_point) };
 
-		let handle = unsafe { crate::runtime_tasks::spawn(func_ptr as u32, payload) };
+		let handle = unsafe {
+			crate::runtime_tasks::spawn(
+				(dispatch_wrapper as extern "C" fn(u32, u32, u32) -> u64 as usize) as u32,
+				func_ptr as u32,
+				payload,
+			)
+		};
 		DataJoinHandle { handle }
 	}