From b708f2e737cfa1ef2d4ea506abfce334eeb6c342 Mon Sep 17 00:00:00 2001
From: StackOverflowExcept1on
 <109800286+StackOverflowExcept1on@users.noreply.github.com>
Date: Tue, 16 Jan 2024 23:31:01 +0300
Subject: [PATCH] use % == 0, add extra check for init_called

---
 utils/wasm-gen/src/generator/entry_points.rs  | 29 +++++++++++++++++--
 .../src/generator/syscalls/invocator.rs       | 11 +++++--
 2 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/utils/wasm-gen/src/generator/entry_points.rs b/utils/wasm-gen/src/generator/entry_points.rs
index 94c8dfb287e..224bd18ebcb 100644
--- a/utils/wasm-gen/src/generator/entry_points.rs
+++ b/utils/wasm-gen/src/generator/entry_points.rs
@@ -20,13 +20,15 @@
 
 use crate::{
     generator::{CallIndexes, FrozenGearWasmGenerator, GearWasmGenerator, ModuleWithCallIndexes},
-    EntryPointsSet, WasmModule,
+    wasm::{PageCount as WasmPageCount, WasmModule},
+    EntryPointsSet,
 };
 use arbitrary::{Result, Unstructured};
 use gear_wasm_instrument::parity_wasm::{
     builder,
     elements::{FunctionType, Instruction, Instructions, Type, ValueType},
 };
+use std::mem;
 
 /// Gear wasm entry points generator.
 ///
@@ -164,8 +166,31 @@ impl<'a, 'b> EntryPointsGenerator<'a, 'b> {
             (module, func_type)
         });
 
-        let export_body_instructions =
+        let mut export_body_instructions =
             self.generate_export_body(export_body_call_idx, export_body_call_func_type)?;
+
+        // after initializing the program, we will write about this in a special pointer
+        if name == "init" {
+            if let Some(memory_size_pages) = self.module.initial_mem_size() {
+                let mem_size = Into::<WasmPageCount>::into(memory_size_pages).memory_size();
+
+                let upper_limit = mem_size.saturating_sub(1) as i32;
+                let wait_called_ptr = upper_limit.saturating_sub(50);
+                let init_called_ptr = wait_called_ptr + mem::size_of::<bool>() as i32;
+
+                let end_instruction = export_body_instructions.pop();
+                export_body_instructions.extend_from_slice(&[
+                    // *init_called_ptr = true
+                    Instruction::I32Const(init_called_ptr),
+                    Instruction::I32Const(1),
+                    Instruction::I32Store8(0, 0),
+                ]);
+                if let Some(end_instruction) = end_instruction {
+                    export_body_instructions.push(end_instruction);
+                }
+            }
+        }
+
         self.module.with(|module| {
             let module = builder::from_module(module)
                 .function()
diff --git a/utils/wasm-gen/src/generator/syscalls/invocator.rs b/utils/wasm-gen/src/generator/syscalls/invocator.rs
index f316b8d513b..7dda86ace60 100644
--- a/utils/wasm-gen/src/generator/syscalls/invocator.rs
+++ b/utils/wasm-gen/src/generator/syscalls/invocator.rs
@@ -338,17 +338,22 @@ impl<'a, 'b> SyscallsInvocator<'a, 'b> {
 
             let upper_limit = mem_size.saturating_sub(1) as i32;
             let wait_called_ptr = upper_limit.saturating_sub(50);
+            let init_called_ptr = wait_called_ptr + mem::size_of::<bool>() as i32;
 
             // add instructions before calling wait syscall
             instructions.splice(
                 0..0,
                 [
+                    Instruction::I32Const(init_called_ptr),
+                    Instruction::I32Load8U(0, 0),
+                    // if *init_called_ptr { .. }
+                    Instruction::If(BlockType::NoResult),
                     Instruction::I32Const(wait_called_ptr),
                     Instruction::I32Load(2, 0),
                     Instruction::I32Const(wait_frequency as i32),
                     Instruction::I32RemU,
-                    Instruction::I32Const(1),
-                    Instruction::I32Eq,
+                    Instruction::I32Eqz,
+                    // if *wait_called_ptr % wait_frequency == 0 { orig_wait_syscall(); }
                     Instruction::If(BlockType::NoResult),
                 ],
             );
@@ -356,12 +361,14 @@ impl<'a, 'b> SyscallsInvocator<'a, 'b> {
             // add instructions after calling wait syscall
             instructions.extend_from_slice(&[
                 Instruction::End,
+                // *wait_called_ptr += 1
                 Instruction::I32Const(wait_called_ptr),
                 Instruction::I32Const(wait_called_ptr),
                 Instruction::I32Load(2, 0),
                 Instruction::I32Const(1),
                 Instruction::I32Add,
                 Instruction::I32Store(2, 0),
+                Instruction::End,
             ]);
         }