Merge pull request #141 from paritytech/at-fixes

Fix issues with the stack height metering

cli/build/main.rs

@@ -230,7 +230,7 @@ mod tests {
 
 			let wasm_path = target_path.join("example_wasm.wasm");
 			let mut f = fs::File::create(wasm_path).expect("create fail failed");
-			f.write(b"\0asm").expect("write file failed");
+			f.write_all(b"\0asm").expect("write file failed");
 		}
 
 		let path = tmp_dir.path().to_string_lossy();

src/stack_height/max_height.rs

@@ -288,6 +288,9 @@ pub(crate) fn compute(func_idx: u32, module: &elements::Module) -> Result<u32, E
 					.get(x as usize)
 					.ok_or_else(|| Error("Type not found".into()))?;
 
+				// Pop the offset into the function table.
+				stack.pop_values(1)?;
+
 				// Pop values for arguments of the function.
 				stack.pop_values(ty.params().len() as u32)?;
 
@@ -523,4 +526,28 @@ mod tests {
 		let height = compute(0, &module).unwrap();
 		assert_eq!(height, 2);
 	}
+
+	#[test]
+	fn call_indirect() {
+		let module = parse_wat(
+			r#"
+(module
+	(table $ptr 1 1 funcref)
+	(elem $ptr (i32.const 0) func 1)
+	(func $main
+		(call_indirect (i32.const 0))
+		(call_indirect (i32.const 0))
+		(call_indirect (i32.const 0))
+	)
+	(func $callee
+		i64.const 42
+		drop
+	)
+)
+"#,
+		);
+
+		let height = compute(0, &module).unwrap();
+		assert_eq!(height, 1);
+	}
 }

src/stack_height/thunk.rs

@@ -13,7 +13,6 @@ struct Thunk {
 	signature: FunctionType,
 	// Index in function space of this thunk.
 	idx: Option<u32>,
-	original_func_idx: u32,
 	callee_stack_cost: u32,
 }
 
@@ -23,9 +22,6 @@ pub(crate) fn generate_thunks(
 ) -> Result<elements::Module, Error> {
 	// First, we need to collect all function indices that should be replaced by thunks
 
-	// Function indices which needs to generate thunks.
-	let mut need_thunks: Vec<u32> = Vec::new();
-
 	let mut replacement_map: Map<u32, Thunk> = {
 		let exports = module
 			.export_section()
@@ -57,12 +53,10 @@ pub(crate) fn generate_thunks(
 
 			// Don't generate a thunk if stack_cost of a callee is zero.
 			if callee_stack_cost != 0 {
-				need_thunks.push(func_idx);
 				replacement_map.insert(func_idx, Thunk {
 					signature: resolve_func_type(func_idx, &module)?.clone(),
 					idx: None,
 					callee_stack_cost,
-					original_func_idx: func_idx,
 				});
 			}
 		}
@@ -76,17 +70,9 @@ pub(crate) fn generate_thunks(
 	let mut next_func_idx = module.functions_space() as u32;
 
 	let mut mbuilder = builder::from_module(module);
-	for func_idx in need_thunks {
-		let mut thunk = replacement_map
-			.get_mut(&func_idx)
-			.expect(
-				"`func_idx` should come from `need_thunks`;
-				`need_thunks` is populated with the same items that in `replacement_map`;
-				qed"
-			);
-
+	for (func_idx, thunk) in replacement_map.iter_mut() {
 		let instrumented_call = instrument_call!(
-			thunk.original_func_idx as u32,
+			*func_idx,
 			thunk.callee_stack_cost as i32,
 			ctx.stack_height_global_idx(),
 			ctx.stack_limit()
@@ -113,7 +99,7 @@ pub(crate) fn generate_thunks(
 				// Signature of the thunk should match the original function signature.
 				.signature()
 					.with_params(thunk.signature.params().to_vec())
-					.with_return_type(thunk.signature.return_type().clone())
+					.with_return_type(thunk.signature.return_type())
 					.build()
 				.body()
 					.with_instructions(elements::Instructions::new(

tests/expectations/stack-height/start.wat

@@ -22,23 +22,7 @@
     i32.const 1
     i32.sub
     global.set 0)
-  (func (;4;) (type 1)
-    global.get 0
-    i32.const 1
-    i32.add
-    global.set 0
-    global.get 0
-    i32.const 1024
-    i32.gt_u
-    if  ;; label = @1
-      unreachable
-    end
-    call 1
-    global.get 0
-    i32.const 1
-    i32.sub
-    global.set 0)
   (global (;0;) (mut i32) (i32.const 0))
-  (export "exported_start" (func 4))
+  (export "exported_start" (func 3))
   (export "call" (func 2))
-  (start 4))
+  (start 3))

tests/expectations/stack-height/table.wat

@@ -26,25 +26,7 @@
     local.get 0
     local.get 1
     i32.add)
-  (func (;3;) (type 2) (param i32 i32) (result i32)
-    local.get 0
-    local.get 1
-    global.get 0
-    i32.const 2
-    i32.add
-    global.set 0
-    global.get 0
-    i32.const 1024
-    i32.gt_u
-    if  ;; label = @1
-      unreachable
-    end
-    call 2
-    global.get 0
-    i32.const 2
-    i32.sub
-    global.set 0)
-  (func (;4;) (type 1) (param i32)
+  (func (;3;) (type 1) (param i32)
     local.get 0
     global.get 0
     i32.const 2
@@ -61,7 +43,7 @@
     i32.const 2
     i32.sub
     global.set 0)
-  (func (;5;) (type 2) (param i32 i32) (result i32)
+  (func (;4;) (type 2) (param i32 i32) (result i32)
     local.get 0
     local.get 1
     global.get 0
@@ -81,5 +63,5 @@
     global.set 0)
   (table (;0;) 10 funcref)
   (global (;0;) (mut i32) (i32.const 0))
-  (export "i32.add" (func 5))
-  (elem (;0;) (i32.const 0) func 0 4 5))
+  (export "i32.add" (func 4))
+  (elem (;0;) (i32.const 0) func 0 3 4))