From 3792fd9b53f29dc7b47f7bf7a53cc13f72e5fffd Mon Sep 17 00:00:00 2001
From: Adam Spofford <adam.spofford@dfinity.org>
Date: Fri, 1 Nov 2024 09:13:11 -0700
Subject: [PATCH] Implement low wasm memory hook

---
 src/ic-cdk-macros/src/export.rs        | 24 +++++++++++++++++-------
 src/ic-cdk-macros/src/lib.rs           | 10 ++++++++++
 src/ic-cdk/src/macros.rs               | 21 +++++++++++++++++++++
 src/ic-cdk/tests/pass/blank_methods.rs |  7 ++++++-
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/src/ic-cdk-macros/src/export.rs b/src/ic-cdk-macros/src/export.rs
index 6d087d398..320c94d77 100644
--- a/src/ic-cdk-macros/src/export.rs
+++ b/src/ic-cdk-macros/src/export.rs
@@ -33,18 +33,20 @@ enum MethodType {
     Query,
     Heartbeat,
     InspectMessage,
+    OnLowWasmMemory,
 }
 
 impl MethodType {
     pub fn is_lifecycle(&self) -> bool {
-        matches!(
-            self,
+        match self {
             MethodType::Init
-                | MethodType::PreUpgrade
-                | MethodType::PostUpgrade
-                | MethodType::Heartbeat
-                | MethodType::InspectMessage
-        )
+            | MethodType::PreUpgrade
+            | MethodType::PostUpgrade
+            | MethodType::Heartbeat
+            | MethodType::InspectMessage
+            | MethodType::OnLowWasmMemory => true,
+            MethodType::Update | MethodType::Query => false,
+        }
     }
 }
 
@@ -58,6 +60,7 @@ impl std::fmt::Display for MethodType {
             MethodType::Update => f.write_str("update"),
             MethodType::Heartbeat => f.write_str("heartbeat"),
             MethodType::InspectMessage => f.write_str("inspect_message"),
+            MethodType::OnLowWasmMemory => f.write_str("on_low_wasm_memory"),
         }
     }
 }
@@ -309,6 +312,13 @@ pub(crate) fn ic_inspect_message(
     dfn_macro(MethodType::InspectMessage, attr, item)
 }
 
+pub(crate) fn ic_on_low_wasm_memory(
+    attr: TokenStream,
+    item: TokenStream,
+) -> Result<TokenStream, Error> {
+    dfn_macro(MethodType::OnLowWasmMemory, attr, item)
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/src/ic-cdk-macros/src/lib.rs b/src/ic-cdk-macros/src/lib.rs
index d9e75e690..00a0d566d 100644
--- a/src/ic-cdk-macros/src/lib.rs
+++ b/src/ic-cdk-macros/src/lib.rs
@@ -91,3 +91,13 @@ pub fn heartbeat(attr: TokenStream, item: TokenStream) -> TokenStream {
 pub fn inspect_message(attr: TokenStream, item: TokenStream) -> TokenStream {
     handle_debug_and_errors(export::ic_inspect_message, "ic_inspect_message", attr, item)
 }
+
+#[proc_macro_attribute]
+pub fn on_low_wasm_memory(attr: TokenStream, item: TokenStream) -> TokenStream {
+    handle_debug_and_errors(
+        export::ic_on_low_wasm_memory,
+        "ic_on_low_wasm_memory",
+        attr,
+        item,
+    )
+}
diff --git a/src/ic-cdk/src/macros.rs b/src/ic-cdk/src/macros.rs
index 8c3162875..e2f88c5b9 100644
--- a/src/ic-cdk/src/macros.rs
+++ b/src/ic-cdk/src/macros.rs
@@ -313,3 +313,24 @@ pub use ic_cdk_macros::heartbeat;
 /// }
 /// ```
 pub use ic_cdk_macros::inspect_message;
+
+/// Register the `canister_on_low_wasm_memory` entry point of a canister.
+///
+/// This attribute macro will export the function `canister_on_low_wasm_memory`
+/// in the canister module.
+///
+/// The function under this attribute must have no return value.
+///
+/// Each canister can only have one `canister_on_low_wasm_memory` entry point.
+///
+/// # Example
+///
+/// ```rust
+/// # use ic_cdk::on_low_wasm_memory;
+/// #[on_low_wasm_memory]
+/// fn low_memory_handler() {
+///     // ...
+/// # unimplemented!()
+/// }
+/// ```
+pub use ic_cdk_macros::on_low_wasm_memory;
diff --git a/src/ic-cdk/tests/pass/blank_methods.rs b/src/ic-cdk/tests/pass/blank_methods.rs
index bd3922901..32e4440e6 100644
--- a/src/ic-cdk/tests/pass/blank_methods.rs
+++ b/src/ic-cdk/tests/pass/blank_methods.rs
@@ -1,4 +1,6 @@
-use ic_cdk::{heartbeat, init, inspect_message, post_upgrade, pre_upgrade, query, update};
+use ic_cdk::{
+    heartbeat, init, inspect_message, on_low_wasm_memory, post_upgrade, pre_upgrade, query, update,
+};
 
 #[init]
 fn init() {}
@@ -30,4 +32,7 @@ fn heartbeat() {}
 #[inspect_message]
 fn inspect_message() {}
 
+#[on_low_wasm_memory]
+fn on_low_wasm_memory() {}
+
 fn main() {}