Flat Stack (#98)

* Define Instruction Set.

* WIP

* WIP 2

* Tests

* Working

* Bunch of other tests.

* WIP

* WIP

* Use Vec instead of VecDeque.

* Calibrate the limits.

* Clean

* Clean

* Another round of cleaning.

* Ignore traces.

* Optimize value stack

* Optimize a bit more.

* Cache memory index.

* Inline always instruction dispatch function.

* Comments.

* Clean

* Clean

* Use vector to keep unresolved references.

* Estimate resulting size.

* do refactoring

* Validate the locals count in the begging

* Introduce Keep and DropKeep structs in isa

* Rename/Split Validator into Reader

* Document stack layout

* Remove println!

* Fix typo.

* Use .last / .last_mut in stack

* Update docs for BrTable.

* Review fixes.

* Merge.

* Add an assert that stack is empty after the exec

benches/.gitignore

@@ -1 +1,3 @@
 /target
+*.trace
+

benches/Cargo.toml

@@ -7,3 +7,6 @@ authors = ["Sergey Pepyakin <s.pepyakin@gmail.com>"]
 wasmi = { path = ".." }
 assert_matches = "1.2"
 wabt = "0.3"
+
+[profile.bench]
+debug = true

src/common/mod.rs

@@ -1,4 +1,3 @@
-use parity_wasm::elements::BlockType;
 
 pub mod stack;
 
@@ -7,38 +6,4 @@ pub const DEFAULT_MEMORY_INDEX: u32 = 0;
 /// Index of default table.
 pub const DEFAULT_TABLE_INDEX: u32 = 0;
 
-/// Control stack frame.
-#[derive(Debug, Clone)]
-pub struct BlockFrame {
-	/// Frame type.
-	pub frame_type: BlockFrameType,
-	/// A signature, which is a block signature type indicating the number and types of result values of the region.
-	pub block_type: BlockType,
-	/// A label for reference to block instruction.
-	pub begin_position: usize,
-	/// A label for reference from branch instructions.
-	pub branch_position: usize,
-	/// A label for reference from end instructions.
-	pub end_position: usize,
-	/// A limit integer value, which is an index into the value stack indicating where to reset it to on a branch to that label.
-	pub value_stack_len: usize,
-	/// Boolean which signals whether value stack became polymorphic. Value stack starts in non-polymorphic state and
-	/// becomes polymorphic only after an instruction that never passes control further is executed,
-	/// i.e. `unreachable`, `br` (but not `br_if`!), etc.
-	pub polymorphic_stack: bool,
-}
-
-/// Type of block frame.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum BlockFrameType {
-	/// Function frame.
-	Function,
-	/// Usual block frame.
-	Block,
-	/// Loop frame (branching to the beginning of block).
-	Loop,
-	/// True-subblock of if expression.
-	IfTrue,
-	/// False-subblock of if expression.
-	IfFalse,
-}
+// TODO: Move BlockFrame under validation.

src/common/stack.rs

@@ -1,5 +1,4 @@
 
-use std::collections::VecDeque;
 use std::error;
 use std::fmt;
 
