From c9439739a7b800a2137dca555efee35a8cf57971 Mon Sep 17 00:00:00 2001
From: Yan Chen <48968912+chenyan-dfinity@users.noreply.github.com>
Date: Tue, 27 Aug 2024 16:09:36 -0700
Subject: [PATCH] fix memory id in limit_resource (#67)

* fix memory id in limit_resource

* bump version
---
 Cargo.lock            |  2 +-
 Cargo.toml            |  2 +-
 src/limit_resource.rs | 13 +++++++------
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 47f50b2..c4ef050 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -487,7 +487,7 @@ checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70"
 
 [[package]]
 name = "ic-wasm"
-version = "0.8.2"
+version = "0.8.3"
 dependencies = [
  "anyhow",
  "assert_cmd",
diff --git a/Cargo.toml b/Cargo.toml
index 7b78ec9..6357069 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "ic-wasm"
-version = "0.8.2"
+version = "0.8.3"
 authors = ["DFINITY Stiftung"]
 edition = "2021"
 description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer"
diff --git a/src/limit_resource.rs b/src/limit_resource.rs
index f3092e7..1d34a87 100644
--- a/src/limit_resource.rs
+++ b/src/limit_resource.rs
@@ -171,13 +171,14 @@ pub fn limit_resource(m: &mut Module, config: &Config) {
 }
 
 fn limit_heap_memory(m: &mut Module, limit: u32) {
-    let memory_id = m.get_memory_id().expect("unable to get memory id");
-    let memory = m.memories.get_mut(memory_id);
-    let limit = limit as u64;
-    if memory.initial > limit {
-        memory.initial = limit
+    if let Ok(memory_id) = m.get_memory_id() {
+        let memory = m.memories.get_mut(memory_id);
+        let limit = limit as u64;
+        if memory.initial > limit {
+            memory.initial = limit
+        }
+        memory.maximum = Some(limit);
     }
-    memory.maximum = Some(limit);
 }
 
 fn make_cycles_add128(m: &mut Module) -> FunctionId {