From 4d78c6bf380486500dedfc039673e1d829ed2b84 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 4 Oct 2023 22:00:43 +0300
Subject: [PATCH 01/17] add incomplete implementation

---
 utils/wasm-gen/src/tests.rs | 33 +++++++++++++++++
 utils/wasm-gen/src/utils.rs | 72 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs
index fa5887bbe6a..b37e007401f 100644
--- a/utils/wasm-gen/src/tests.rs
+++ b/utils/wasm-gen/src/tests.rs
@@ -47,6 +47,39 @@ use std::mem;
 
 const UNSTRUCTURED_SIZE: usize = 1_000_000;
 
+#[test]
+fn instrument_recursions() {
+    let wat1 = r#"
+    (module
+        (func $import0 (import "env" "gr_leave"))
+        (memory $memory0 (import "env" "memory") 16)
+        (export "handle" (func $handle))
+        (func $handle
+            call $f
+            drop
+        )
+        (func $f (result i64)
+            call $f
+        )
+    )"#;
+
+    let wasm_bytes = wat::parse_str(wat1).expect("invalid wat");
+    let module =
+        parity_wasm::deserialize_buffer::<Module>(&wasm_bytes).expect("invalid wasm bytes");
+    let limited_recursions_module = utils::instrument_recursion(module);
+
+    let wasm_bytes = limited_recursions_module
+        .into_bytes()
+        .expect("invalid pw module");
+    assert!(wasmparser::validate(&wasm_bytes).is_ok());
+
+    let wat = wasmprinter::print_bytes(&wasm_bytes).expect("failed printing bytes");
+    println!("wat = {wat}");
+
+    let code_res = Code::try_new(wasm_bytes, 1, |_| CustomConstantCostRules::default(), None);
+    assert!(code_res.is_ok());
+}
+
 #[test]
 fn remove_trivial_recursions() {
     let wat1 = r#"
diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index c89440967a6..dae90ae572e 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -16,9 +16,12 @@
 // 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 crate::wasm::PageCount as WasmPageCount;
 use gear_wasm_instrument::parity_wasm::{
     builder,
-    elements::{self, FuncBody, ImportCountType, Instruction, Module, Type, ValueType},
+    elements::{
+        self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type, ValueType,
+    },
 };
 use gsys::HashWithValue;
 use std::{
@@ -215,6 +218,73 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
+pub fn instrument_recursion(mut module: Module) -> Module {
+    let Some(mem_size) = module
+        .import_section()
+        .and_then(|section| {
+            section
+                .entries()
+                .iter()
+                .find_map(|entry| match entry.external() {
+                    External::Memory(mem_ty) => Some(mem_ty.limits().initial()),
+                    _ => None,
+                })
+        })
+        .map(Into::<WasmPageCount>::into).map(|page_count| page_count.memory_size()) else {
+            return module;
+        };
+
+    let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
+
+    let Some(code_section) = module.code_section_mut() else {
+        return module;
+    };
+
+    for func_body in code_section.bodies_mut() {
+        let instructions = func_body.code_mut().elements_mut();
+
+        instructions.splice(
+            0..0,
+            [
+                //if call_depth > 256 { call_depth = 0; gr_leave(); }
+                Instruction::I32Const(0),
+                Instruction::I32Load(2, call_depth_ptr),
+                Instruction::I32Const(256),
+                Instruction::I32GtU,
+                Instruction::If(BlockType::NoResult),
+                Instruction::I32Const(0),
+                Instruction::I32Const(0),
+                Instruction::I32Store(2, call_depth_ptr),
+                Instruction::Unreachable, //TODO: Instruction::Call(gr_leave_index),
+                Instruction::End,
+                //call_depth += 1;
+                Instruction::I32Const(0),
+                Instruction::I32Const(0),
+                Instruction::I32Load(2, call_depth_ptr),
+                Instruction::I32Const(1),
+                Instruction::I32Add,
+                Instruction::I32Store(2, call_depth_ptr),
+            ],
+        );
+
+        let last = instructions.len() - 1;
+        instructions.splice(
+            last..last,
+            [
+                //call_depth -= 1;
+                Instruction::I32Const(0),
+                Instruction::I32Const(0),
+                Instruction::I32Load(2, call_depth_ptr),
+                Instruction::I32Const(-1),
+                Instruction::I32Add,
+                Instruction::I32Store(2, call_depth_ptr),
+            ],
+        );
+    }
+
+    module
+}
+
 pub(crate) fn hash_with_value_to_vec(hash_with_value: &HashWithValue) -> Vec<u8> {
     let address_data_size = mem::size_of::<HashWithValue>();
     let address_data_slice = unsafe {

From eaa0976acea682ad34f49fdff5c89c9939322856 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 4 Oct 2023 23:37:07 +0300
Subject: [PATCH 02/17] complete impl & use it in wasm-gen

---
 utils/wasm-gen/src/generator.rs |  2 +-
 utils/wasm-gen/src/tests.rs     |  1 -
 utils/wasm-gen/src/utils.rs     | 79 ++++++++++++++++++++++++++++++---
 3 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/utils/wasm-gen/src/generator.rs b/utils/wasm-gen/src/generator.rs
index ffa44d21a90..0f38d92641c 100644
--- a/utils/wasm-gen/src/generator.rs
+++ b/utils/wasm-gen/src/generator.rs
@@ -141,7 +141,7 @@ impl<'a, 'b> GearWasmGenerator<'a, 'b> {
 
         Ok(if config.remove_recursions {
             log::trace!("Removing recursions");
-            utils::remove_recursion(module)
+            utils::instrument_recursion(module)
         } else {
             log::trace!("Leaving recursions");
             module
diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs
index b37e007401f..02c0f66ccfc 100644
--- a/utils/wasm-gen/src/tests.rs
+++ b/utils/wasm-gen/src/tests.rs
@@ -51,7 +51,6 @@ const UNSTRUCTURED_SIZE: usize = 1_000_000;
 fn instrument_recursions() {
     let wat1 = r#"
     (module
-        (func $import0 (import "env" "gr_leave"))
         (memory $memory0 (import "env" "memory") 16)
         (export "handle" (func $handle))
         (func $handle
diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index dae90ae572e..1afd4a069f1 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -16,12 +16,18 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
+#![allow(dead_code)]
+
 use crate::wasm::PageCount as WasmPageCount;
-use gear_wasm_instrument::parity_wasm::{
-    builder,
-    elements::{
-        self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type, ValueType,
+use gear_wasm_instrument::{
+    parity_wasm::{
+        builder,
+        elements::{
+            self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type,
+            ValueType,
+        },
     },
+    syscalls::SysCallName,
 };
 use gsys::HashWithValue;
 use std::{
@@ -218,7 +224,7 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
-pub fn instrument_recursion(mut module: Module) -> Module {
+pub fn instrument_recursion(module: Module) -> Module {
     let Some(mem_size) = module
         .import_section()
         .and_then(|section| {
@@ -236,6 +242,27 @@ pub fn instrument_recursion(mut module: Module) -> Module {
 
     let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
 
+    let mut mbuilder = builder::from_module(module);
+
+    // fn gr_leave() -> !;
+    let import_sig = mbuilder.push_signature(builder::signature().build_sig());
+
+    mbuilder.push_import(
+        builder::import()
+            .module("env")
+            .field(SysCallName::Leave.to_str())
+            .external()
+            .func(import_sig)
+            .build(),
+    );
+
+    // back to plain module
+    let mut module = mbuilder.build();
+
+    let import_count = module.import_count(ImportCountType::Function);
+    let gr_leave_index @ inserted_index = import_count as u32 - 1;
+    let inserted_count = 1;
+
     let Some(code_section) = module.code_section_mut() else {
         return module;
     };
@@ -243,6 +270,14 @@ pub fn instrument_recursion(mut module: Module) -> Module {
     for func_body in code_section.bodies_mut() {
         let instructions = func_body.code_mut().elements_mut();
 
+        for instruction in instructions.iter_mut() {
+            if let Instruction::Call(call_index) = instruction {
+                if *call_index >= inserted_index {
+                    *call_index += inserted_count
+                }
+            }
+        }
+
         instructions.splice(
             0..0,
             [
@@ -255,7 +290,7 @@ pub fn instrument_recursion(mut module: Module) -> Module {
                 Instruction::I32Const(0),
                 Instruction::I32Const(0),
                 Instruction::I32Store(2, call_depth_ptr),
-                Instruction::Unreachable, //TODO: Instruction::Call(gr_leave_index),
+                Instruction::Call(gr_leave_index),
                 Instruction::End,
                 //call_depth += 1;
                 Instruction::I32Const(0),
@@ -282,6 +317,38 @@ pub fn instrument_recursion(mut module: Module) -> Module {
         );
     }
 
+    let sections = module.sections_mut();
+    sections.retain(|section| !matches!(section, elements::Section::Custom(_)));
+
+    for section in sections {
+        match section {
+            elements::Section::Export(export_section) => {
+                for export in export_section.entries_mut() {
+                    if let elements::Internal::Function(func_index) = export.internal_mut() {
+                        if *func_index >= inserted_index {
+                            *func_index += inserted_count
+                        }
+                    }
+                }
+            }
+            elements::Section::Element(elements_section) => {
+                for segment in elements_section.entries_mut() {
+                    for func_index in segment.members_mut() {
+                        if *func_index >= inserted_index {
+                            *func_index += inserted_count
+                        }
+                    }
+                }
+            }
+            elements::Section::Start(start_idx) => {
+                if *start_idx >= inserted_index {
+                    *start_idx += inserted_count
+                }
+            }
+            _ => {}
+        }
+    }
+
     module
 }
 

From eb665dc875713b0ffc4d62ee918925f3a089c79c Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Fri, 6 Oct 2023 12:53:38 +0300
Subject: [PATCH 03/17] move instructions into 2 functions: prolog and epilog

---
 utils/wasm-gen/src/utils.rs | 84 ++++++++++++++++++++++++-------------
 1 file changed, 56 insertions(+), 28 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 1afd4a069f1..182e6b3d429 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -23,8 +23,8 @@ use gear_wasm_instrument::{
     parity_wasm::{
         builder,
         elements::{
-            self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type,
-            ValueType,
+            self, BlockType, External, FuncBody, ImportCountType, Instruction, Instructions,
+            Module, Type, ValueType,
         },
     },
     syscalls::SysCallName,
@@ -104,7 +104,7 @@ pub fn remove_recursion(module: Module) -> Module {
                     .with_results(signature.results().to_vec())
                     .build()
                     .body()
-                    .with_instructions(elements::Instructions::new(body))
+                    .with_instructions(Instructions::new(body))
                     .build()
                     .build(),
             )
@@ -257,30 +257,22 @@ pub fn instrument_recursion(module: Module) -> Module {
     );
 
     // back to plain module
-    let mut module = mbuilder.build();
+    let module = mbuilder.build();
 
     let import_count = module.import_count(ImportCountType::Function);
     let gr_leave_index @ inserted_index = import_count as u32 - 1;
     let inserted_count = 1;
 
-    let Some(code_section) = module.code_section_mut() else {
-        return module;
-    };
-
-    for func_body in code_section.bodies_mut() {
-        let instructions = func_body.code_mut().elements_mut();
+    let prolog_index = module.functions_space();
+    let epilog_index = prolog_index + 1;
 
-        for instruction in instructions.iter_mut() {
-            if let Instruction::Call(call_index) = instruction {
-                if *call_index >= inserted_index {
-                    *call_index += inserted_count
-                }
-            }
-        }
+    let mut mbuilder = builder::from_module(module);
 
-        instructions.splice(
-            0..0,
-            [
+    // prolog function
+    mbuilder.push_function(
+        builder::function()
+            .body()
+            .with_instructions(Instructions::new(vec![
                 //if call_depth > 256 { call_depth = 0; gr_leave(); }
                 Instruction::I32Const(0),
                 Instruction::I32Load(2, call_depth_ptr),
@@ -299,13 +291,17 @@ pub fn instrument_recursion(module: Module) -> Module {
                 Instruction::I32Const(1),
                 Instruction::I32Add,
                 Instruction::I32Store(2, call_depth_ptr),
-            ],
-        );
+                Instruction::End,
+            ]))
+            .build()
+            .build(),
+    );
 
-        let last = instructions.len() - 1;
-        instructions.splice(
-            last..last,
-            [
+    // epilog function
+    mbuilder.push_function(
+        builder::function()
+            .body()
+            .with_instructions(Instructions::new(vec![
                 //call_depth -= 1;
                 Instruction::I32Const(0),
                 Instruction::I32Const(0),
@@ -313,8 +309,40 @@ pub fn instrument_recursion(module: Module) -> Module {
                 Instruction::I32Const(-1),
                 Instruction::I32Add,
                 Instruction::I32Store(2, call_depth_ptr),
-            ],
-        );
+                Instruction::End,
+            ]))
+            .build()
+            .build(),
+    );
+
+    // back to plain module
+    let mut module = mbuilder.build();
+
+    let Some(code_section) = module.code_section_mut() else {
+        return module;
+    };
+
+    for (i, func_body) in code_section.bodies_mut().iter_mut().enumerate() {
+        if i + import_count >= prolog_index {
+            continue;
+        }
+
+        let instructions = func_body.code_mut().elements_mut();
+
+        for instruction in instructions.iter_mut() {
+            if let Instruction::Call(call_index) = instruction {
+                if *call_index >= inserted_index {
+                    *call_index += inserted_count
+                }
+            }
+        }
+
+        instructions.insert(0, Instruction::Call(prolog_index as u32));
+
+        if let Some(end @ Instruction::End) = instructions.pop() {
+            instructions.push(Instruction::Call(epilog_index as u32));
+            instructions.push(end);
+        }
     }
 
     let sections = module.sections_mut();

From 4f0cae814fc2ce68a30f99408fdaaf94dfb4b7c1 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Sat, 7 Oct 2023 21:02:19 +0300
Subject: [PATCH 04/17] exit from func using `return` instruction

---
 utils/wasm-gen/src/utils.rs | 191 ++++++++++++------------------------
 1 file changed, 65 insertions(+), 126 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 182e6b3d429..8badc67663d 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -19,15 +19,11 @@
 #![allow(dead_code)]
 
 use crate::wasm::PageCount as WasmPageCount;
-use gear_wasm_instrument::{
-    parity_wasm::{
-        builder,
-        elements::{
-            self, BlockType, External, FuncBody, ImportCountType, Instruction, Instructions,
-            Module, Type, ValueType,
-        },
+use gear_wasm_instrument::parity_wasm::{
+    builder,
+    elements::{
+        self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type, ValueType,
     },
-    syscalls::SysCallName,
 };
 use gsys::HashWithValue;
 use std::{
@@ -104,7 +100,7 @@ pub fn remove_recursion(module: Module) -> Module {
                     .with_results(signature.results().to_vec())
                     .build()
                     .body()
-                    .with_instructions(Instructions::new(body))
+                    .with_instructions(elements::Instructions::new(body))
                     .build()
                     .build(),
             )
@@ -224,7 +220,7 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
-pub fn instrument_recursion(module: Module) -> Module {
+pub fn instrument_recursion(mut module: Module) -> Module {
     let Some(mem_size) = module
         .import_section()
         .and_then(|section| {
@@ -242,66 +238,70 @@ pub fn instrument_recursion(module: Module) -> Module {
 
     let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
 
-    let mut mbuilder = builder::from_module(module);
+    let (Some(type_section), Some(function_section)) = (module.type_section(), module.function_section()) else {
+        return module;
+    };
 
-    // fn gr_leave() -> !;
-    let import_sig = mbuilder.push_signature(builder::signature().build_sig());
+    let types = type_section.types().to_vec();
+    let signature_entries = function_section.entries().to_vec();
 
-    mbuilder.push_import(
-        builder::import()
-            .module("env")
-            .field(SysCallName::Leave.to_str())
-            .external()
-            .func(import_sig)
-            .build(),
-    );
+    let Some(code_section) = module.code_section_mut() else {
+        return module;
+    };
 
-    // back to plain module
-    let module = mbuilder.build();
+    for (index, func_body) in code_section.bodies_mut().iter_mut().enumerate() {
+        let signature_index = &signature_entries[index];
+        let signature = &types[signature_index.type_ref() as usize];
+        let Type::Function(signature) = signature;
+        let results = signature.results();
+        let mut body = Vec::with_capacity(results.len() + 2);
+        for result in results {
+            let instruction = match result {
+                ValueType::I32 => Instruction::I32Const(u32::MAX as i32),
+                ValueType::I64 => Instruction::I64Const(u64::MAX as i64),
+                ValueType::F32 | ValueType::F64 => unreachable!("f32/64 types are not supported"),
+            };
 
-    let import_count = module.import_count(ImportCountType::Function);
-    let gr_leave_index @ inserted_index = import_count as u32 - 1;
-    let inserted_count = 1;
+            body.push(instruction);
+        }
 
-    let prolog_index = module.functions_space();
-    let epilog_index = prolog_index + 1;
+        body.push(Instruction::Return);
+        body.push(Instruction::End);
 
-    let mut mbuilder = builder::from_module(module);
+        let instructions = func_body.code_mut().elements_mut();
 
-    // prolog function
-    mbuilder.push_function(
-        builder::function()
-            .body()
-            .with_instructions(Instructions::new(vec![
-                //if call_depth > 256 { call_depth = 0; gr_leave(); }
-                Instruction::I32Const(0),
-                Instruction::I32Load(2, call_depth_ptr),
-                Instruction::I32Const(256),
-                Instruction::I32GtU,
-                Instruction::If(BlockType::NoResult),
-                Instruction::I32Const(0),
-                Instruction::I32Const(0),
-                Instruction::I32Store(2, call_depth_ptr),
-                Instruction::Call(gr_leave_index),
-                Instruction::End,
-                //call_depth += 1;
-                Instruction::I32Const(0),
-                Instruction::I32Const(0),
-                Instruction::I32Load(2, call_depth_ptr),
-                Instruction::I32Const(1),
-                Instruction::I32Add,
-                Instruction::I32Store(2, call_depth_ptr),
-                Instruction::End,
-            ]))
-            .build()
-            .build(),
-    );
-
-    // epilog function
-    mbuilder.push_function(
-        builder::function()
-            .body()
-            .with_instructions(Instructions::new(vec![
+        instructions.splice(
+            0..0,
+            [
+                vec![
+                    //if call_depth > 256 { call_depth = 0; return; }
+                    Instruction::I32Const(0),
+                    Instruction::I32Load(2, call_depth_ptr),
+                    Instruction::I32Const(256),
+                    Instruction::I32GtU,
+                    Instruction::If(BlockType::NoResult),
+                    Instruction::I32Const(0),
+                    Instruction::I32Const(0),
+                    Instruction::I32Store(2, call_depth_ptr),
+                ],
+                body,
+                vec![
+                    //call_depth += 1;
+                    Instruction::I32Const(0),
+                    Instruction::I32Const(0),
+                    Instruction::I32Load(2, call_depth_ptr),
+                    Instruction::I32Const(1),
+                    Instruction::I32Add,
+                    Instruction::I32Store(2, call_depth_ptr),
+                ],
+            ]
+            .concat(),
+        );
+
+        let last = instructions.len() - 1;
+        instructions.splice(
+            last..last,
+            [
                 //call_depth -= 1;
                 Instruction::I32Const(0),
                 Instruction::I32Const(0),
@@ -309,74 +309,13 @@ pub fn instrument_recursion(module: Module) -> Module {
                 Instruction::I32Const(-1),
                 Instruction::I32Add,
                 Instruction::I32Store(2, call_depth_ptr),
-                Instruction::End,
-            ]))
-            .build()
-            .build(),
-    );
-
-    // back to plain module
-    let mut module = mbuilder.build();
-
-    let Some(code_section) = module.code_section_mut() else {
-        return module;
-    };
-
-    for (i, func_body) in code_section.bodies_mut().iter_mut().enumerate() {
-        if i + import_count >= prolog_index {
-            continue;
-        }
-
-        let instructions = func_body.code_mut().elements_mut();
-
-        for instruction in instructions.iter_mut() {
-            if let Instruction::Call(call_index) = instruction {
-                if *call_index >= inserted_index {
-                    *call_index += inserted_count
-                }
-            }
-        }
-
-        instructions.insert(0, Instruction::Call(prolog_index as u32));
-
-        if let Some(end @ Instruction::End) = instructions.pop() {
-            instructions.push(Instruction::Call(epilog_index as u32));
-            instructions.push(end);
-        }
+            ],
+        );
     }
 
     let sections = module.sections_mut();
     sections.retain(|section| !matches!(section, elements::Section::Custom(_)));
 
-    for section in sections {
-        match section {
-            elements::Section::Export(export_section) => {
-                for export in export_section.entries_mut() {
-                    if let elements::Internal::Function(func_index) = export.internal_mut() {
-                        if *func_index >= inserted_index {
-                            *func_index += inserted_count
-                        }
-                    }
-                }
-            }
-            elements::Section::Element(elements_section) => {
-                for segment in elements_section.entries_mut() {
-                    for func_index in segment.members_mut() {
-                        if *func_index >= inserted_index {
-                            *func_index += inserted_count
-                        }
-                    }
-                }
-            }
-            elements::Section::Start(start_idx) => {
-                if *start_idx >= inserted_index {
-                    *start_idx += inserted_count
-                }
-            }
-            _ => {}
-        }
-    }
-
     module
 }
 

From f6de6e2e7ab6875590a94980dd0675a00b3f82b9 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Sun, 8 Oct 2023 14:50:54 +0300
Subject: [PATCH 05/17] [ci skip] cleanup

---
 utils/wasm-gen/src/utils.rs | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 8badc67663d..e69e5d79a5b 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -313,9 +313,6 @@ pub fn instrument_recursion(mut module: Module) -> Module {
         );
     }
 
-    let sections = module.sections_mut();
-    sections.retain(|section| !matches!(section, elements::Section::Custom(_)));
-
     module
 }
 

From 85dd6e711f9bc7bcbb37c066342ab441d364c0d1 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Mon, 9 Oct 2023 13:19:03 +0300
Subject: [PATCH 06/17] also instrument loops

---
 utils/wasm-gen/src/tests.rs |   5 ++
 utils/wasm-gen/src/utils.rs | 132 +++++++++++++++++++++++++++++++++---
 2 files changed, 126 insertions(+), 11 deletions(-)

diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs
index 4a195935407..cd2ffe22aaf 100644
--- a/utils/wasm-gen/src/tests.rs
+++ b/utils/wasm-gen/src/tests.rs
@@ -60,6 +60,11 @@ fn instrument_recursions() {
         (func $f (result i64)
             call $f
         )
+        (func $g
+            (loop $my_loop
+                br $my_loop
+            )
+        )
     )"#;
 
     let wasm_bytes = wat::parse_str(wat1).expect("invalid wat");
diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index e69e5d79a5b..14a6649549a 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -19,11 +19,15 @@
 #![allow(dead_code)]
 
 use crate::wasm::PageCount as WasmPageCount;
-use gear_wasm_instrument::parity_wasm::{
-    builder,
-    elements::{
-        self, BlockType, External, FuncBody, ImportCountType, Instruction, Module, Type, ValueType,
+use gear_wasm_instrument::{
+    parity_wasm::{
+        builder,
+        elements::{
+            BlockType, External, FuncBody, ImportCountType, Instruction, Instructions, Internal,
+            Module, Section, Type, ValueType,
+        },
     },
+    syscalls::SysCallName,
 };
 use gsys::HashWithValue;
 use std::{
@@ -100,7 +104,7 @@ pub fn remove_recursion(module: Module) -> Module {
                     .with_results(signature.results().to_vec())
                     .build()
                     .body()
-                    .with_instructions(elements::Instructions::new(body))
+                    .with_instructions(Instructions::new(body))
                     .build()
                     .build(),
             )
@@ -220,7 +224,7 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
-pub fn instrument_recursion(mut module: Module) -> Module {
+pub fn instrument_recursion(module: Module) -> Module {
     let Some(mem_size) = module
         .import_section()
         .and_then(|section| {
@@ -237,6 +241,48 @@ pub fn instrument_recursion(mut module: Module) -> Module {
         };
 
     let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
+    let gas_ptr = call_depth_ptr - mem::size_of::<u64>() as u32;
+
+    let maybe_gr_gas_available_index = module.import_section().and_then(|section| {
+        section
+            .entries()
+            .iter()
+            .filter(|entry| matches!(entry.external(), External::Function(_)))
+            .enumerate()
+            .find_map(|(i, entry)| {
+                (entry.module() == "env" && entry.field() == SysCallName::GasAvailable.to_str())
+                    .then_some(i as u32)
+            })
+    });
+    let rewrite_calls = maybe_gr_gas_available_index.is_none();
+
+    let (gr_gas_available_index, mut module) = match maybe_gr_gas_available_index {
+        Some(gr_gas_available_index) => (gr_gas_available_index, module),
+        None => {
+            let mut mbuilder = builder::from_module(module);
+
+            // fn gr_gas_available(gas: *mut u64);
+            let import_sig = mbuilder
+                .push_signature(builder::signature().with_param(ValueType::I32).build_sig());
+
+            mbuilder.push_import(
+                builder::import()
+                    .module("env")
+                    .field(SysCallName::GasAvailable.to_str())
+                    .external()
+                    .func(import_sig)
+                    .build(),
+            );
+
+            // back to plain module
+            let module = mbuilder.build();
+
+            let import_count = module.import_count(ImportCountType::Function);
+            let gr_gas_available_index = import_count as u32 - 1;
+
+            (gr_gas_available_index, module)
+        }
+    };
 
     let (Some(type_section), Some(function_section)) = (module.type_section(), module.function_section()) else {
         return module;
@@ -265,8 +311,7 @@ pub fn instrument_recursion(mut module: Module) -> Module {
             body.push(instruction);
         }
 
-        body.push(Instruction::Return);
-        body.push(Instruction::End);
+        body.extend([Instruction::Return, Instruction::End]);
 
         let instructions = func_body.code_mut().elements_mut();
 
@@ -274,17 +319,17 @@ pub fn instrument_recursion(mut module: Module) -> Module {
             0..0,
             [
                 vec![
-                    //if call_depth > 256 { call_depth = 0; return; }
+                    //if call_depth > 512 { call_depth = 0; return; }
                     Instruction::I32Const(0),
                     Instruction::I32Load(2, call_depth_ptr),
-                    Instruction::I32Const(256),
+                    Instruction::I32Const(512),
                     Instruction::I32GtU,
                     Instruction::If(BlockType::NoResult),
                     Instruction::I32Const(0),
                     Instruction::I32Const(0),
                     Instruction::I32Store(2, call_depth_ptr),
                 ],
-                body,
+                body.clone(),
                 vec![
                     //call_depth += 1;
                     Instruction::I32Const(0),
@@ -311,6 +356,71 @@ pub fn instrument_recursion(mut module: Module) -> Module {
                 Instruction::I32Store(2, call_depth_ptr),
             ],
         );
+
+        let original_instructions =
+            mem::replace(instructions, Vec::with_capacity(instructions.len()));
+        let new_instructions = instructions;
+
+        for instruction in original_instructions {
+            match instruction {
+                Instruction::Loop(_) => {
+                    new_instructions.push(instruction);
+
+                    new_instructions.extend([
+                        Instruction::I32Const(gas_ptr as i32),
+                        Instruction::Call(gr_gas_available_index),
+                        Instruction::I32Const(gas_ptr as i32),
+                        Instruction::I64Load(3, 0),
+                        Instruction::I64Const(1_000_000),
+                        Instruction::I64LeU,
+                        Instruction::If(BlockType::NoResult),
+                    ]);
+                    new_instructions.extend(body.clone());
+                }
+                Instruction::Call(call_index)
+                    if rewrite_calls && call_index >= gr_gas_available_index =>
+                {
+                    new_instructions.push(Instruction::Call(call_index + 1));
+                }
+                _ => {
+                    new_instructions.push(instruction);
+                }
+            }
+        }
+    }
+
+    if rewrite_calls {
+        let sections = module.sections_mut();
+        sections.retain(|section| !matches!(section, Section::Custom(_)));
+
+        for section in sections {
+            match section {
+                Section::Export(export_section) => {
+                    for export in export_section.entries_mut() {
+                        if let Internal::Function(func_index) = export.internal_mut() {
+                            if *func_index >= gr_gas_available_index {
+                                *func_index += 1
+                            }
+                        }
+                    }
+                }
+                Section::Element(elements_section) => {
+                    for segment in elements_section.entries_mut() {
+                        for func_index in segment.members_mut() {
+                            if *func_index >= gr_gas_available_index {
+                                *func_index += 1
+                            }
+                        }
+                    }
+                }
+                Section::Start(start_idx) => {
+                    if *start_idx >= gr_gas_available_index {
+                        *start_idx += 1;
+                    }
+                }
+                _ => {}
+            }
+        }
     }
 
     module

From 1517c2062d1be4282269eedf72c0aa78d828b568 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Mon, 9 Oct 2023 15:35:38 +0300
Subject: [PATCH 07/17] [ci skip] enable remove_recursion

---
 utils/wasm-gen/src/config.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index a8edb35f91a..beb39d7933d 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -164,7 +164,7 @@ impl<T> Default for StandardGearWasmConfigsBundle<T> {
         Self {
             log_info: Some("StandardGearWasmConfigsBundle".into()),
             existing_addresses: None,
-            remove_recursion: false,
+            remove_recursion: true,
             call_indirect_enabled: true,
             injection_types: SysCallsInjectionTypes::all_once(),
             entry_points_set: Default::default(),

From 7256ce6ee85f5b4c2c980f677dbe8d378d77a932 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Fri, 20 Oct 2023 19:09:58 +0300
Subject: [PATCH 08/17] fix fmt

---
 utils/wasm-gen/src/utils.rs | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 14a6649549a..f9a6d383ce1 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -236,9 +236,11 @@ pub fn instrument_recursion(module: Module) -> Module {
                     _ => None,
                 })
         })
-        .map(Into::<WasmPageCount>::into).map(|page_count| page_count.memory_size()) else {
-            return module;
-        };
+        .map(Into::<WasmPageCount>::into)
+        .map(|page_count| page_count.memory_size())
+    else {
+        return module;
+    };
 
     let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
     let gas_ptr = call_depth_ptr - mem::size_of::<u64>() as u32;
@@ -284,7 +286,9 @@ pub fn instrument_recursion(module: Module) -> Module {
         }
     };
 
-    let (Some(type_section), Some(function_section)) = (module.type_section(), module.function_section()) else {
+    let (Some(type_section), Some(function_section)) =
+        (module.type_section(), module.function_section())
+    else {
         return module;
     };
 

From da9cb0128c8c5ce416497f3635dac1ba7009f233 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Thu, 9 Nov 2023 13:34:07 +0300
Subject: [PATCH 09/17] add changes

---
 utils/wasm-gen/src/config.rs |  8 ++++
 utils/wasm-gen/src/tests.rs  |  3 --
 utils/wasm-gen/src/utils.rs  | 76 +++++++++---------------------------
 3 files changed, 26 insertions(+), 61 deletions(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index 9fba80a8757..020e51de22a 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -103,7 +103,10 @@ pub use generator::*;
 pub use module::*;
 pub use syscalls::*;
 
+use crate::InvocableSysCall;
+
 use gear_utils::NonEmpty;
+use gear_wasm_instrument::syscalls::SysCallName;
 use gsys::Hash;
 
 /// Trait which describes a type that stores and manages data for generating
@@ -198,6 +201,11 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
             ..SelectableParams::default()
         };
 
+        let mut injection_types = injection_types;
+        if remove_recursion {
+            injection_types.enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
+        }
+
         let mut syscalls_config_builder = SysCallsConfigBuilder::new(injection_types);
         if let Some(log_info) = log_info {
             syscalls_config_builder = syscalls_config_builder.with_log_info(log_info);
diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs
index 0f2c6fe54f6..98782e04a7f 100644
--- a/utils/wasm-gen/src/tests.rs
+++ b/utils/wasm-gen/src/tests.rs
@@ -79,9 +79,6 @@ fn instrument_recursions() {
 
     let wat = wasmprinter::print_bytes(&wasm_bytes).expect("failed printing bytes");
     println!("wat = {wat}");
-
-    let code_res = Code::try_new(wasm_bytes, 1, |_| CustomConstantCostRules::default(), None);
-    assert!(code_res.is_ok());
 }
 
 #[test]
diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index f9a6d383ce1..6faa8119cd7 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -242,8 +242,7 @@ pub fn instrument_recursion(module: Module) -> Module {
         return module;
     };
 
-    let call_depth_ptr = mem_size - mem::size_of::<u32>() as u32;
-    let gas_ptr = call_depth_ptr - mem::size_of::<u64>() as u32;
+    let gas_ptr = mem_size - mem::size_of::<u64>() as u32;
 
     let maybe_gr_gas_available_index = module.import_section().and_then(|section| {
         section
@@ -304,7 +303,18 @@ pub fn instrument_recursion(module: Module) -> Module {
         let signature = &types[signature_index.type_ref() as usize];
         let Type::Function(signature) = signature;
         let results = signature.results();
-        let mut body = Vec::with_capacity(results.len() + 2);
+
+        let mut body = Vec::with_capacity(results.len() + 9);
+        body.extend_from_slice(&[
+            Instruction::I32Const(gas_ptr as i32),
+            Instruction::Call(gr_gas_available_index),
+            Instruction::I32Const(gas_ptr as i32),
+            Instruction::I64Load(3, 0),
+            Instruction::I64Const(1_000_000),
+            Instruction::I64LeU,
+            Instruction::If(BlockType::NoResult),
+        ]);
+
         for result in results {
             let instruction = match result {
                 ValueType::I32 => Instruction::I32Const(u32::MAX as i32),
@@ -315,71 +325,21 @@ pub fn instrument_recursion(module: Module) -> Module {
             body.push(instruction);
         }
 
-        body.extend([Instruction::Return, Instruction::End]);
+        body.extend_from_slice(&[Instruction::Return, Instruction::End]);
 
         let instructions = func_body.code_mut().elements_mut();
 
-        instructions.splice(
-            0..0,
-            [
-                vec![
-                    //if call_depth > 512 { call_depth = 0; return; }
-                    Instruction::I32Const(0),
-                    Instruction::I32Load(2, call_depth_ptr),
-                    Instruction::I32Const(512),
-                    Instruction::I32GtU,
-                    Instruction::If(BlockType::NoResult),
-                    Instruction::I32Const(0),
-                    Instruction::I32Const(0),
-                    Instruction::I32Store(2, call_depth_ptr),
-                ],
-                body.clone(),
-                vec![
-                    //call_depth += 1;
-                    Instruction::I32Const(0),
-                    Instruction::I32Const(0),
-                    Instruction::I32Load(2, call_depth_ptr),
-                    Instruction::I32Const(1),
-                    Instruction::I32Add,
-                    Instruction::I32Store(2, call_depth_ptr),
-                ],
-            ]
-            .concat(),
-        );
-
-        let last = instructions.len() - 1;
-        instructions.splice(
-            last..last,
-            [
-                //call_depth -= 1;
-                Instruction::I32Const(0),
-                Instruction::I32Const(0),
-                Instruction::I32Load(2, call_depth_ptr),
-                Instruction::I32Const(-1),
-                Instruction::I32Add,
-                Instruction::I32Store(2, call_depth_ptr),
-            ],
-        );
-
         let original_instructions =
             mem::replace(instructions, Vec::with_capacity(instructions.len()));
         let new_instructions = instructions;
 
+        new_instructions.extend_from_slice(&body);
+
         for instruction in original_instructions {
             match instruction {
-                Instruction::Loop(_) => {
+                Instruction::Block(_) | Instruction::Loop(_) | Instruction::If(_) => {
                     new_instructions.push(instruction);
-
-                    new_instructions.extend([
-                        Instruction::I32Const(gas_ptr as i32),
-                        Instruction::Call(gr_gas_available_index),
-                        Instruction::I32Const(gas_ptr as i32),
-                        Instruction::I64Load(3, 0),
-                        Instruction::I64Const(1_000_000),
-                        Instruction::I64LeU,
-                        Instruction::If(BlockType::NoResult),
-                    ]);
-                    new_instructions.extend(body.clone());
+                    new_instructions.extend_from_slice(&body);
                 }
                 Instruction::Call(call_index)
                     if rewrite_calls && call_index >= gr_gas_available_index =>

From 87b97dc8f091210bd5ae2d50ee78700efe2e9d17 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Thu, 9 Nov 2023 14:23:39 +0300
Subject: [PATCH 10/17] fix cargo fmt

---
 utils/wasm-gen/src/config.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index 4ba0e58f0db..131c78b412e 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -187,7 +187,8 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
 
         let mut injection_types = injection_types;
         if remove_recursion {
-            injection_types.enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
+            injection_types
+                .enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
         }
 
         let mut syscalls_config_builder = SysCallsConfigBuilder::new(injection_types);

From 4802c06ddacd8b0e1a0d896d62c95c060f42e0d9 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Mon, 13 Nov 2023 18:12:23 +0300
Subject: [PATCH 11/17] refactoring

---
 utils/wasm-gen/src/config.rs           | 14 ++++++++++++--
 utils/wasm-gen/src/config/generator.rs | 10 ++++++++++
 utils/wasm-gen/src/generator.rs        | 10 +++++++++-
 utils/wasm-gen/src/tests.rs            |  6 +++---
 utils/wasm-gen/src/utils.rs            |  6 ++----
 5 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index 131c78b412e..5eca31bcd4e 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -143,6 +143,13 @@ pub struct StandardGearWasmConfigsBundle<T = [u8; 32]> {
     pub existing_addresses: Option<NonEmpty<T>>,
     /// Flag which signals whether recursions must be removed.
     pub remove_recursion: bool,
+    /// If the limit is set to `Some(_)`, programs will try to stop execution
+    /// after reaching a critical gas limit, which can be useful to exit from
+    /// heavy loops and recursions that waste all gas.
+    ///
+    /// The `gr_gas_available` syscall is called at the beginning of each
+    /// function and for each control instruction (blocks, loops, conditions).
+    pub critical_gas_limit: Option<u64>,
     /// Injection type for each syscall.
     pub injection_types: SysCallsInjectionTypes,
     /// Config of gear wasm call entry-points (exports).
@@ -160,7 +167,8 @@ impl<T> Default for StandardGearWasmConfigsBundle<T> {
         Self {
             log_info: Some("StandardGearWasmConfigsBundle".into()),
             existing_addresses: None,
-            remove_recursion: true,
+            remove_recursion: false,
+            critical_gas_limit: Some(1_000_000),
             injection_types: SysCallsInjectionTypes::all_once(),
             entry_points_set: Default::default(),
             initial_pages: DEFAULT_INITIAL_SIZE,
@@ -176,6 +184,7 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
             log_info,
             existing_addresses,
             remove_recursion,
+            critical_gas_limit,
             injection_types,
             entry_points_set,
             initial_pages,
@@ -186,7 +195,7 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
         let selectable_params = SelectableParams::default();
 
         let mut injection_types = injection_types;
-        if remove_recursion {
+        if critical_gas_limit.is_some() {
             injection_types
                 .enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
         }
@@ -208,6 +217,7 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
             upper_limit: None,
         };
         let gear_wasm_generator_config = GearWasmGeneratorConfigBuilder::new()
+            .with_critical_gas_limit(critical_gas_limit)
             .with_recursions_removed(remove_recursion)
             .with_syscalls_config(syscalls_config_builder.build())
             .with_entry_points_config(entry_points_set)
diff --git a/utils/wasm-gen/src/config/generator.rs b/utils/wasm-gen/src/config/generator.rs
index 2c1da46750f..2b47e0058f7 100644
--- a/utils/wasm-gen/src/config/generator.rs
+++ b/utils/wasm-gen/src/config/generator.rs
@@ -60,6 +60,13 @@ impl GearWasmGeneratorConfigBuilder {
         self
     }
 
+    /// Defines whether programs should have a critical gas limit.
+    pub fn with_critical_gas_limit(mut self, critical_gas_limit: Option<u64>) -> Self {
+        self.0.critical_gas_limit = critical_gas_limit;
+
+        self
+    }
+
     /// Build the gear wasm generator.
     pub fn build(self) -> GearWasmGeneratorConfig {
         self.0
@@ -81,6 +88,9 @@ pub struct GearWasmGeneratorConfig {
     /// Flag, signalizing whether recursions
     /// should be removed from resulting module.
     pub remove_recursions: bool,
+    /// The critical gas limit after which the program
+    /// will attempt to terminate successfully.
+    pub critical_gas_limit: Option<u64>,
 }
 
 /// Memory pages config used by [`crate::MemoryGenerator`].
diff --git a/utils/wasm-gen/src/generator.rs b/utils/wasm-gen/src/generator.rs
index 88dcbcd066b..dd6642f6708 100644
--- a/utils/wasm-gen/src/generator.rs
+++ b/utils/wasm-gen/src/generator.rs
@@ -139,9 +139,17 @@ impl<'a, 'b> GearWasmGenerator<'a, 'b> {
             .into_wasm_module()
             .into_inner();
 
+        let module = if let Some(critical_gas_limit) = config.critical_gas_limit {
+            log::trace!("Injecting critical gas limit");
+            utils::inject_critical_gas_limit(module, critical_gas_limit)
+        } else {
+            log::trace!("Critical gas limit is not set");
+            module
+        };
+
         Ok(if config.remove_recursions {
             log::trace!("Removing recursions");
-            utils::instrument_recursion(module)
+            utils::remove_recursion(module)
         } else {
             log::trace!("Leaving recursions");
             module
diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs
index a60e2bf5d47..03926826690 100644
--- a/utils/wasm-gen/src/tests.rs
+++ b/utils/wasm-gen/src/tests.rs
@@ -48,7 +48,7 @@ use std::{mem, num::NonZeroUsize};
 const UNSTRUCTURED_SIZE: usize = 1_000_000;
 
 #[test]
-fn instrument_recursions() {
+fn inject_critical_gas_limit_works() {
     let wat1 = r#"
     (module
         (memory $memory0 (import "env" "memory") 16)
@@ -70,9 +70,9 @@ fn instrument_recursions() {
     let wasm_bytes = wat::parse_str(wat1).expect("invalid wat");
     let module =
         parity_wasm::deserialize_buffer::<Module>(&wasm_bytes).expect("invalid wasm bytes");
-    let limited_recursions_module = utils::instrument_recursion(module);
+    let module_with_critical_gas_limit = utils::inject_critical_gas_limit(module, 1_000_000);
 
-    let wasm_bytes = limited_recursions_module
+    let wasm_bytes = module_with_critical_gas_limit
         .into_bytes()
         .expect("invalid pw module");
     assert!(wasmparser::validate(&wasm_bytes).is_ok());
diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 6faa8119cd7..b66a22a4963 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -16,8 +16,6 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-#![allow(dead_code)]
-
 use crate::wasm::PageCount as WasmPageCount;
 use gear_wasm_instrument::{
     parity_wasm::{
@@ -224,7 +222,7 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
-pub fn instrument_recursion(module: Module) -> Module {
+pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Module {
     let Some(mem_size) = module
         .import_section()
         .and_then(|section| {
@@ -310,7 +308,7 @@ pub fn instrument_recursion(module: Module) -> Module {
             Instruction::Call(gr_gas_available_index),
             Instruction::I32Const(gas_ptr as i32),
             Instruction::I64Load(3, 0),
-            Instruction::I64Const(1_000_000),
+            Instruction::I64Const(critical_gas_limit as i64),
             Instruction::I64LeU,
             Instruction::If(BlockType::NoResult),
         ]);

From 3bd71a81ed4d87df77183fbe1a664ed7cc00e0e9 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Mon, 13 Nov 2023 18:27:41 +0300
Subject: [PATCH 12/17] fix doc test

---
 utils/wasm-gen/src/config.rs | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index 5eca31bcd4e..be8fab2faeb 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -62,7 +62,8 @@
 //! let wasm_gen_config = GearWasmGeneratorConfig {
 //!     memory_config: memory_pages_config,
 //!     entry_points_config: entry_points_set,
-//!     remove_recursions: true,
+//!     remove_recursions: false,
+//!     critical_gas_limit: Some(1_000_000),
 //!     syscalls_config,
 //! };
 //! ```

From d793944027c1b23af747be9f3ded0bb7f6562990 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 15 Nov 2023 17:12:17 +0300
Subject: [PATCH 13/17] move enabling gr_gas_available syscall to config

---
 utils/wasm-gen/src/config.rs          | 12 +++---------
 utils/wasm-gen/src/config/syscalls.rs |  9 +++++++++
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index be8fab2faeb..b30417d0274 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -102,10 +102,7 @@ pub use generator::*;
 pub use module::*;
 pub use syscalls::*;
 
-use crate::InvocableSysCall;
-
 use gear_utils::NonEmpty;
-use gear_wasm_instrument::syscalls::SysCallName;
 use gsys::Hash;
 
 /// Trait which describes a type that stores and manages data for generating
@@ -195,12 +192,6 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
 
         let selectable_params = SelectableParams::default();
 
-        let mut injection_types = injection_types;
-        if critical_gas_limit.is_some() {
-            injection_types
-                .enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
-        }
-
         let mut syscalls_config_builder = SysCallsConfigBuilder::new(injection_types);
         if let Some(log_info) = log_info {
             syscalls_config_builder = syscalls_config_builder.with_log_info(log_info);
@@ -210,6 +201,9 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
         } else {
             syscalls_config_builder = syscalls_config_builder.with_source_msg_dest();
         }
+        if critical_gas_limit.is_some() {
+            syscalls_config_builder = syscalls_config_builder.with_critical_gas_limit();
+        }
         syscalls_config_builder = syscalls_config_builder.with_params_config(params_config);
 
         let memory_pages_config = MemoryPagesConfig {
diff --git a/utils/wasm-gen/src/config/syscalls.rs b/utils/wasm-gen/src/config/syscalls.rs
index 143f0ec826f..6d19b6d2ec6 100644
--- a/utils/wasm-gen/src/config/syscalls.rs
+++ b/utils/wasm-gen/src/config/syscalls.rs
@@ -104,6 +104,15 @@ impl SysCallsConfigBuilder {
         self
     }
 
+    /// Set whether programs should have a critical gas limit.
+    pub fn with_critical_gas_limit(mut self) -> Self {
+        self.0
+            .injection_types
+            .enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
+
+        self
+    }
+
     /// Setup fallible syscalls error processing options.
     pub fn set_error_processing_config(mut self, config: ErrorProcessingConfig) -> Self {
         self.0.error_processing_config = config;

From 00bf3602a70ad63d0f7a1afa14297c52946765a9 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 15 Nov 2023 17:36:10 +0300
Subject: [PATCH 14/17] more docs

---
 utils/wasm-gen/src/utils.rs | 50 ++++++++++++++++++++++++++++++++++---
 1 file changed, 47 insertions(+), 3 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index b66a22a4963..88e57906e4a 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -222,7 +222,38 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
+/// Injects a critical gas limit to a given wasm module.
+///
+/// Code before injection gas limiter:
+/// ```no_run
+/// fn func() {
+///     func();
+///     loop { }
+/// }
+/// ```
+/// 
+/// Code after injection gas limiter:
+/// ```no_run
+/// use gcore::exec;
+///
+/// const CRITICAL_GAS_LIMIT: u64 = 1_000_000;
+/// 
+/// fn func() {
+///     // exit from recursions
+///     if exec::gas_available() <= CRITICAL_GAS_LIMIT {
+///         return;
+///     }
+///     func();
+///     loop {
+///         // exit from heavy loops
+///         if exec::gas_available() <= CRITICAL_GAS_LIMIT {
+///             return;
+///         }
+///     }
+/// }
+/// ```
 pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Module {
+    // get initial memory size of program
     let Some(mem_size) = module
         .import_section()
         .and_then(|section| {
@@ -240,8 +271,10 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
         return module;
     };
 
+    // store available gas pointer on the last memory page
     let gas_ptr = mem_size - mem::size_of::<u64>() as u32;
 
+    // add gr_gas_available import if needed
     let maybe_gr_gas_available_index = module.import_section().and_then(|section| {
         section
             .entries()
@@ -253,7 +286,8 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
                     .then_some(i as u32)
             })
     });
-    let rewrite_calls = maybe_gr_gas_available_index.is_none();
+    // sections should only be rewritten if the module did not previously have gr_gas_available import
+    let rewrite_sections = maybe_gr_gas_available_index.is_none();
 
     let (gr_gas_available_index, mut module) = match maybe_gr_gas_available_index {
         Some(gr_gas_available_index) => (gr_gas_available_index, module),
@@ -302,17 +336,22 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
         let Type::Function(signature) = signature;
         let results = signature.results();
 
+        // create the body of the gas limiter:
         let mut body = Vec::with_capacity(results.len() + 9);
         body.extend_from_slice(&[
+            // gr_gas_available(gas_ptr)
             Instruction::I32Const(gas_ptr as i32),
             Instruction::Call(gr_gas_available_index),
+            // gas_available = *gas_ptr
             Instruction::I32Const(gas_ptr as i32),
             Instruction::I64Load(3, 0),
             Instruction::I64Const(critical_gas_limit as i64),
+            // if gas_available <= critical_gas_limit { return result; }
             Instruction::I64LeU,
             Instruction::If(BlockType::NoResult),
         ]);
 
+        // exit the current function with dummy results
         for result in results {
             let instruction = match result {
                 ValueType::I32 => Instruction::I32Const(u32::MAX as i32),
@@ -331,8 +370,11 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
             mem::replace(instructions, Vec::with_capacity(instructions.len()));
         let new_instructions = instructions;
 
+        // insert gas limiter at the beginning of each function to limit recursions
         new_instructions.extend_from_slice(&body);
 
+        // also insert gas limiter at the beginning of each block, loop and condition
+        // to limit control instructions
         for instruction in original_instructions {
             match instruction {
                 Instruction::Block(_) | Instruction::Loop(_) | Instruction::If(_) => {
@@ -340,8 +382,9 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
                     new_instructions.extend_from_slice(&body);
                 }
                 Instruction::Call(call_index)
-                    if rewrite_calls && call_index >= gr_gas_available_index =>
+                    if rewrite_sections && call_index >= gr_gas_available_index =>
                 {
+                    // fix function indexes if import gr_gas_available was inserted
                     new_instructions.push(Instruction::Call(call_index + 1));
                 }
                 _ => {
@@ -351,7 +394,8 @@ pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Mod
         }
     }
 
-    if rewrite_calls {
+    // fix other sections if import gr_gas_available was inserted
+    if rewrite_sections {
         let sections = module.sections_mut();
         sections.retain(|section| !matches!(section, Section::Custom(_)));
 

From 7d42eb02f391a6d1bad7bbc6337b63f3bb4159fc Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 15 Nov 2023 17:36:48 +0300
Subject: [PATCH 15/17] fix fmt

---
 utils/wasm-gen/src/utils.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 88e57906e4a..2d82529a871 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -231,13 +231,13 @@ fn find_recursion_impl<Callback>(
 ///     loop { }
 /// }
 /// ```
-/// 
+///
 /// Code after injection gas limiter:
 /// ```no_run
 /// use gcore::exec;
 ///
 /// const CRITICAL_GAS_LIMIT: u64 = 1_000_000;
-/// 
+///
 /// fn func() {
 ///     // exit from recursions
 ///     if exec::gas_available() <= CRITICAL_GAS_LIMIT {

From 4a4cf6e57de56b0cbcfe517dd01a2e39446ac305 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Wed, 15 Nov 2023 18:01:46 +0300
Subject: [PATCH 16/17] fixes

---
 utils/wasm-gen/src/utils.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs
index 2d82529a871..0dc4b1e510c 100644
--- a/utils/wasm-gen/src/utils.rs
+++ b/utils/wasm-gen/src/utils.rs
@@ -225,7 +225,7 @@ fn find_recursion_impl<Callback>(
 /// Injects a critical gas limit to a given wasm module.
 ///
 /// Code before injection gas limiter:
-/// ```no_run
+/// ```ignore
 /// fn func() {
 ///     func();
 ///     loop { }
@@ -233,7 +233,7 @@ fn find_recursion_impl<Callback>(
 /// ```
 ///
 /// Code after injection gas limiter:
-/// ```no_run
+/// ```ignore
 /// use gcore::exec;
 ///
 /// const CRITICAL_GAS_LIMIT: u64 = 1_000_000;

From 60a89cbfd35a0558ff09818293f930a09dda0630 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Fri, 17 Nov 2023 13:12:41 +0300
Subject: [PATCH 17/17] revert changes

---
 utils/wasm-gen/src/config.rs          | 3 ---
 utils/wasm-gen/src/config/syscalls.rs | 9 ---------
 2 files changed, 12 deletions(-)

diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs
index b30417d0274..b715baf62b8 100644
--- a/utils/wasm-gen/src/config.rs
+++ b/utils/wasm-gen/src/config.rs
@@ -201,9 +201,6 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
         } else {
             syscalls_config_builder = syscalls_config_builder.with_source_msg_dest();
         }
-        if critical_gas_limit.is_some() {
-            syscalls_config_builder = syscalls_config_builder.with_critical_gas_limit();
-        }
         syscalls_config_builder = syscalls_config_builder.with_params_config(params_config);
 
         let memory_pages_config = MemoryPagesConfig {
diff --git a/utils/wasm-gen/src/config/syscalls.rs b/utils/wasm-gen/src/config/syscalls.rs
index 6d19b6d2ec6..143f0ec826f 100644
--- a/utils/wasm-gen/src/config/syscalls.rs
+++ b/utils/wasm-gen/src/config/syscalls.rs
@@ -104,15 +104,6 @@ impl SysCallsConfigBuilder {
         self
     }
 
-    /// Set whether programs should have a critical gas limit.
-    pub fn with_critical_gas_limit(mut self) -> Self {
-        self.0
-            .injection_types
-            .enable_syscall_import(InvocableSysCall::Loose(SysCallName::GasAvailable));
-
-        self
-    }
-
     /// Setup fallible syscalls error processing options.
     pub fn set_error_processing_config(mut self, config: ErrorProcessingConfig) -> Self {
         self.0.error_processing_config = config;