fix page_size check

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ic-wasm"
-version = "0.1.4"
+version = "0.1.5"
 authors = ["DFINITY Stiftung"]
 edition = "2021"
 description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer"
@@ -22,7 +22,7 @@ required-features = ["exe"]
 walrus = "0.19.0"
 candid = "0.7"
 
-ic-types = "0.4.0"
+ic-types = "0.4.2"
 
 anyhow = { version = "1.0.34", optional = true }
 clap = { version = "3.0.14", features = ["derive", "cargo"], optional = true }

src/instrumentation.rs

@@ -73,6 +73,7 @@ pub fn instrument(m: &mut Module) {
     let leb = make_leb128_encoder(m);
     make_stable_getter(m, &vars, leb);
     make_getter(m, &vars);
+    make_toggle_tracing(m, &vars);
     let name = make_name_section(m);
     m.customs.add(name);
 }
@@ -337,16 +338,6 @@ fn make_stable_writer(m: &mut Module, vars: &Variables) -> FunctionId {
     let size = m.locals.add(ValType::I32);
     builder
         .func_body()
-        .global_get(vars.page_size)
-        .i32_const(32)
-        .binop(BinaryOp::I32GtS) // trace >= 2M
-        .if_else(
-            None,
-            |then| {
-                then.return_();
-            },
-            |_| {},
-        )
         .local_get(offset)
         .local_get(size)
         .binop(BinaryOp::I32Add)
@@ -358,13 +349,25 @@ fn make_stable_writer(m: &mut Module, vars: &Variables) -> FunctionId {
             None,
             |then| {
                 // TODO: This assumes user code doesn't use stable memory
-                then.i32_const(1)
-                    .call(grow)
-                    .drop()
-                    .global_get(vars.page_size)
-                    .i32_const(1)
-                    .binop(BinaryOp::I32Add)
-                    .global_set(vars.page_size);
+                then.global_get(vars.page_size)
+                    .i32_const(32)
+                    .binop(BinaryOp::I32GtS) // trace >= 2M
+                    .if_else(
+                        None,
+                        |then| {
+                            then.return_();
+                        },
+                        |else_| {
+                            else_
+                                .i32_const(1)
+                                .call(grow)
+                                .drop()
+                                .global_get(vars.page_size)
+                                .i32_const(1)
+                                .binop(BinaryOp::I32Add)
+                                .global_set(vars.page_size);
+                        },
+                    );
             },
             |_| {},
         )
@@ -587,3 +590,33 @@ fn make_getter(m: &mut Module, vars: &Variables) {
     let getter = getter.finish(vec![], &mut m.funcs);
     m.exports.add("canister_query __get_cycles", getter);
 }
+fn make_toggle_tracing(m: &mut Module, vars: &Variables) {
+    let memory = get_memory_id(m);
+    let reply_data = get_ic_func_id(m, "msg_reply_data_append");
+    let reply = get_ic_func_id(m, "msg_reply");
+    let tmp = m.locals.add(ValType::I64);
+    let mut builder = FunctionBuilder::new(&mut m.types, &[], &[]);
+    builder.name("__toggle_tracing".to_string());
+    #[rustfmt::skip]
+    builder
+        .func_body()
+        .global_get(vars.is_init)
+        .i32_const(1)
+        .binop(BinaryOp::I32Xor)
+        .global_set(vars.is_init)
+        .i32_const(0)
+        .load(memory, LoadKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
+        .local_set(tmp)
+        .i32_const(0)
+        .i64_const(0x4c444944) // "DIDL0000xxxx"
+        .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
+        .i32_const(0)
+        .i32_const(6)
+        .call(reply_data)
+        .i32_const(0)
+        .local_get(tmp)
+        .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
+        .call(reply);
+    let id = builder.finish(vec![], &mut m.funcs);
+    m.exports.add("canister_update __toggle_tracing", id);
+}

src/utils.rs

@@ -1,7 +1,7 @@
 use std::collections::HashMap;
 use walrus::*;
 
-#[derive(Clone, Copy, PartialEq)]
+#[derive(Clone, Copy, PartialEq, Eq)]
 pub enum InjectionKind {
     Static,
     Dynamic,