// This file is part of Substrate.

// Copyright (C) 2021-2022 Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.

// 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 codec::{Decode as _, Encode as _};
use sc_executor_common::{error::Error, runtime_blob::RuntimeBlob, wasm_runtime::WasmModule};
use sc_runtime_test::wasm_binary_unwrap;

use crate::InstantiationStrategy;

type HostFunctions = sp_io::SubstrateHostFunctions;

#[macro_export]
macro_rules! test_wasm_execution {
	(@no_legacy_instance_reuse $method_name:ident) => {
		paste::item! {
			#[test]
			fn [<$method_name _recreate_instance_cow>]() {
				$method_name(
					InstantiationStrategy::RecreateInstanceCopyOnWrite
				);
			}

			#[test]
			fn [<$method_name _recreate_instance_vanilla>]() {
				$method_name(
					InstantiationStrategy::RecreateInstance
				);
			}

			#[test]
			fn [<$method_name _pooling_cow>]() {
				$method_name(
					InstantiationStrategy::PoolingCopyOnWrite
				);
			}

			#[test]
			fn [<$method_name _pooling_vanilla>]() {
				$method_name(
					InstantiationStrategy::Pooling
				);
			}
		}
	};

	($method_name:ident) => {
		test_wasm_execution!(@no_legacy_instance_reuse $method_name);

		paste::item! {
			#[test]
			fn [<$method_name _legacy_instance_reuse>]() {
				$method_name(
					InstantiationStrategy::LegacyInstanceReuse
				);
			}
		}
	};
}

struct RuntimeBuilder {
	code: Option<String>,
	instantiation_strategy: InstantiationStrategy,
	canonicalize_nans: bool,
	deterministic_stack: bool,
	extra_heap_pages: u64,
	max_memory_size: Option<usize>,
	precompile_runtime: bool,
	tmpdir: Option<tempfile::TempDir>,
}

impl RuntimeBuilder {
	fn new(instantiation_strategy: InstantiationStrategy) -> Self {
		Self {
			code: None,
			instantiation_strategy,
			canonicalize_nans: false,
			deterministic_stack: false,
			extra_heap_pages: 1024,
			max_memory_size: None,
			precompile_runtime: false,
			tmpdir: None,
		}
	}

	fn use_wat(mut self, code: String) -> Self {
		self.code = Some(code);
		self
	}

	fn canonicalize_nans(mut self, canonicalize_nans: bool) -> Self {
		self.canonicalize_nans = canonicalize_nans;
		self
	}

	fn deterministic_stack(mut self, deterministic_stack: bool) -> Self {
		self.deterministic_stack = deterministic_stack;
		self
	}

	fn precompile_runtime(mut self, precompile_runtime: bool) -> Self {
		self.precompile_runtime = precompile_runtime;
		self
	}

	fn max_memory_size(mut self, max_memory_size: Option<usize>) -> Self {
		self.max_memory_size = max_memory_size;
		self
	}

	fn build(&mut self) -> impl WasmModule + '_ {
		let blob = {
			let wasm: Vec<u8>;

			let wasm = match self.code {
				None => wasm_binary_unwrap(),
				Some(ref wat) => {
					wasm = wat::parse_str(wat).expect("wat parsing failed");
					&wasm
				},
			};

			RuntimeBlob::uncompress_if_needed(&wasm)
				.expect("failed to create a runtime blob out of test runtime")
		};

		let config = crate::Config {
			allow_missing_func_imports: true,
			cache_path: None,
			semantics: crate::Semantics {
				instantiation_strategy: self.instantiation_strategy,
				deterministic_stack_limit: match self.deterministic_stack {
					true => Some(crate::DeterministicStackLimit {
						logical_max: 65536,
						native_stack_max: 256 * 1024 * 1024,
					}),
					false => None,
				},
				canonicalize_nans: self.canonicalize_nans,
				parallel_compilation: true,
				extra_heap_pages: self.extra_heap_pages,
				max_memory_size: self.max_memory_size,
			},
		};

		if self.precompile_runtime {
			let dir = tempfile::tempdir().unwrap();
			let path = dir.path().join("runtime.bin");

			// Delay the removal of the temporary directory until we're dropped.
			self.tmpdir = Some(dir);

			let artifact = crate::prepare_runtime_artifact(blob, &config.semantics).unwrap();
			std::fs::write(&path, artifact).unwrap();
			unsafe { crate::create_runtime_from_artifact::<HostFunctions>(&path, config) }
		} else {
			crate::create_runtime::<HostFunctions>(blob, config)
		}
		.expect("cannot create runtime")
	}
}