@@ -22,7 +21,7 @@ impl error::Error for Error {
 #[derive(Debug)]
 pub struct StackWithLimit<T> where T: Clone {
 	/// Stack values.
-	values: VecDeque<T>,
+	values: Vec<T>,
 	/// Stack limit (maximal stack len).
 	limit: usize,
 }
@@ -30,7 +29,7 @@ pub struct StackWithLimit<T> where T: Clone {
 impl<T> StackWithLimit<T> where T: Clone {
 	pub fn with_limit(limit: usize) -> Self {
 		StackWithLimit {
-			values: VecDeque::new(),
+			values: Vec::new(),
 			limit: limit
 		}
 	}
@@ -43,19 +42,15 @@ impl<T> StackWithLimit<T> where T: Clone {
 		self.values.len()
 	}
 
-	pub fn limit(&self) -> usize {
-		self.limit
-	}
-
 	pub fn top(&self) -> Result<&T, Error> {
 		self.values
-			.back()
+			.last()
 			.ok_or_else(|| Error("non-empty stack expected".into()))
 	}
 
 	pub fn top_mut(&mut self) -> Result<&mut T, Error> {
 		self.values
-			.back_mut()
+			.last_mut()
 			.ok_or_else(|| Error("non-empty stack expected".into()))
 	}
 
@@ -72,13 +67,13 @@ impl<T> StackWithLimit<T> where T: Clone {
 			return Err(Error(format!("exceeded stack limit {}", self.limit)));
 		}
 
-		self.values.push_back(value);
+		self.values.push(value);
 		Ok(())
 	}
 
 	pub fn pop(&mut self) -> Result<T, Error> {
 		self.values
-			.pop_back()
+			.pop()
 			.ok_or_else(|| Error("non-empty stack expected".into()))
 	}
 

src/func.rs

@@ -1,12 +1,12 @@
 use std::rc::{Rc, Weak};
 use std::fmt;
-use std::collections::HashMap;
-use parity_wasm::elements::{Local, Instructions};
+use parity_wasm::elements::Local;
 use {Trap, TrapKind, Signature};
 use host::Externals;
 use runner::{check_function_args, Interpreter};
 use value::RuntimeValue;
 use module::ModuleInstance;
+use isa;
 
 /// Reference to a function (See [`FuncInstance`] for details).
 ///
@@ -158,6 +158,5 @@ impl FuncInstance {
 #[derive(Clone, Debug)]
 pub struct FuncBody {
 	pub locals: Vec<Local>,
-	pub instructions: Instructions,
-	pub labels: HashMap<usize, usize>,
+	pub code: isa::Instructions,
 }

src/isa.rs

@@ -0,0 +1,304 @@
+//! An instruction set used by wasmi.
+//!
+//! The instruction set is mostly derived from Wasm. However,
+//! there is a substantial difference.
+//!
+//! # Structured Stack Machine vs Plain One
+//!
+//! Wasm is a structured stack machine. Wasm encodes control flow in structures
+//! similar to that commonly found in a programming languages
+//! such as if, while. That contrasts to a plain stack machine which
+//!  encodes all control flow with goto-like instructions.
+//!
+//! Structured stack machine code aligns well with goals of Wasm,
+//! namely providing fast validation of Wasm code and compilation to native code.
+//!
+//! Unfortunately, the downside of structured stack machine code is
+//! that it is less convenient to interpret. For example, let's look at
+//! the following example in hypothetical structured stack machine:
+//!
+//! ```plain
+//! loop
+//!   ...
+//!   if_true_jump_to_end
+//!   ...
+//! end
+//! ```
+//!
+//! To execute `if_true_jump_to_end` , the interpreter needs to skip all instructions
+//! until it reaches the *matching* `end`. That's quite inefficient compared
+//! to a plain goto to the specific position.
+//!
+//! Because of this, the translation from the Wasm structured stack machine into a
+//! plain one is taking place.
+//!
+//! # Locals
+//!
+//! In a plain stack machine local variables and arguments live on the stack. Instead of
+//! accessing predifined locals slots in a plain stack machine locals are addressed relative
+//! to the current stack pointer. Because of this instead of taking an index of a local
+//! in {get,set,tee}_local operations, they take a relative depth as immediate. This works
+//! because at each instruction we always know the current stack height.
+//!
+//! Roughly, the stack layout looks like this
+//!
+//! | caller arguments |
+//! |  - arg 1         |
+//! |  - arg 2         |
+//! +------------------+
+//! | callee locals    |
+//! |  - var 1         |
+//! |  - var 2         |
+//! +------------------+
+//! | operands         |
+//! |  - op 1          |
+//! |  - op 2          |
+//! |                  |  <-- current stack pointer
+//! +------------------+
+//!
+//! # Differences from Wasm
+//!
+//! - There is no `nop` instruction.
+//! - All control flow strucutres are flattened to plain gotos.
+//! - Implicit returns via reaching function scope `End` are replaced with an explicit `return` instruction.
+//! - Locals live on the value stack now.
+//! - Load/store instructions doesn't take `align` parameter.
+//! - *.const store value in straight encoding.
+//! - Reserved immediates are ignored for `call_indirect`, `current_memory`, `grow_memory`.
+//!
+
+/// Should we keep a value before "discarding" a stack frame?
+///
+/// Note that this is a `enum` since Wasm doesn't support multiple return
+/// values at the moment.
+#[repr(u8)]
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum Keep {
+	None,
+	/// Pop one value from the yet-to-be-discarded stack frame to the
+	/// current stack frame.
+	Single,
+}
+
+/// Specifies how many values we should keep and how many we should drop.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub struct DropKeep {
+	pub drop: u32,
+	pub keep: Keep,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Target {
+	pub dst_pc: u32,
+	pub drop_keep: DropKeep,
+}
+
+#[allow(unused)] // TODO: Remove
+#[derive(Debug, Clone, PartialEq)]
+pub enum Instruction {
+	/// Push a local variable or an argument from the specified depth.
+	GetLocal(u32),
+
+	/// Pop a value and put it in at the specified depth.
+	SetLocal(u32),
+
+	/// Copy a value to the specified depth.
+	TeeLocal(u32),
+
+	/// Similar to the Wasm ones, but instead of a label depth
+	/// they specify direct PC.
+	Br(Target),
+	BrIfEqz(Target),
+	BrIfNez(Target),
+
+	/// br_table [t1 t2 t3 .. tn] tdefault
+	///
+	/// Pops the value from the stack. Then this value is used as an index
+	/// to the branch table.
+	///
+	/// However, the last target represents the default target. So if the index
+	/// is greater than length of the branch table, then the last index will be used.
+	///
+	/// Validation ensures that there should be at least one target.
+	BrTable(Box<[Target]>),
+
+	Unreachable,
+	Return(DropKeep),
+
+	Call(u32),
+	CallIndirect(u32),
+
+	Drop,
+	Select,
+
+	GetGlobal(u32),
+	SetGlobal(u32),
+
+	I32Load(u32),
+	I64Load(u32),
+	F32Load(u32),
+	F64Load(u32),
+	I32Load8S(u32),
+	I32Load8U(u32),
+	I32Load16S(u32),
+	I32Load16U(u32),
+	I64Load8S(u32),
+	I64Load8U(u32),
+	I64Load16S(u32),
+	I64Load16U(u32),
+	I64Load32S(u32),
+	I64Load32U(u32),
+	I32Store(u32),
+	I64Store(u32),
+	F32Store(u32),
+	F64Store(u32),
+	I32Store8(u32),
+	I32Store16(u32),
+	I64Store8(u32),
+	I64Store16(u32),
+	I64Store32(u32),
+
+	CurrentMemory,
+	GrowMemory,
+
+	I32Const(i32),
+	I64Const(i64),
+	F32Const(u32),
+	F64Const(u64),
+
+	I32Eqz,
+	I32Eq,
+	I32Ne,
+	I32LtS,
+	I32LtU,
+	I32GtS,
+	I32GtU,
+	I32LeS,
+	I32LeU,
+	I32GeS,
+	I32GeU,
+
+	I64Eqz,
+	I64Eq,
+	I64Ne,
+	I64LtS,
+	I64LtU,
+	I64GtS,
+	I64GtU,
+	I64LeS,
+	I64LeU,
+	I64GeS,
+	I64GeU,
+
+	F32Eq,
+	F32Ne,
+	F32Lt,
+	F32Gt,
+	F32Le,
+	F32Ge,
+
+	F64Eq,
+	F64Ne,
+	F64Lt,
+	F64Gt,
+	F64Le,
+	F64Ge,
+
+	I32Clz,
+	I32Ctz,
+	I32Popcnt,
+	I32Add,
+	I32Sub,
+	I32Mul,
+	I32DivS,
+	I32DivU,
+	I32RemS,
+	I32RemU,
+	I32And,
+	I32Or,
+	I32Xor,
+	I32Shl,
+	I32ShrS,
+	I32ShrU,
+	I32Rotl,
+	I32Rotr,
+
+	I64Clz,
+	I64Ctz,
+	I64Popcnt,
+	I64Add,
+	I64Sub,
+	I64Mul,
+	I64DivS,
+	I64DivU,
+	I64RemS,
+	I64RemU,
+	I64And,
+	I64Or,
+	I64Xor,
+	I64Shl,
+	I64ShrS,
+	I64ShrU,
+	I64Rotl,
+	I64Rotr,
+	F32Abs,
+	F32Neg,
+	F32Ceil,
+	F32Floor,
+	F32Trunc,
+	F32Nearest,
+	F32Sqrt,
+	F32Add,
+	F32Sub,
+	F32Mul,
+	F32Div,
+	F32Min,
+	F32Max,
+	F32Copysign,
+	F64Abs,
+	F64Neg,
+	F64Ceil,
+	F64Floor,
+	F64Trunc,
+	F64Nearest,
+	F64Sqrt,
+	F64Add,
+	F64Sub,
+	F64Mul,
+	F64Div,
+	F64Min,
+	F64Max,
+	F64Copysign,
+
+	I32WrapI64,
+	I32TruncSF32,
+	I32TruncUF32,
+	I32TruncSF64,
+	I32TruncUF64,
+	I64ExtendSI32,
+	I64ExtendUI32,
+	I64TruncSF32,
+	I64TruncUF32,
+	I64TruncSF64,
+	I64TruncUF64,
+	F32ConvertSI32,
+	F32ConvertUI32,
+	F32ConvertSI64,
+	F32ConvertUI64,
+	F32DemoteF64,
+	F64ConvertSI32,
+	F64ConvertUI32,
+	F64ConvertSI64,
+	F64ConvertUI64,
+	F64PromoteF32,
+
+	I32ReinterpretF32,
+	I64ReinterpretF64,
+	F32ReinterpretI32,
+	F64ReinterpretI64,
+}
+
+#[derive(Debug, Clone)]
+pub struct Instructions {
+	pub code: Vec<Instruction>,
+}

src/lib.rs

@@ -109,7 +109,6 @@ extern crate nan_preserving_float;
 
 use std::fmt;
 use std::error;
-use std::collections::HashMap;
 
 /// Error type which can thrown by wasm code or by host environment.
 ///
@@ -356,6 +355,7 @@ mod imports;
 mod global;
 mod func;
 mod types;
+mod isa;
 
 #[cfg(test)]
 mod tests;
@@ -378,7 +378,7 @@ pub mod memory_units {
 
 /// Deserialized module prepared for instantiation.
 pub struct Module {
-	labels: HashMap<usize, HashMap<usize, usize>>,
+	code_map: Vec<isa::Instructions>,
 	module: parity_wasm::elements::Module,
 }
 
@@ -418,12 +418,12 @@ impl Module {
 	pub fn from_parity_wasm_module(module: parity_wasm::elements::Module) -> Result<Module, Error> {
 		use validation::{validate_module, ValidatedModule};
 		let ValidatedModule {
-			labels,
+			code_map,
 			module,
 		} = validate_module(module)?;
 
 		Ok(Module {
-			labels,
+			code_map,
 			module,
 		})
 	}
@@ -524,7 +524,7 @@ impl Module {
 		&self.module
 	}
 
-	pub(crate) fn labels(&self) -> &HashMap<usize, HashMap<usize, usize>> {
-		&self.labels
+	pub(crate) fn code(&self) -> &Vec<isa::Instructions> {
+		&self.code_map
 	}
 }

src/module.rs

@@ -291,7 +291,7 @@ impl ModuleInstance {
 			}
 		}
 
-		let labels = loaded_module.labels();
+		let code = loaded_module.code();
 		{
 			let funcs = module.function_section().map(|fs| fs.entries()).unwrap_or(
 				&[],
@@ -308,13 +308,12 @@ impl ModuleInstance {
 				let signature = instance.signature_by_index(ty.type_ref()).expect(
 					"Due to validation type should exists",
 				);
-				let labels = labels.get(&index).expect(
+				let code = code.get(index).expect(
 					"At func validation time labels are collected; Collected labels are added by index; qed",
 				).clone();
 				let func_body = FuncBody {
 					locals: body.locals().to_vec(),
-					instructions: body.code().clone(),
-					labels: labels,
+					code: code,
 				};
 				let func_instance =
 					FuncInstance::alloc_internal(Rc::downgrade(&instance.0), signature, func_body);

src/validation/mod.rs

@@ -1,14 +1,15 @@
 use std::error;
 use std::fmt;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 use parity_wasm::elements::{
 	BlockType, External, GlobalEntry, GlobalType, Internal, MemoryType, Module, Instruction,
 	ResizableLimits, TableType, ValueType, InitExpr, Type,
 };
 use common::stack;
 use self::context::ModuleContextBuilder;
-use self::func::Validator;
+use self::func::FunctionReader;
 use memory_units::Pages;
+use isa;
 
 mod context;
 mod func;
@@ -40,7 +41,7 @@ impl From<stack::Error> for Error {
 
 #[derive(Clone)]
 pub struct ValidatedModule {
-	pub labels: HashMap<usize, HashMap<usize, usize>>,
+	pub code_map: Vec<isa::Instructions>,
 	pub module: Module,
 }
 
@@ -167,7 +168,7 @@ pub fn deny_floating_point(module: &Module) -> Result<(), Error> {
 pub fn validate_module(module: Module) -> Result<ValidatedModule, Error> {
 	let mut context_builder = ModuleContextBuilder::new();
 	let mut imported_globals = Vec::new();
-	let mut labels = HashMap::new();
+	let mut code_map = Vec::new();
 
 	// Copy types from module as is.
 	context_builder.set_types(
@@ -257,12 +258,12 @@ pub fn validate_module(module: Module) -> Result<ValidatedModule, Error> {
 					index
 				)),
 			)?;
-			let func_labels = Validator::validate_function(&context, function, function_body)
+			let code = FunctionReader::read_function(&context, function, function_body)
 				.map_err(|e| {
 					let Error(ref msg) = e;
-					Error(format!("Function #{} validation error: {}", index, msg))
+					Error(format!("Function #{} reading/validation error: {}", index, msg))
 				})?;
-			labels.insert(index, func_labels);
+			code_map.push(code);
 		}
 	}
 
@@ -374,7 +375,7 @@ pub fn validate_module(module: Module) -> Result<ValidatedModule, Error> {
 
 	Ok(ValidatedModule {
 		module,
-		labels
+		code_map,
 	})
 }
 

src/validation/tests.rs

@@ -1,9 +1,12 @@
-use super::validate_module;
+use super::{validate_module, ValidatedModule};
 use parity_wasm::builder::module;
 use parity_wasm::elements::{
-    External, GlobalEntry, GlobalType, ImportEntry, InitExpr, MemoryType,
-    Instruction, Instructions, TableType, ValueType, BlockType
+	External, GlobalEntry, GlobalType, ImportEntry, InitExpr, MemoryType,
+	Instruction, Instructions, TableType, ValueType, BlockType, deserialize_buffer,
+	Module,
 };
+use isa;
+use wabt;
 
 #[test]
 fn empty_is_valid() {
@@ -299,3 +302,617 @@ fn if_else_with_return_type_validation() {
 		.build();
 	validate_module(m).unwrap();
 }
+
+fn validate(wat: &str) -> ValidatedModule {
+	let wasm = wabt::wat2wasm(wat).unwrap();
+	let module = deserialize_buffer::<Module>(&wasm).unwrap();
+	let validated_module = validate_module(module).unwrap();
+	validated_module
+}
+
+fn compile(wat: &str) -> Vec<isa::Instruction> {
+	let validated_module = validate(wat);
+	let code = &validated_module.code_map[0];
+	code.code.clone()
+}
+
+#[test]
+fn implicit_return_no_value() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			})
+		]
+	)
+}
+
+#[test]
+fn implicit_return_with_value() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (result i32)
+				i32.const 0
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(0),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}
+
+#[test]
+fn implicit_return_param() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32)
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn get_local() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32) (result i32)
+				get_local 0
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::GetLocal(1),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}
+
+#[test]
+fn explicit_return() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32) (result i32)
+				get_local 0
+				return
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::GetLocal(1),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}
+
+#[test]
+fn add_params() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32) (param i32) (result i32)
+				get_local 0
+				get_local 1
+				i32.add
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			// This is tricky. Locals are now loaded from the stack. The load
+			// happens from address relative of the current stack pointer. The first load
+			// takes the value below the previous one (i.e the second argument) and then, it increments
+			// the stack pointer. And then the same thing hapens with the value below the previous one
+			// (which happens to be the value loaded by the first get_local).
+			isa::Instruction::GetLocal(2),
+			isa::Instruction::GetLocal(2),
+			isa::Instruction::I32Add,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 2,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}
+
+#[test]
+fn drop_locals() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32)
+				(local i32)
+				get_local 0
+				set_local 1
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::GetLocal(2),
+			isa::Instruction::SetLocal(1),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 2,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn if_without_else() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32) (result i32)
+				i32.const 1
+				if
+					i32.const 2
+					return
+				end
+				i32.const 3
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfEqz(isa::Target {
+				dst_pc: 4,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1, // 1 param
+				keep: isa::Keep::Single, // 1 result
+			}),
+			isa::Instruction::I32Const(3),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}
+
+#[test]
+fn if_else() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				(local i32)
+				i32.const 1
+				if
+					i32.const 2
+					set_local 0
+				else
+					i32.const 3
+					set_local 0
+				end
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfEqz(isa::Target {
+				dst_pc: 5,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::SetLocal(1),
+			isa::Instruction::Br(isa::Target {
+				dst_pc: 7,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(3),
+			isa::Instruction::SetLocal(1),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn if_else_returns_result() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				i32.const 1
+				if (result i32)
+					i32.const 2
+				else
+					i32.const 3
+				end
+				drop
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfEqz(isa::Target {
+				dst_pc: 4,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::Br(isa::Target {
+				dst_pc: 5,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(3),
+			isa::Instruction::Drop,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn if_else_branch_from_true_branch() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				i32.const 1
+				if (result i32)
+					i32.const 1
+					i32.const 1
+					br_if 0
+					drop
+					i32.const 2
+				else
+					i32.const 3
+				end
+				drop
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfEqz(isa::Target {
+				dst_pc: 8,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfNez(isa::Target {
+				dst_pc: 9,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::Single,
+				},
+			}),
+			isa::Instruction::Drop,
+			isa::Instruction::I32Const(2),
+			isa::Instruction::Br(isa::Target {
+				dst_pc: 9,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(3),
+			isa::Instruction::Drop,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn if_else_branch_from_false_branch() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				i32.const 1
+				if (result i32)
+					i32.const 1
+				else
+					i32.const 2
+					i32.const 1
+					br_if 0
+					drop
+					i32.const 3
+				end
+				drop
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfEqz(isa::Target {
+				dst_pc: 4,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::Br(isa::Target {
+				dst_pc: 9,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfNez(isa::Target {
+				dst_pc: 9,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::Single,
+				},
+			}),
+			isa::Instruction::Drop,
+			isa::Instruction::I32Const(3),
+			isa::Instruction::Drop,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn loop_() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				loop (result i32)
+					i32.const 1
+					br_if 0
+					i32.const 2
+				end
+				drop
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrIfNez(isa::Target {
+				dst_pc: 0,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::Drop,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn loop_empty() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				loop
+				end
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn brtable() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				block $1
+					loop $2
+						i32.const 0
+						br_table $2 $1
+					end
+				end
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(0),
+			isa::Instruction::BrTable(
+				vec![
+					isa::Target {
+						dst_pc: 0,
+						drop_keep: isa::DropKeep {
+							drop: 0,
+							keep: isa::Keep::None,
+						},
+					},
+					isa::Target {
+						dst_pc: 2,
+						drop_keep: isa::DropKeep {
+							drop: 0,
+							keep: isa::Keep::None,
+						},
+					},
+				].into_boxed_slice()
+			),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn brtable_returns_result() {
+	let code = compile(r#"
+		(module
+			(func (export "call")
+				block $1 (result i32)
+					block $2 (result i32)
+						i32.const 0
+						i32.const 1
+						br_table $2 $1
+					end
+					unreachable
+				end
+				drop
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::I32Const(0),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::BrTable(
+				vec![
+					isa::Target {
+						dst_pc: 3,
+						drop_keep: isa::DropKeep {
+							drop: 0,
+							keep: isa::Keep::Single,
+						},
+					},
+					isa::Target {
+						dst_pc: 4,
+						drop_keep: isa::DropKeep {
+							keep: isa::Keep::Single,
+							drop: 0,
+						},
+					},
+				].into_boxed_slice()
+			),
+			isa::Instruction::Unreachable,
+			isa::Instruction::Drop,
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 0,
+				keep: isa::Keep::None,
+			}),
+		]
+	)
+}
+
+#[test]
+fn wabt_example() {
+	let code = compile(r#"
+		(module
+			(func (export "call") (param i32) (result i32)
+				block $exit
+					get_local 0
+					br_if $exit
+					i32.const 1
+					return
+				end
+				i32.const 2
+				return
+			)
+		)
+	"#);
+	assert_eq!(
+		code,
+		vec![
+			isa::Instruction::GetLocal(1),
+			isa::Instruction::BrIfNez(isa::Target {
+				dst_pc: 4,
+				drop_keep: isa::DropKeep {
+					drop: 0,
+					keep: isa::Keep::None,
+				},
+			}),
+			isa::Instruction::I32Const(1),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1, // 1 parameter
+				keep: isa::Keep::Single,
+			}),
+			isa::Instruction::I32Const(2),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+			isa::Instruction::Return(isa::DropKeep {
+				drop: 1,
+				keep: isa::Keep::Single,
+			}),
+		]
+	)
+}

src/validation/util.rs

@@ -10,14 +10,36 @@ use validation::Error;
 pub struct Locals<'a> {
 	params: &'a [ValueType],
 	local_groups: &'a [Local],
+	count: u32,
 }
 
 impl<'a> Locals<'a> {
-	pub fn new(params: &'a [ValueType], local_groups: &'a [Local]) -> Locals<'a> {
-		Locals {
+	/// Create a new wrapper around declared variables and parameters.
+	pub fn new(params: &'a [ValueType], local_groups: &'a [Local]) -> Result<Locals<'a>, Error> {
+		let mut acc = params.len() as u32;
+		for locals_group in local_groups {
+			acc = acc
+				.checked_add(locals_group.count())
+				.ok_or_else(||
+					Error(String::from("Locals range not in 32-bit range"))
+				)?;
+		}
+
+		Ok(Locals {
 			params,
 			local_groups,
-		}
+			count: acc,
+		})
+	}
+
+	/// Returns parameter count.
+	pub fn param_count(&self) -> u32 {
+		self.params.len() as u32
+	}
+
+	/// Returns total count of all declared locals and paramaterers.
+	pub fn count(&self) -> u32 {
+		self.count
 	}
 
 	/// Returns the type of a local variable (either a declared local or a param).
@@ -29,7 +51,7 @@ impl<'a> Locals<'a> {
 		}
 
 		// If an index doesn't point to a param, then we have to look into local declarations.
-		let mut start_idx = self.params.len() as u32;
+		let mut start_idx = self.param_count();
 		for locals_group in self.local_groups {
 			let end_idx = start_idx
 				.checked_add(locals_group.count())
@@ -62,7 +84,7 @@ mod tests {
 	fn locals_it_works() {
 		let params = vec![ValueType::I32, ValueType::I64];
 		let local_groups = vec![Local::new(2, ValueType::F32), Local::new(2, ValueType::F64)];
-		let locals = Locals::new(&params, &local_groups);
+		let locals = Locals::new(&params, &local_groups).unwrap();
 
 		assert_matches!(locals.type_of_local(0), Ok(ValueType::I32));
 		assert_matches!(locals.type_of_local(1), Ok(ValueType::I64));
@@ -76,7 +98,7 @@ mod tests {
 	#[test]
 	fn locals_no_declared_locals() {
 		let params = vec![ValueType::I32];
-		let locals = Locals::new(&params, &[]);
+		let locals = Locals::new(&params, &[]).unwrap();
 
 		assert_matches!(locals.type_of_local(0), Ok(ValueType::I32));
 		assert_matches!(locals.type_of_local(1), Err(_));
@@ -85,7 +107,7 @@ mod tests {
 	#[test]
 	fn locals_no_params() {
 		let local_groups = vec![Local::new(2, ValueType::I32), Local::new(3, ValueType::I64)];
-		let locals = Locals::new(&[], &local_groups);
+		let locals = Locals::new(&[], &local_groups).unwrap();
 
 		assert_matches!(locals.type_of_local(0), Ok(ValueType::I32));
 		assert_matches!(locals.type_of_local(1), Ok(ValueType::I32));
@@ -101,12 +123,9 @@ mod tests {
 			Local::new(u32::max_value(), ValueType::I32),
 			Local::new(1, ValueType::I64),
 		];
-		let locals = Locals::new(&[], &local_groups);
-
 		assert_matches!(
-			locals.type_of_local(u32::max_value() - 1),
-			Ok(ValueType::I32)
+			Locals::new(&[], &local_groups),
+			Err(_)
 		);
-		assert_matches!(locals.type_of_local(u32::max_value()), Err(_));
 	}
 }

tests/spec/run.rs

@@ -360,6 +360,8 @@ fn try_spec(name: &str) -> Result<(), Error> {
 				}};
 		}
 
+		println!("Running spec cmd {}: {:?}", line, kind);
+
 		match kind {
 			CommandKind::Module { name, module, .. } => {
 				load_module(&module.into_vec()?, &name, &mut spec_driver)