From fe575b4eaa3a9d23c4ee8e3352cb25929fbf290a Mon Sep 17 00:00:00 2001
From: NikVolf <nikvolf@gmail.com>
Date: Wed, 12 Apr 2017 19:41:36 +0300
Subject: [PATCH] working gas counter

---
 gas/src/main.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/gas/src/main.rs b/gas/src/main.rs
index 95e272f..e1f44ec 100644
--- a/gas/src/main.rs
+++ b/gas/src/main.rs
@@ -18,6 +18,22 @@ pub fn update_call_index(opcodes: &mut elements::Opcodes, inserted_index: u32) {
     }
 }
 
+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) => {
+                inject_counter(block, gas_func)
+            },
+            _ => { }
+        }
+    }
+
+    let ops = opcodes.elements_mut().len() as u32;
+    opcodes.elements_mut().insert(0, I32Const(ops));
+    opcodes.elements_mut().insert(1, Call(gas_func));
+}
+
 fn main() {
 
     let args = env::args().collect::<Vec<_>>();
@@ -66,6 +82,7 @@ fn main() {
             &mut elements::Section::Code(ref mut code_section) => {
                 for ref mut func_body in code_section.bodies_mut() {
                     update_call_index(func_body.code_mut(), gas_func);
+                    inject_counter(func_body.code_mut(), gas_func);
                 }
             },
             _ => { }