use 64bit API for stable memory

dfinity · Sep 26, 2024 · 1f16e57 · 1f16e57
1 parent 12f4cb7
commit 1f16e57
Show file tree

Hide file tree

Showing 10 changed files with 117 additions and 145 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ic-wasm"
-version = "0.8.6"
+version = "0.9.0"
 authors = ["DFINITY Stiftung"]
 edition = "2021"
 description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer"
@@ -20,7 +20,7 @@ required-features = ["exe"]
 [dependencies]
 # Major version bump of walrus should result in a major version bump of ic-wasm.
 # Because we expose walrus types in ic-wasm public API.
-walrus = "0.21.1" 
+walrus = "0.22.0" 
 candid = "0.10"
 rustc-demangle = "0.1"
 thiserror = "1.0.35"

diff --git a/src/bin/main.rs b/src/bin/main.rs
@@ -90,7 +90,7 @@ enum SubCommand {
         trace_only: Option<Vec<String>>,
         /// If the canister preallocates a stable memory region, specify the starting page. Required if you want to profile upgrades, or the canister uses stable memory
         #[clap(short, long)]
-        start_page: Option<i32>,
+        start_page: Option<i64>,
         /// The number of pages of the preallocated stable memory
         #[clap(short, long, requires("start_page"))]
         page_limit: Option<i32>,

diff --git a/src/instrumentation.rs b/src/instrumentation.rs
@@ -36,23 +36,23 @@ struct Variables {
 
 pub struct Config {
     pub trace_only_funcs: Vec<String>,
-    pub start_address: Option<i32>,
+    pub start_address: Option<i64>,
     pub page_limit: Option<i32>,
 }
 impl Config {
     pub fn is_preallocated(&self) -> bool {
         self.start_address.is_some()
     }
-    pub fn log_start_address(&self) -> i32 {
-        self.start_address.unwrap_or(0) + METADATA_SIZE
+    pub fn log_start_address(&self) -> i64 {
+        self.start_address.unwrap_or(0) + METADATA_SIZE as i64
     }
-    pub fn metadata_start_address(&self) -> i32 {
+    pub fn metadata_start_address(&self) -> i64 {
         self.start_address.unwrap_or(0)
     }
-    pub fn page_limit(&self) -> i32 {
+    pub fn page_limit(&self) -> i64 {
         self.page_limit
             .map(|x| x - 1)
-            .unwrap_or(DEFAULT_PAGE_LIMIT - 1) // minus 1 because of metadata
+            .unwrap_or(DEFAULT_PAGE_LIMIT - 1) as i64 // minus 1 because of metadata
     }
 }
 
@@ -440,38 +440,39 @@ fn make_dynamic_counter64(
     builder.finish(vec![size], &mut m.funcs)
 }
 fn make_stable_writer(m: &mut Module, vars: &Variables, config: &Config) -> FunctionId {
-    let writer = get_ic_func_id(m, "stable_write");
-    let grow = get_ic_func_id(m, "stable_grow");
+    let writer = get_ic_func_id(m, "stable64_write");
+    let grow = get_ic_func_id(m, "stable64_grow");
     let mut builder = FunctionBuilder::new(
         &mut m.types,
-        &[ValType::I32, ValType::I32, ValType::I32],
+        &[ValType::I64, ValType::I64, ValType::I64],
         &[],
     );
     let start_address = config.log_start_address();
     let size_limit = config.page_limit() * 65536;
     let is_preallocated = config.is_preallocated();
-    let offset = m.locals.add(ValType::I32);
-    let src = m.locals.add(ValType::I32);
-    let size = m.locals.add(ValType::I32);
+    let offset = m.locals.add(ValType::I64);
+    let src = m.locals.add(ValType::I64);
+    let size = m.locals.add(ValType::I64);
     builder
         .func_body()
         .local_get(offset)
         .local_get(size)
-        .binop(BinaryOp::I32Add);
+        .binop(BinaryOp::I64Add);
     if is_preallocated {
-        builder.func_body().i32_const(size_limit);
+        builder.func_body().i64_const(size_limit);
     } else {
         builder
             .func_body()
             .global_get(vars.page_size)
             .i32_const(65536)
             .binop(BinaryOp::I32Mul)
             .i32_const(METADATA_SIZE)
-            .binop(BinaryOp::I32Sub);
+            .binop(BinaryOp::I32Sub)
+            .unop(UnaryOp::I64ExtendUI32);
     }
     builder
         .func_body()
-        .binop(BinaryOp::I32GtS)
+        .binop(BinaryOp::I64GtS)
         .if_else(
             None,
             |then| {
@@ -489,7 +490,7 @@ fn make_stable_writer(m: &mut Module, vars: &Variables, config: &Config) -> Func
                             },
                             |else_| {
                                 else_
-                                    .i32_const(1)
+                                    .i64_const(1)
                                     .call(grow)
                                     .drop()
                                     .global_get(vars.page_size)
@@ -502,9 +503,9 @@ fn make_stable_writer(m: &mut Module, vars: &Variables, config: &Config) -> Func
             },
             |_| {},
         )
-        .i32_const(start_address)
+        .i64_const(start_address)
         .local_get(offset)
-        .binop(BinaryOp::I32Add)
+        .binop(BinaryOp::I64Add)
         .local_get(src)
         .local_get(size)
         .call(writer)
@@ -544,10 +545,11 @@ fn make_printer(m: &mut Module, vars: &Variables, writer: FunctionId) -> Functio
                 .global_get(vars.total_counter)
                 .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
                 .global_get(vars.log_size)
-                .i32_const(LOG_ITEM_SIZE)
-                .binop(BinaryOp::I32Mul)
-                .i32_const(0)
-                .i32_const(LOG_ITEM_SIZE)
+                .unop(UnaryOp::I64ExtendUI32)
+                .i64_const(LOG_ITEM_SIZE as i64)
+                .binop(BinaryOp::I64Mul)
+                .i64_const(0)
+                .i64_const(LOG_ITEM_SIZE as i64)
                 .call(writer)
                 // restore memory
                 .i32_const(0)
@@ -617,7 +619,7 @@ fn inject_init(m: &mut Module, is_init: GlobalId) {
     builder.i32_const(0).global_set(is_init);
 }
 fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
-    let writer = get_ic_func_id(m, "stable_write");
+    let writer = get_ic_func_id(m, "stable64_write");
     let memory = get_memory_id(m);
     let a = m.locals.add(ValType::I64);
     let b = m.locals.add(ValType::I64);
@@ -651,9 +653,9 @@ fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
         .i32_const(20)
         .global_get(vars.is_entry)
         .store(memory, StoreKind::I32 { atomic: false }, MemArg { offset: 0, align: 4 })
-        .i32_const(config.metadata_start_address())
-        .i32_const(0)
-        .i32_const(METADATA_SIZE)
+        .i64_const(config.metadata_start_address())
+        .i64_const(0)
+        .i64_const(METADATA_SIZE as i64)
         .call(writer)
         // restore memory
         .i32_const(0)
@@ -667,7 +669,7 @@ fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
         .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 });
 }
 fn inject_post_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
-    let reader = get_ic_func_id(m, "stable_read");
+    let reader = get_ic_func_id(m, "stable64_read");
     let memory = get_memory_id(m);
     let a = m.locals.add(ValType::I64);
     let b = m.locals.add(ValType::I64);
@@ -686,9 +688,9 @@ fn inject_post_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
         Load { memory, kind: LoadKind::I64 { atomic: false }, arg: MemArg { offset: 0, align: 8 } }.into(),
         LocalSet { local: c }.into(),
         // load from stable memory
-        Const { value: Value::I32(0) }.into(),
-        Const { value: Value::I32(config.metadata_start_address()) }.into(),
-        Const { value: Value::I32(METADATA_SIZE) }.into(),
+        Const { value: Value::I64(0) }.into(),
+        Const { value: Value::I64(config.metadata_start_address()) }.into(),
+        Const { value: Value::I64(METADATA_SIZE as i64) }.into(),
         Call { func: reader }.into(),
         Const { value: Value::I32(0) }.into(),
         Load { memory, kind: LoadKind::I64 { atomic: false }, arg: MemArg { offset: 0, align: 8 } }.into(),
@@ -725,7 +727,7 @@ fn make_stable_getter(m: &mut Module, vars: &Variables, leb: FunctionId, config:
     let reply_data = get_ic_func_id(m, "msg_reply_data_append");
     let reply = get_ic_func_id(m, "msg_reply");
     let trap = get_ic_func_id(m, "trap");
-    let reader = get_ic_func_id(m, "stable_read");
+    let reader = get_ic_func_id(m, "stable64_read");
     let idx = m.locals.add(ValType::I32);
     let len = m.locals.add(ValType::I32);
     let next_idx = m.locals.add(ValType::I32);
@@ -808,15 +810,15 @@ fn make_stable_getter(m: &mut Module, vars: &Variables, leb: FunctionId, config:
         .i32_const(5)
         .call(reply_data)
         // read stable logs
-        .i32_const(0)
-        .i32_const(config.log_start_address())
+        .i64_const(0)
+        .i64_const(config.log_start_address())
         .local_get(idx)
-        .i32_const(LOG_ITEM_SIZE)
-        .binop(BinaryOp::I32Mul)
-        .binop(BinaryOp::I32Add)
+        .i64_const(LOG_ITEM_SIZE as i64)
+        .binop(BinaryOp::I64Mul)
+        .binop(BinaryOp::I64Add)
         .local_get(len)
-        .i32_const(LOG_ITEM_SIZE)
-        .binop(BinaryOp::I32Mul)
+        .i64_const(LOG_ITEM_SIZE as i64)
+        .binop(BinaryOp::I64Mul)
         .call(reader)
         .i32_const(0)
         .local_get(len)

diff --git a/tests/ok/motoko-gc-instrument.wasm b/tests/ok/motoko-gc-instrument.wasm
diff --git a/tests/ok/motoko-instrument.wasm b/tests/ok/motoko-instrument.wasm
diff --git a/tests/ok/motoko-region-instrument.wasm b/tests/ok/motoko-region-instrument.wasm
diff --git a/tests/ok/rust-instrument.wasm b/tests/ok/rust-instrument.wasm
diff --git a/tests/ok/rust-region-instrument.wasm b/tests/ok/rust-region-instrument.wasm
diff --git a/tests/ok/wat-instrument.wasm b/tests/ok/wat-instrument.wasm