From c5ac4cba56d377e8348733d818dd26022ee59490 Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Mon, 26 Jun 2017 14:27:29 +0300
Subject: [PATCH] updated to new parity-wasm api

---
 src/ext.rs       |  3 --
 src/gas.rs       | 73 +++++++++++++++++++++++++++++++++++++++---------
 src/optimizer.rs |  9 ------
 src/symbols.rs   |  3 --
 4 files changed, 60 insertions(+), 28 deletions(-)

diff --git a/src/ext.rs b/src/ext.rs
index 5891d6a..a8a592b 100644
--- a/src/ext.rs
+++ b/src/ext.rs
@@ -6,9 +6,6 @@ pub fn update_call_index(opcodes: &mut elements::Opcodes, original_imports: usiz
 	use parity_wasm::elements::Opcode::*;
 	for opcode in opcodes.elements_mut().iter_mut() {
 		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				update_call_index(block, original_imports, inserts)
-			},
 			&mut Call(ref mut call_index) => {
 				if let Some(pos) = inserts.iter().position(|x| x.1 == *call_index) {
 					*call_index = (original_imports + pos) as u32; 
diff --git a/src/gas.rs b/src/gas.rs
index 3b469ec..ac02de3 100644
--- a/src/gas.rs
+++ b/src/gas.rs
@@ -4,31 +4,78 @@ pub fn update_call_index(opcodes: &mut elements::Opcodes, inserted_index: u32) {
 	use parity_wasm::elements::Opcode::*;
 	for opcode in opcodes.elements_mut().iter_mut() {
 		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				update_call_index(block, inserted_index)
-			},
 			&mut Call(ref mut call_index) => {
 				if *call_index >= inserted_index { *call_index += 1}
 			},
-			_ => { }
+			_ => { },
 		}
 	}
 }
 
+enum InjectAction {
+	Spawn,
+	Continue,
+	Increment,
+	IncrementSpawn,
+}
+
 pub fn inject_counter(opcodes: &mut elements::Opcodes, gas_func: u32) {
 	use parity_wasm::elements::Opcode::*;
-	for opcode in opcodes.elements_mut().iter_mut() {
-		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				inject_counter(block, gas_func)
+
+	let mut stack: Vec<(usize, usize)> = Vec::new(); 
+	let mut cursor = 0;
+	stack.push((0, 1));
+
+	loop {
+		if cursor >= opcodes.elements().len() {
+			break;
+		}
+
+		let last_entry = stack.pop().expect("There should be at least one entry on stack");
+
+		let action = {
+			let opcode = &opcodes.elements()[cursor];
+			match *opcode {
+				Block(_) | If(_) | Loop(_) => {
+					InjectAction::Spawn
+				},
+				Else => {
+					InjectAction::IncrementSpawn
+				},
+				End => {
+					InjectAction::Increment
+				},
+				_ => {
+					InjectAction::Continue
+				}
+			}
+		};
+
+		match action {
+			InjectAction::Increment => {
+				let (pos, ops) = last_entry;
+				opcodes.elements_mut().insert(pos, I32Const(ops as i32));
+				opcodes.elements_mut().insert(pos, Call(gas_func));
+				cursor += 3;
+			},
+			InjectAction::IncrementSpawn => {
+				let (pos, ops) = last_entry;
+				opcodes.elements_mut().insert(pos, I32Const(ops as i32));
+				opcodes.elements_mut().insert(pos, Call(gas_func));
+				cursor += 3;
+				stack.push((cursor, 1));
+			},
+			InjectAction::Continue => {
+				cursor += 1;
+				let (pos, ops) = last_entry;
+				stack.push((pos, ops+1));
+			},
+			InjectAction::Spawn => {
+				cursor += 1;
+				stack.push((cursor, 1));
 			},
-			_ => { }
 		}
 	}
-
-	let ops = opcodes.elements_mut().len() as u32;
-	opcodes.elements_mut().insert(0, I32Const(ops as i32));
-	opcodes.elements_mut().insert(1, Call(gas_func));
 }
 
 pub fn inject_gas_counter(module: elements::Module) -> elements::Module {
diff --git a/src/optimizer.rs b/src/optimizer.rs
index 992c411..bdb169e 100644
--- a/src/optimizer.rs
+++ b/src/optimizer.rs
@@ -262,9 +262,6 @@ pub fn update_call_index(opcodes: &mut elements::Opcodes, eliminated_indices: &[
 	use parity_wasm::elements::Opcode::*;
 	for opcode in opcodes.elements_mut().iter_mut() {
 		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				update_call_index(block, eliminated_indices)
-			},
 			&mut Call(ref mut call_index) => {
 				let totalle = eliminated_indices.iter().take_while(|i| (**i as u32) < *call_index).count();
 				trace!("rewired call {} -> call {}", *call_index, *call_index - totalle as u32);
@@ -280,9 +277,6 @@ pub fn update_global_index(opcodes: &mut Vec<elements::Opcode>, eliminated_indic
 	use parity_wasm::elements::Opcode::*;
 	for opcode in opcodes.iter_mut() {
 		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				update_global_index(block.elements_mut(), eliminated_indices)
-			},
 			&mut GetGlobal(ref mut index) | &mut SetGlobal(ref mut index) => {
 				let totalle = eliminated_indices.iter().take_while(|i| (**i as u32) < *index).count();
 				trace!("rewired global {} -> global {}", *index, *index - totalle as u32);
@@ -298,9 +292,6 @@ pub fn update_type_index(opcodes: &mut elements::Opcodes, eliminated_indices: &[
 	use parity_wasm::elements::Opcode::*;
 	for opcode in opcodes.elements_mut().iter_mut() {
 		match opcode {
-			&mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => {
-				update_type_index(block, eliminated_indices)
-			},
 			&mut CallIndirect(ref mut call_index, _) => {
 				let totalle = eliminated_indices.iter().take_while(|i| (**i as u32) < *call_index).count();
 				trace!("rewired call_indrect {} -> call_indirect {}", *call_index, *call_index - totalle as u32);
diff --git a/src/symbols.rs b/src/symbols.rs
index 0b567ac..3348f1e 100644
--- a/src/symbols.rs
+++ b/src/symbols.rs
@@ -62,9 +62,6 @@ pub fn push_code_symbols(module: &elements::Module, opcodes: &[elements::Opcode]
 			&GetGlobal(idx) | &SetGlobal(idx) => {
 				dest.push(resolve_global(module, idx))
 			},
-			&If(_, ref block) | &Loop(_, ref block) | &Block(_, ref block) => {
-				push_code_symbols(module, block.elements(), dest);
-			},
 			_ => { },
 		} 
 	}