fn deep_call_stack_wat(depth: usize) -> String {
	format!(
		r#"
			(module
			  (memory $0 32)
			  (export "memory" (memory $0))
			  (global (export "__heap_base") i32 (i32.const 0))
			  (func (export "overflow") call 0)

			  (func $overflow (param $0 i32)
			    (block $label$1
			      (br_if $label$1
			        (i32.ge_u
			          (local.get $0)
			          (i32.const {depth})
			        )
			      )
			      (call $overflow
			        (i32.add
			          (local.get $0)
			          (i32.const 1)
			        )
			      )
			    )
			  )

			  (func (export "main")
			    (param i32 i32) (result i64)
			    (call $overflow (i32.const 0))
			    (i64.const 0)
			  )
			)
		"#
	)
}

// These two tests ensure that the `wasmtime`'s stack size limit and the amount of
// stack space used by a single stack frame doesn't suddenly change without us noticing.
//
// If they do (e.g. because we've pulled in a new version of `wasmtime`) we want to know
// that it did, regardless of how small the change was.
//
// If these tests starting failing it doesn't necessarily mean that something is broken;
// what it means is that one (or multiple) of the following has to be done:
//   a) the tests may need to be updated for the new call depth,
//   b) the stack limit may need to be changed to maintain backwards compatibility,
//   c) the root cause of the new call depth limit determined, and potentially fixed,
//   d) the new call depth limit may need to be validated to ensure it doesn't prevent any
//      existing chain from syncing (if it was effectively decreased)

// We need two limits here since depending on whether the code is compiled in debug
// or in release mode the maximum call depth is slightly different.
const CALL_DEPTH_LOWER_LIMIT: usize = 65478;
const CALL_DEPTH_UPPER_LIMIT: usize = 65514;

test_wasm_execution!(test_consume_under_1mb_of_stack_does_not_trap);
fn test_consume_under_1mb_of_stack_does_not_trap(instantiation_strategy: InstantiationStrategy) {
	let wat = deep_call_stack_wat(CALL_DEPTH_LOWER_LIMIT);
	let mut builder = RuntimeBuilder::new(instantiation_strategy).use_wat(wat);
	let runtime = builder.build();
	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");
	instance.call_export("main", &[]).unwrap();
}

test_wasm_execution!(test_consume_over_1mb_of_stack_does_trap);
fn test_consume_over_1mb_of_stack_does_trap(instantiation_strategy: InstantiationStrategy) {
	let wat = deep_call_stack_wat(CALL_DEPTH_UPPER_LIMIT + 1);
	let mut builder = RuntimeBuilder::new(instantiation_strategy).use_wat(wat);
	let runtime = builder.build();
	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");
	match instance.call_export("main", &[]).unwrap_err() {
		Error::AbortedDueToTrap(error) => {
			let expected = "wasm trap: call stack exhausted";
			assert_eq!(error.message, expected);
		},
		error => panic!("unexpected error: {:?}", error),
	}
}

test_wasm_execution!(test_nan_canonicalization);
fn test_nan_canonicalization(instantiation_strategy: InstantiationStrategy) {
	let mut builder = RuntimeBuilder::new(instantiation_strategy).canonicalize_nans(true);
	let runtime = builder.build();

	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");

	/// A NaN with canonical payload bits.
	const CANONICAL_NAN_BITS: u32 = 0x7fc00000;
	/// A NaN value with an abitrary payload.
	const ARBITRARY_NAN_BITS: u32 = 0x7f812345;

	// This test works like this: we essentially do
	//
	//     a + b
	//
	// where
	//
	//     * a is a nan with arbitrary bits in its payload
	//     * b is 1.
	//
	// according to the wasm spec, if one of the inputs to the operation is a non-canonical NaN
	// then the value be a NaN with non-deterministic payload bits.
	//
	// However, with the `canonicalize_nans` option turned on above, we expect that the output will
	// be a canonical NaN.
	//
	// We exterpolate the results of this tests so that we assume that all intermediate computations
	// that involve floats are sanitized and cannot produce a non-deterministic NaN.

	let params = (u32::to_le_bytes(ARBITRARY_NAN_BITS), u32::to_le_bytes(1)).encode();
	let res = {
		let raw_result = instance.call_export("test_fp_f32add", &params).unwrap();
		u32::from_le_bytes(<[u8; 4]>::decode(&mut &raw_result[..]).unwrap())
	};
	assert_eq!(res, CANONICAL_NAN_BITS);
}

test_wasm_execution!(test_stack_depth_reaching);
fn test_stack_depth_reaching(instantiation_strategy: InstantiationStrategy) {
	const TEST_GUARD_PAGE_SKIP: &str = include_str!("test-guard-page-skip.wat");

	let mut builder = RuntimeBuilder::new(instantiation_strategy)
		.use_wat(TEST_GUARD_PAGE_SKIP.to_string())
		.deterministic_stack(true);

	let runtime = builder.build();
	let mut instance = runtime.new_instance().expect("failed to instantiate a runtime");

	match instance.call_export("test-many-locals", &[]).unwrap_err() {
		Error::AbortedDueToTrap(error) => {
			let expected = "wasm trap: wasm `unreachable` instruction executed";
			assert_eq!(error.message, expected);
		},
		error => panic!("unexpected error: {:?}", error),
	}
}

