From bfd54a115bee689321063c87b23d12db56516617 Mon Sep 17 00:00:00 2001
From: Yan Chen <yan.chen@dfinity.org>
Date: Sat, 7 Oct 2023 11:54:42 -0700
Subject: [PATCH] fix

---
 src/bin/main.rs        |  2 +-
 src/instrumentation.rs | 31 ++++++++++++++++++++++++++++---
 src/utils.rs           |  4 ++--
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/bin/main.rs b/src/bin/main.rs
index 6147d47..189b99d 100644
--- a/src/bin/main.rs
+++ b/src/bin/main.rs
@@ -76,7 +76,7 @@ enum SubCommand {
         /// The number of pages of the preallocated stable memory
         #[clap(short, long, requires("start_page"))]
         page_limit: Option<i32>,
-        /// Use the new metering cost, default to false
+        /// Use the new metering cost, default to false. This flag will eventually be removed and set to true after the mainnet has adapted the new metering.
         #[clap(short, long)]
         use_new_metering: bool,
     },
diff --git a/src/instrumentation.rs b/src/instrumentation.rs
index bd5a2d8..302e13a 100644
--- a/src/instrumentation.rs
+++ b/src/instrumentation.rs
@@ -195,7 +195,10 @@ fn inject_metering(
                 }
                 Instr::Call(Call { func }) => {
                     curr.cost += instr_cost(instr, use_new_metering);
-                    match func_cost.get_cost(*func) {
+                    match func_cost
+                        .get_cost(*func)
+                        .unwrap_or((instr_cost(instr, use_new_metering), InjectionKind::Static))
+                    {
                         (cost, InjectionKind::Static) => curr.cost += cost,
                         (cost, kind @ InjectionKind::Dynamic)
                         | (cost, kind @ InjectionKind::Dynamic64) => {
@@ -616,10 +619,22 @@ fn inject_init(m: &mut Module, is_init: GlobalId) {
 fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
     let writer = get_ic_func_id(m, "stable_write");
     let memory = get_memory_id(m);
+    /*let a = m.locals.add(ValType::I64);
+    let b = m.locals.add(ValType::I64);
+    let c = m.locals.add(ValType::I64);*/
     let mut builder = get_or_create_export_func(m, "canister_pre_upgrade");
     #[rustfmt::skip]
     builder
-        // no need to backup memory, since this is the end of the pre-upgrade.
+        // backup memory. This is not strictly needed, since it's at the end of pre-upgrade.
+        /*.i32_const(0)
+        .load(memory, LoadKind::I64 { atomic: false }, MemArg { offset: 0, align: 8})
+        .local_set(a)
+        .i32_const(8)
+        .load(memory, LoadKind::I64 { atomic: false }, MemArg { offset: 0, align: 8})
+        .local_set(b)
+        .i32_const(16)
+        .load(memory, LoadKind::I64 { atomic: false }, MemArg { offset: 0, align: 8})
+        .local_set(c)*/
         // persist globals
         .i32_const(0)
         .global_get(vars.total_counter)
@@ -639,7 +654,17 @@ fn inject_pre_upgrade(m: &mut Module, vars: &Variables, config: &Config) {
         .i32_const(config.metadata_start_address())
         .i32_const(0)
         .i32_const(METADATA_SIZE)
-        .call(writer);
+        .call(writer)
+        // restore memory
+        /*.i32_const(0)
+        .local_get(a)
+        .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
+        .i32_const(8)
+        .local_get(b)
+        .store(memory, StoreKind::I64 { atomic: false }, MemArg { offset: 0, align: 8 })
+        .i32_const(16)
+        .local_get(c)
+        .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");
diff --git a/src/utils.rs b/src/utils.rs
index 8098f63..eb7ae94 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -114,8 +114,8 @@ impl FunctionCost {
         }
         Self(res)
     }
-    pub fn get_cost(&self, id: FunctionId) -> (i64, InjectionKind) {
-        *self.0.get(&id).unwrap_or(&(1, InjectionKind::Static))
+    pub fn get_cost(&self, id: FunctionId) -> Option<(i64, InjectionKind)> {
+        self.0.get(&id).copied()
     }
 }
 pub(crate) fn instr_cost(i: &ir::Instr, use_new_metering: bool) -> i64 {