Refactor WASM module instantiation (#10480)

* Refactor WASM module instantiation; enable WASM instance pooling

* Disable the `uffd` feature on `wasmtime`

* Restore the original behavior regarding the initial WASM memory size

* Adjust error message

* Remove unnecessary import in the benchmarks

* Preinstantiate the WASM runtime for a slight speedup

* Delete the asserts in `convert_memory_import_into_export`

* `return` -> `break`

* Revert WASM instance pooling for now

* Have `convert_memory_import_into_export` return an error instead of panic

* Update the warning when an import is missing

* Rustfmt and clippy fix

* Fix executor benchmarks' compilation without `wasmtime` being enabled

* rustfmt again

* Align to review comments

* Extend tests so that both imported and exported memories are tested

* Increase the number of heap pages for exported memories too

* Fix `decommit_works` test

substrate/client/executor/wasmtime/src/imports.rs

@@ -16,37 +16,24 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-use crate::{
-	host::HostContext,
-	runtime::{Store, StoreData},
-};
+use crate::{host::HostContext, runtime::StoreData};
 use sc_executor_common::error::WasmError;
 use sp_wasm_interface::{FunctionContext, HostFunctions};
-use std::{collections::HashMap, convert::TryInto};
-use wasmtime::{Extern, ExternType, Func, FuncType, ImportType, Memory, MemoryType, Module, Trap};
+use std::collections::HashMap;
+use wasmtime::{ExternType, FuncType, ImportType, Linker, Module, Trap};
 
-pub struct Imports {
-	/// Contains the index into `externs` where the memory import is stored if any. `None` if there
-	/// is none.
-	pub memory_import_index: Option<usize>,
-	pub externs: Vec<Extern>,
-}
-
-/// Goes over all imports of a module and prepares a vector of `Extern`s that can be used for
-/// instantiation of the module. Returns an error if there are imports that cannot be satisfied.
-pub(crate) fn resolve_imports<H>(
-	store: &mut Store,
+/// Goes over all imports of a module and prepares the given linker for instantiation of the module.
+/// Returns an error if there are imports that cannot be satisfied.
+pub(crate) fn prepare_imports<H>(
+	linker: &mut Linker<StoreData>,
 	module: &Module,
-	heap_pages: u64,
 	allow_missing_func_imports: bool,
-) -> Result<Imports, WasmError>
+) -> Result<(), WasmError>
 where
 	H: HostFunctions,
 {
-	let mut externs = vec![];
-	let mut memory_import_index = None;
 	let mut pending_func_imports = HashMap::new();
-	for (index, import_ty) in module.imports().enumerate() {
+	for import_ty in module.imports() {
 		let name = import_name(&import_ty)?;
 
 		if import_ty.module() != "env" {
@@ -57,41 +44,36 @@ where
 			)))
 		}
 
-		if name == "memory" {
-			memory_import_index = Some(index);
-			externs.push((index, resolve_memory_import(store, &import_ty, heap_pages)?));
-			continue
-		}
-
 		match import_ty.ty() {
 			ExternType::Func(func_ty) => {
-				pending_func_imports.insert(name.to_owned(), (index, import_ty, func_ty));
+				pending_func_imports.insert(name.to_owned(), (import_ty, func_ty));
 			},
 			_ =>
 				return Err(WasmError::Other(format!(
-					"host doesn't provide any non function imports besides 'memory': {}:{}",
+					"host doesn't provide any non function imports: {}:{}",
 					import_ty.module(),
 					name,
 				))),
 		};
 	}
 
-	let mut registry = Registry { store, externs, pending_func_imports };
-
+	let mut registry = Registry { linker, pending_func_imports };
 	H::register_static(&mut registry)?;