test_wasm_execution!(test_max_memory_pages_imported_memory_without_precompilation);
fn test_max_memory_pages_imported_memory_without_precompilation(
	instantiation_strategy: InstantiationStrategy,
) {
	test_max_memory_pages(instantiation_strategy, true, false);
}

test_wasm_execution!(test_max_memory_pages_exported_memory_without_precompilation);
fn test_max_memory_pages_exported_memory_without_precompilation(
	instantiation_strategy: InstantiationStrategy,
) {
	test_max_memory_pages(instantiation_strategy, false, false);
}

test_wasm_execution!(@no_legacy_instance_reuse test_max_memory_pages_imported_memory_with_precompilation);
fn test_max_memory_pages_imported_memory_with_precompilation(
	instantiation_strategy: InstantiationStrategy,
) {
	test_max_memory_pages(instantiation_strategy, true, true);
}

test_wasm_execution!(@no_legacy_instance_reuse test_max_memory_pages_exported_memory_with_precompilation);
fn test_max_memory_pages_exported_memory_with_precompilation(
	instantiation_strategy: InstantiationStrategy,
) {
	test_max_memory_pages(instantiation_strategy, false, true);
}

fn test_max_memory_pages(
	instantiation_strategy: InstantiationStrategy,
	import_memory: bool,
	precompile_runtime: bool,
) {
	fn try_instantiate(
		max_memory_size: Option<usize>,
		wat: String,
		instantiation_strategy: InstantiationStrategy,
		precompile_runtime: bool,
	) -> Result<(), Box<dyn std::error::Error>> {
		let mut builder = RuntimeBuilder::new(instantiation_strategy)
			.use_wat(wat)
			.max_memory_size(max_memory_size)
			.precompile_runtime(precompile_runtime);

		let runtime = builder.build();
		let mut instance = runtime.new_instance()?;
		let _ = instance.call_export("main", &[])?;
		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,
		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)
		),
		instantiation_strategy,
		precompile_runtime,
	)
	.unwrap();

	// max is not specified, therefore it's implied to be 65536 pages (4 GiB).
	//
	// max_memory_size = (1 (initial) + 1024 (heap_pages)) * WASM_PAGE_SIZE
	try_instantiate(
		Some((1 + 1024) * WASM_PAGE_SIZE),
		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)
		),
		instantiation_strategy,
		precompile_runtime,
	)
	.unwrap();

	// max is specified explicitly to 2048 pages.
	try_instantiate(
		Some((1 + 1024) * WASM_PAGE_SIZE),
		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)
		),
		instantiation_strategy,
		precompile_runtime,
	)
	.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),
		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.eq
							(memory.grow
								(i32.const 25)
							)
							(i32.const -1)
						)
						(unreachable)
					)

					(i64.const 0)
				)
			)
			"#,
			// Zero starting pages.
			memory(0, None, import_memory)
		),
		instantiation_strategy,
		precompile_runtime,
	)
	.unwrap();

	// We start with 1025 pages and try to grow at least one.
	try_instantiate(
		Some((1 + 1024) * WASM_PAGE_SIZE),
		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)
							)
							(i32.const -1)
						)
						(unreachable)
					)

					(i64.const 0)
				)
			)
			"#,
			// Initial=1, meaning after heap pages mount the total will be already 1025.
			memory(1, None, import_memory)
		),
		instantiation_strategy,
		precompile_runtime,
	)
	.unwrap();
}

// This test takes quite a while to execute in a debug build (over 6 minutes on a TR 3970x)
// so it's ignored by default unless it was compiled with `--release`.
#[cfg_attr(build_type = "debug", ignore)]
#[test]
fn test_instances_without_reuse_are_not_leaked() {
	let runtime = crate::create_runtime::<HostFunctions>(
		RuntimeBlob::uncompress_if_needed(wasm_binary_unwrap()).unwrap(),
		crate::Config {
			allow_missing_func_imports: true,
			cache_path: None,
			semantics: crate::Semantics {
				instantiation_strategy: InstantiationStrategy::RecreateInstance,
				deterministic_stack_limit: None,
				canonicalize_nans: false,
				parallel_compilation: true,
				extra_heap_pages: 2048,
				max_memory_size: None,
			},
		},
	)
	.unwrap();

	// As long as the `wasmtime`'s `Store` lives the instances spawned through it
	// will live indefinitely. Currently it has a maximum limit of 10k instances,
	// so let's spawn 10k + 1 of them to make sure our code doesn't keep the `Store`
	// alive longer than it is necessary. (And since we disabled instance reuse
	// a new instance will be spawned on each call.)
	let mut instance = runtime.new_instance().unwrap();
	for _ in 0..10001 {
		instance.call_export("test_empty_return", &[0]).unwrap();
	}
}