-	let mut externs = registry.externs;
 
 	if !registry.pending_func_imports.is_empty() {
 		if allow_missing_func_imports {
-			for (_, (index, import_ty, func_ty)) in registry.pending_func_imports {
-				externs.push((
-					index,
-					MissingHostFuncHandler::new(&import_ty)?.into_extern(store, &func_ty),
-				));
+			for (name, (import_ty, func_ty)) in registry.pending_func_imports {
+				let error = format!("call to a missing function {}:{}", import_ty.module(), name);
+				log::debug!("Missing import: '{}' {:?}", name, func_ty);
+				linker
+					.func_new("env", &name, func_ty.clone(), move |_, _, _| {
+						Err(Trap::new(error.clone()))
+					})
+					.expect("adding a missing import stub can only fail when the item already exists, and it is missing here; qed");
 			}
 		} else {
 			let mut names = Vec::new();
-			for (name, (_, import_ty, _)) in registry.pending_func_imports {
+			for (name, (import_ty, _)) in registry.pending_func_imports {
 				names.push(format!("'{}:{}'", import_ty.module(), name));
 			}
 			let names = names.join(", ");
@@ -102,16 +84,12 @@ where
 		}
 	}
 
-	externs.sort_unstable_by_key(|&(index, _)| index);
-	let externs = externs.into_iter().map(|(_, ext)| ext).collect();
-
-	Ok(Imports { memory_import_index, externs })
+	Ok(())
 }
 
 struct Registry<'a, 'b> {
-	store: &'a mut Store,
-	externs: Vec<(usize, Extern)>,
-	pending_func_imports: HashMap<String, (usize, ImportType<'b>, FuncType)>,
+	linker: &'a mut Linker<StoreData>,
+	pending_func_imports: HashMap<String, (ImportType<'b>, FuncType)>,
 }
 
 impl<'a, 'b> sp_wasm_interface::HostFunctionRegistry for Registry<'a, 'b> {
@@ -131,9 +109,13 @@ impl<'a, 'b> sp_wasm_interface::HostFunctionRegistry for Registry<'a, 'b> {
 		fn_name: &str,
 		func: impl wasmtime::IntoFunc<Self::State, Params, Results>,
 	) -> Result<(), Self::Error> {
-		if let Some((index, _, _)) = self.pending_func_imports.remove(fn_name) {
-			let func = Func::wrap(&mut *self.store, func);
-			self.externs.push((index, Extern::Func(func)));
+		if self.pending_func_imports.remove(fn_name).is_some() {
+			self.linker.func_wrap("env", fn_name, func).map_err(|error| {
+				WasmError::Other(format!(
+					"failed to register host function '{}' with the WASM linker: {}",
+					fn_name, error
+				))
+			})?;
 		}
 
 		Ok(())
@@ -149,85 +131,3 @@ fn import_name<'a, 'b: 'a>(import: &'a ImportType<'b>) -> Result<&'a str, WasmEr
 	})?;
 	Ok(name)
 }
-
-fn resolve_memory_import(
-	store: &mut Store,
-	import_ty: &ImportType,
-	heap_pages: u64,
-) -> Result<Extern, WasmError> {
-	let requested_memory_ty = match import_ty.ty() {
-		ExternType::Memory(memory_ty) => memory_ty,
-		_ =>
-			return Err(WasmError::Other(format!(
-				"this import must be of memory type: {}:{}",
-				import_ty.module(),
-				import_name(&import_ty)?,
-			))),
-	};
-
-	// Increment the min (a.k.a initial) number of pages by `heap_pages` and check if it exceeds the
-	// maximum specified by the import.
-	let initial = requested_memory_ty.minimum().saturating_add(heap_pages);
-	if let Some(max) = requested_memory_ty.maximum() {
-		if initial > max {
-			return Err(WasmError::Other(format!(
-				"incremented number of pages by heap_pages (total={}) is more than maximum requested\
-				by the runtime wasm module {}",
-				initial,
-				max,
-			)))
-		}
-	}
-
-	// Note that the return value of `maximum` and `minimum`, while a u64,
-	// will always fit into a u32 for 32-bit memories.
-	// 64-bit memories are part of the memory64 proposal for WebAssembly which is not standardized
-	// yet.
-	let minimum: u32 = initial.try_into().map_err(|_| {
-		WasmError::Other(format!(
-			"minimum number of memory pages ({}) doesn't fit into u32",
-			initial
-		))
-	})?;
-	let maximum: Option<u32> = match requested_memory_ty.maximum() {
-		Some(max) => Some(max.try_into().map_err(|_| {
-			WasmError::Other(format!(
-				"maximum number of memory pages ({}) doesn't fit into u32",
-				max
-			))
-		})?),
-		None => None,
-	};
-
-	let memory_ty = MemoryType::new(minimum, maximum);
-	let memory = Memory::new(store, memory_ty).map_err(|e| {
-		WasmError::Other(format!(
-			"failed to create a memory during resolving of memory import: {}",
-			e,
-		))
-	})?;
-	Ok(Extern::Memory(memory))
-}
-
-/// A `Callable` handler for missing functions.
-struct MissingHostFuncHandler {
-	module: String,
-	name: String,
-}
-
-impl MissingHostFuncHandler {
-	fn new(import_ty: &ImportType) -> Result<Self, WasmError> {
-		Ok(Self {
-			module: import_ty.module().to_string(),
-			name: import_name(import_ty)?.to_string(),
-		})
-	}
-
-	fn into_extern(self, store: &mut Store, func_ty: &FuncType) -> Extern {
-		let Self { module, name } = self;
-		let func = Func::new(store, func_ty.clone(), move |_, _, _| {
-			Err(Trap::new(format!("call to a missing function {}:{}", module, name)))
-		});
-		Extern::Func(func)
-	}
-}

substrate/client/executor/wasmtime/src/instance_wrapper.rs

@@ -21,12 +21,13 @@
 
 use crate::runtime::{Store, StoreData};
 use sc_executor_common::{
-	error::{Backtrace, Error, MessageWithBacktrace, Result},
+	error::{Backtrace, Error, MessageWithBacktrace, Result, WasmError},
 	wasm_runtime::InvokeMethod,
 };
-use sp_wasm_interface::{HostFunctions, Pointer, Value, WordSize};
+use sp_wasm_interface::{Pointer, Value, WordSize};
 use wasmtime::{
-	AsContext, AsContextMut, Extern, Func, Global, Instance, Memory, Module, Table, Val,
+	AsContext, AsContextMut, Engine, Extern, Func, Global, Instance, InstancePre, Memory, Table,
+	Val,
 };
 
 /// Invoked entrypoint format.
@@ -162,62 +163,41 @@ fn extern_func(extern_: &Extern) -> Option<&Func> {
 	}
 }
 
+pub(crate) fn create_store(engine: &wasmtime::Engine, max_memory_size: Option<usize>) -> Store {
+	let limits = if let Some(max_memory_size) = max_memory_size {
+		wasmtime::StoreLimitsBuilder::new().memory_size(max_memory_size).build()
+	} else {
+		Default::default()
+	};
+
+	let mut store =
+		Store::new(engine, StoreData { limits, host_state: None, memory: None, table: None });
+	if max_memory_size.is_some() {
+		store.limiter(|s| &mut s.limits);
+	}
+	store
+}
+
 impl InstanceWrapper {
-	/// Create a new instance wrapper from the given wasm module.
-	pub fn new<H>(
-		module: &Module,
-		heap_pages: u64,
-		allow_missing_func_imports: bool,
+	pub(crate) fn new(
+		engine: &Engine,
+		instance_pre: &InstancePre<StoreData>,
 		max_memory_size: Option<usize>,
-	) -> Result<Self>
-	where
-		H: HostFunctions,
-	{
-		let limits = if let Some(max_memory_size) = max_memory_size {
-			wasmtime::StoreLimitsBuilder::new().memory_size(max_memory_size).build()
-		} else {
-			Default::default()
-		};
-
-		let mut store = Store::new(
-			module.engine(),
-			StoreData { limits, host_state: None, memory: None, table: None },
-		);
-		if max_memory_size.is_some() {
-			store.limiter(|s| &mut s.limits);
-		}
-
-		// Scan all imports, find the matching host functions, and create stubs that adapt arguments
-		// and results.
-		let imports = crate::imports::resolve_imports::<H>(
-			&mut store,
-			module,
-			heap_pages,
-			allow_missing_func_imports,
-		)?;
-
-		let instance = Instance::new(&mut store, module, &imports.externs)
-			.map_err(|e| Error::from(format!("cannot instantiate: {}", e)))?;
-
-		let memory = match imports.memory_import_index {
-			Some(memory_idx) => extern_memory(&imports.externs[memory_idx])
-				.expect("only memory can be at the `memory_idx`; qed")
-				.clone(),
-			None => {
-				let memory = get_linear_memory(&instance, &mut store)?;
-				if !memory.grow(&mut store, heap_pages).is_ok() {
-					return Err("failed top increase the linear memory size".into())
-				}
-				memory
-			},
-		};
+	) -> Result<Self> {
+		let mut store = create_store(engine, max_memory_size);
+		let instance = instance_pre.instantiate(&mut store).map_err(|error| {
+			WasmError::Other(
+				format!("failed to instantiate a new WASM module instance: {}", error,),
+			)
+		})?;
 
+		let memory = get_linear_memory(&instance, &mut store)?;
 		let table = get_table(&instance, &mut store);
 
 		store.data_mut().memory = Some(memory);
 		store.data_mut().table = table;
 
-		Ok(Self { instance, memory, store })
+		Ok(InstanceWrapper { instance, memory, store })
 	}
 
 	/// Resolves a substrate entrypoint by the given name.
@@ -435,8 +415,11 @@ impl InstanceWrapper {
 fn decommit_works() {
 	let engine = wasmtime::Engine::default();
 	let code = wat::parse_str("(module (memory (export \"memory\") 1 4))").unwrap();
-	let module = Module::new(&engine, code).unwrap();
-	let mut wrapper = InstanceWrapper::new::<()>(&module, 2, true, None).unwrap();
+	let module = wasmtime::Module::new(&engine, code).unwrap();
+	let linker = wasmtime::Linker::new(&engine);
+	let mut store = create_store(&engine, None);
+	let instance_pre = linker.instantiate_pre(&mut store, &module).unwrap();
+	let mut wrapper = InstanceWrapper::new(&engine, &instance_pre, None).unwrap();
 	unsafe { *wrapper.memory.data_ptr(&wrapper.store) = 42 };
 	assert_eq!(unsafe { *wrapper.memory.data_ptr(&wrapper.store) }, 42);
 	wrapper.decommit();

substrate/client/executor/wasmtime/src/runtime.rs

@@ -23,7 +23,6 @@ use crate::{
 	instance_wrapper::{EntryPoint, InstanceWrapper},
 	util,
 };
-use core::marker::PhantomData;
 
 use sc_allocator::FreeingBumpHeapAllocator;
 use sc_executor_common::{
@@ -80,35 +79,25 @@ impl StoreData {
 
 pub(crate) type Store = wasmtime::Store<StoreData>;
 
-enum Strategy<H> {
+enum Strategy {
 	FastInstanceReuse {
 		instance_wrapper: InstanceWrapper,
 		globals_snapshot: GlobalsSnapshot<wasmtime::Global>,
 		data_segments_snapshot: Arc<DataSegmentsSnapshot>,
 		heap_base: u32,
 	},
-	RecreateInstance(InstanceCreator<H>),
+	RecreateInstance(InstanceCreator),
 }
 
-struct InstanceCreator<H> {
-	module: Arc<wasmtime::Module>,
-	heap_pages: u64,
-	allow_missing_func_imports: bool,
+struct InstanceCreator {
+	engine: wasmtime::Engine,
+	instance_pre: Arc<wasmtime::InstancePre<StoreData>>,
 	max_memory_size: Option<usize>,
-	phantom: PhantomData<H>,
 }
 
-impl<H> InstanceCreator<H>
-where
-	H: HostFunctions,
-{
+impl InstanceCreator {
 	fn instantiate(&mut self) -> Result<InstanceWrapper> {
-		InstanceWrapper::new::<H>(
-			&*self.module,
-			self.heap_pages,
-			self.allow_missing_func_imports,
-			self.max_memory_size,
-		)
+		InstanceWrapper::new(&self.engine, &self.instance_pre, self.max_memory_size)
 	}
 }
 
@@ -144,23 +133,19 @@ struct InstanceSnapshotData {
 
 /// A `WasmModule` implementation using wasmtime to compile the runtime module to machine code
 /// and execute the compiled code.
-pub struct WasmtimeRuntime<H> {
-	module: Arc<wasmtime::Module>,
+pub struct WasmtimeRuntime {
+	engine: wasmtime::Engine,
+	instance_pre: Arc<wasmtime::InstancePre<StoreData>>,
 	snapshot_data: Option<InstanceSnapshotData>,
 	config: Config,
-	phantom: PhantomData<H>,
 }
 
-impl<H> WasmModule for WasmtimeRuntime<H>
-where
-	H: HostFunctions,
-{
+impl WasmModule for WasmtimeRuntime {
 	fn new_instance(&self) -> Result<Box<dyn WasmInstance>> {
 		let strategy = if let Some(ref snapshot_data) = self.snapshot_data {
-			let mut instance_wrapper = InstanceWrapper::new::<H>(
-				&self.module,
-				self.config.heap_pages,
-				self.config.allow_missing_func_imports,
+			let mut instance_wrapper = InstanceWrapper::new(
+				&self.engine,
+				&self.instance_pre,
 				self.config.max_memory_size,
 			)?;
 			let heap_base = instance_wrapper.extract_heap_base()?;
@@ -174,19 +159,17 @@ where
 				&mut InstanceGlobals { instance: &mut instance_wrapper },
 			);
 
-			Strategy::<H>::FastInstanceReuse {
+			Strategy::FastInstanceReuse {
 				instance_wrapper,
 				globals_snapshot,
 				data_segments_snapshot: snapshot_data.data_segments_snapshot.clone(),
 				heap_base,
 			}
 		} else {
-			Strategy::<H>::RecreateInstance(InstanceCreator {
-				module: self.module.clone(),
-				heap_pages: self.config.heap_pages,
-				allow_missing_func_imports: self.config.allow_missing_func_imports,
+			Strategy::RecreateInstance(InstanceCreator {
+				engine: self.engine.clone(),
+				instance_pre: self.instance_pre.clone(),
 				max_memory_size: self.config.max_memory_size,
-				phantom: PhantomData,
 			})
 		};
 
@@ -196,14 +179,11 @@ where
 
 /// A `WasmInstance` implementation that reuses compiled module and spawns instances
 /// to execute the compiled code.
-pub struct WasmtimeInstance<H> {
-	strategy: Strategy<H>,
+pub struct WasmtimeInstance {
+	strategy: Strategy,
 }
 
-impl<H> WasmInstance for WasmtimeInstance<H>
-where
-	H: HostFunctions,
-{
+impl WasmInstance for WasmtimeInstance {
 	fn call(&mut self, method: InvokeMethod, data: &[u8]) -> Result<Vec<u8>> {
 		match &mut self.strategy {
 			Strategy::FastInstanceReuse {
@@ -498,7 +478,7 @@ enum CodeSupplyMode<'a> {
 pub fn create_runtime<H>(
 	blob: RuntimeBlob,
 	config: Config,
-) -> std::result::Result<WasmtimeRuntime<H>, WasmError>
+) -> std::result::Result<WasmtimeRuntime, WasmError>
 where
 	H: HostFunctions,
 {
@@ -520,7 +500,7 @@ where
 pub unsafe fn create_runtime_from_artifact<H>(
 	compiled_artifact: &[u8],
 	config: Config,
-) -> std::result::Result<WasmtimeRuntime<H>, WasmError>
+) -> std::result::Result<WasmtimeRuntime, WasmError>
 where
 	H: HostFunctions,
 {
@@ -534,7 +514,7 @@ where
 unsafe fn do_create_runtime<H>(
 	code_supply_mode: CodeSupplyMode<'_>,
 	config: Config,
-) -> std::result::Result<WasmtimeRuntime<H>, WasmError>
+) -> std::result::Result<WasmtimeRuntime, WasmError>
 where
 	H: HostFunctions,
 {
@@ -550,27 +530,39 @@ where
 	}
 
 	let engine = Engine::new(&wasmtime_config)
-		.map_err(|e| WasmError::Other(format!("cannot create the engine for runtime: {}", e)))?;
+		.map_err(|e| WasmError::Other(format!("cannot create the wasmtime engine: {}", e)))?;
 
 	let (module, snapshot_data) = match code_supply_mode {
 		CodeSupplyMode::Verbatim { blob } => {
-			let blob = instrument(blob, &config.semantics)?;
+			let mut blob = instrument(blob, &config.semantics)?;
+
+			// We don't actually need the memory to be imported so we can just convert any memory
+			// import into an export with impunity. This simplifies our code since `wasmtime` will
+			// now automatically take care of creating the memory for us, and it also allows us
+			// to potentially enable `wasmtime`'s instance pooling at a later date. (Imported
+			// memories are ineligible for pooling.)
+			blob.convert_memory_import_into_export()?;
+			blob.add_extra_heap_pages_to_memory_section(
+				config
+					.heap_pages
+					.try_into()
+					.map_err(|e| WasmError::Other(format!("invalid `heap_pages`: {}", e)))?,
+			)?;
+
+			let serialized_blob = blob.clone().serialize();
+
+			let module = wasmtime::Module::new(&engine, &serialized_blob)
+				.map_err(|e| WasmError::Other(format!("cannot create module: {}", e)))?;
 
 			if config.semantics.fast_instance_reuse {
 				let data_segments_snapshot = DataSegmentsSnapshot::take(&blob).map_err(|e| {
 					WasmError::Other(format!("cannot take data segments snapshot: {}", e))
 				})?;
 				let data_segments_snapshot = Arc::new(data_segments_snapshot);
-
 				let mutable_globals = ExposedMutableGlobalsSet::collect(&blob);
 
-				let module = wasmtime::Module::new(&engine, &blob.serialize())
-					.map_err(|e| WasmError::Other(format!("cannot create module: {}", e)))?;
-
 				(module, Some(InstanceSnapshotData { data_segments_snapshot, mutable_globals }))
 			} else {
-				let module = wasmtime::Module::new(&engine, &blob.serialize())
-					.map_err(|e| WasmError::Other(format!("cannot create module: {}", e)))?;
 				(module, None)
 			}
 		},
@@ -584,7 +576,15 @@ where
 		},
 	};
 
-	Ok(WasmtimeRuntime { module: Arc::new(module), snapshot_data, config, phantom: PhantomData })
+	let mut linker = wasmtime::Linker::new(&engine);
+	crate::imports::prepare_imports::<H>(&mut linker, &module, config.allow_missing_func_imports)?;
+
+	let mut store = crate::instance_wrapper::create_store(module.engine(), config.max_memory_size);
+	let instance_pre = linker
+		.instantiate_pre(&mut store, &module)
+		.map_err(|e| WasmError::Other(format!("cannot preinstantiate module: {}", e)))?;
+
+	Ok(WasmtimeRuntime { engine, instance_pre: Arc::new(instance_pre), snapshot_data, config })
 }
 
 fn instrument(

substrate/client/executor/wasmtime/src/tests.rs

@@ -24,7 +24,7 @@ use std::sync::Arc;
 type HostFunctions = sp_io::SubstrateHostFunctions;
 
 struct RuntimeBuilder {
-	code: Option<&'static str>,
+	code: Option<String>,
 	fast_instance_reuse: bool,
 	canonicalize_nans: bool,
 	deterministic_stack: bool,
@@ -46,7 +46,7 @@ impl RuntimeBuilder {
 		}
 	}
 
-	fn use_wat(&mut self, code: &'static str) {
+	fn use_wat(&mut self, code: String) {
 		self.code = Some(code);
 	}
 
@@ -152,7 +152,7 @@ fn test_stack_depth_reaching() {
 
 	let runtime = {
 		let mut builder = RuntimeBuilder::new_on_demand();
-		builder.use_wat(TEST_GUARD_PAGE_SKIP);
+		builder.use_wat(TEST_GUARD_PAGE_SKIP.to_string());
 		builder.deterministic_stack(true);
 		builder.build()
 	};
@@ -168,10 +168,19 @@ fn test_stack_depth_reaching() {
 }
 
 #[test]
-fn test_max_memory_pages() {
+fn test_max_memory_pages_imported_memory() {
+	test_max_memory_pages(true);
+}
+
+#[test]
+fn test_max_memory_pages_exported_memory() {
+	test_max_memory_pages(false);
+}
+
+fn test_max_memory_pages(import_memory: bool) {
 	fn try_instantiate(
 		max_memory_size: Option<usize>,
-		wat: &'static str,
+		wat: String,
 	) -> Result<(), Box<dyn std::error::Error>> {
 		let runtime = {
 			let mut builder = RuntimeBuilder::new_on_demand();
@@ -184,31 +193,48 @@ fn test_max_memory_pages() {
 		Ok(())
 	}
 
+	fn memory(initial: u32, maximum: Option<u32>, import: bool) -> String {
+		let memory = if let Some(maximum) = maximum {
+			format!("(memory $0 {} {})", initial, maximum)
+		} else {
+			format!("(memory $0 {})", initial)
+		};
+
+		if import {
+			format!("(import \"env\" \"memory\" {})", memory)
+		} else {
+			format!("{}\n(export \"memory\" (memory $0))", memory)
+		}
+	}
+
 	const WASM_PAGE_SIZE: usize = 65536;
 
 	// check the old behavior if preserved. That is, if no limit is set we allow 4 GiB of memory.
 	try_instantiate(
 		None,
-		r#"
-		(module
-			;; we want to allocate the maximum number of pages supported in wasm for this test.
-			;;
-			;; However, due to a bug in wasmtime (I think wasmi is also affected) it is only possible
-			;; to allocate 65536 - 1 pages.
-			;;
-			;; Then, during creation of the Substrate Runtime instance, 1024 (heap_pages) pages are
-			;; mounted.
-			;;
-			;; Thus 65535 = 64511 + 1024
-			(import "env" "memory" (memory 64511))
-
-			(global (export "__heap_base") i32 (i32.const 0))
-			(func (export "main")
-				(param i32 i32) (result i64)
-				(i64.const 0)
+		format!(
+			r#"
+			(module
+				{}
+				(global (export "__heap_base") i32 (i32.const 0))
+				(func (export "main")
+					(param i32 i32) (result i64)
+					(i64.const 0)
+				)
 			)
-		)
-		"#,
+			"#,
+			/*
+				We want to allocate the maximum number of pages supported in wasm for this test.
+				However, due to a bug in wasmtime (I think wasmi is also affected) it is only possible
+				to allocate 65536 - 1 pages.
+
+				Then, during creation of the Substrate Runtime instance, 1024 (heap_pages) pages are
+				mounted.
+
+				Thus 65535 = 64511 + 1024
+			*/
+			memory(64511, None, import_memory)
+		),
 	)
 	.unwrap();
 
@@ -217,94 +243,104 @@ fn test_max_memory_pages() {
 	// max_memory_size = (1 (initial) + 1024 (heap_pages)) * WASM_PAGE_SIZE
 	try_instantiate(
 		Some((1 + 1024) * WASM_PAGE_SIZE),
-		r#"
-		(module
-
-			(import "env" "memory" (memory 1)) ;; <- 1 initial, max is not specified
-
-			(global (export "__heap_base") i32 (i32.const 0))
-			(func (export "main")
-				(param i32 i32) (result i64)
-				(i64.const 0)
+		format!(
+			r#"
+			(module
+				{}
+				(global (export "__heap_base") i32 (i32.const 0))
+				(func (export "main")
+					(param i32 i32) (result i64)
+					(i64.const 0)
+				)
 			)
-		)
-		"#,
+			"#,
+			// 1 initial, max is not specified.
+			memory(1, None, import_memory)
+		),
 	)
 	.unwrap();
 
 	// max is specified explicitly to 2048 pages.
 	try_instantiate(
 		Some((1 + 1024) * WASM_PAGE_SIZE),
-		r#"
-		(module
-
-			(import "env" "memory" (memory 1 2048)) ;; <- max is 2048
-
-			(global (export "__heap_base") i32 (i32.const 0))
-			(func (export "main")
-				(param i32 i32) (result i64)
-				(i64.const 0)
+		format!(
+			r#"
+			(module
+				{}
+				(global (export "__heap_base") i32 (i32.const 0))
+				(func (export "main")
+					(param i32 i32) (result i64)
+					(i64.const 0)
+				)
 			)
-		)
-		"#,
+			"#,
+			// Max is 2048.
+			memory(1, Some(2048), import_memory)
+		),
 	)
 	.unwrap();
 
 	// memory grow should work as long as it doesn't exceed 1025 pages in total.
 	try_instantiate(
 		Some((0 + 1024 + 25) * WASM_PAGE_SIZE),
-		r#"
-		(module
-			(import "env" "memory" (memory 0)) ;; <- zero starting pages.
+		format!(
+			r#"
+			(module
+				{}
+				(global (export "__heap_base") i32 (i32.const 0))
+				(func (export "main")
+					(param i32 i32) (result i64)
 
-			(global (export "__heap_base") i32 (i32.const 0))
-			(func (export "main")
-				(param i32 i32) (result i64)
-
-				;; assert(memory.grow returns != -1)
-				(if
-					(i32.eq
-						(memory.grow
-							(i32.const 25)
+					;; assert(memory.grow returns != -1)
+					(if
+						(i32.eq
+							(memory.grow
+								(i32.const 25)
+							)
+							(i32.const -1)
 						)
-						(i32.const -1)
+						(unreachable)
 					)
-					(unreachable)
-				)
 
-				(i64.const 0)
+					(i64.const 0)
+				)
 			)
-		)
-		"#,
+			"#,
+			// Zero starting pages.
+			memory(0, None, import_memory)
+		),
 	)
 	.unwrap();
 
 	// We start with 1025 pages and try to grow at least one.
 	try_instantiate(
 		Some((1 + 1024) * WASM_PAGE_SIZE),
-		r#"
-		(module
-			(import "env" "memory" (memory 1))  ;; <- initial=1, meaning after heap pages mount the
-												;; total will be already 1025
-			(global (export "__heap_base") i32 (i32.const 0))
-			(func (export "main")
-				(param i32 i32) (result i64)
+		format!(
+			r#"
+			(module
+				{}
+				(global (export "__heap_base") i32 (i32.const 0))
+				(func (export "main")
+					(param i32 i32) (result i64)
 
-				;; assert(memory.grow returns == -1)
-				(if
-					(i32.ne
-						(memory.grow
-							(i32.const 1)
+					;; assert(memory.grow returns == -1)
+					(if
+						(i32.ne
+							(memory.grow
+								(i32.const 1)
+							)
+							(i32.const -1)
 						)
-						(i32.const -1)
+						(unreachable)
 					)
-					(unreachable)
-				)
 
-				(i64.const 0)
+					(i64.const 0)
+				)
 			)
-		)
-		"#,
+			"#,
+			// Initial=1, meaning after heap pages mount the total will be already 1025.
+			memory(1, None, import_memory)
+		),
 	)
 	.unwrap();
 }