From 9a65c417f30f2df7000b32feb7cd688a46b2954f Mon Sep 17 00:00:00 2001
From: Adam Spofford <93943719+adamspofford-dfinity@users.noreply.github.com>
Date: Mon, 4 Nov 2024 13:06:23 -0800
Subject: [PATCH] feat: Implement low wasm memory hook (#528)

---
 src/ic-cdk-macros/src/export.rs        | 24 +++++++++++++++++-------
 src/ic-cdk-macros/src/lib.rs           | 10 ++++++++++
 src/ic-cdk/CHANGELOG.md                |  1 +
 src/ic-cdk/README.md                   |  1 +
 src/ic-cdk/src/macros.rs               | 21 +++++++++++++++++++++
 src/ic-cdk/tests/pass/blank_methods.rs |  7 ++++++-
 6 files changed, 56 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/CHANGELOG.md b/src/ic-cdk/CHANGELOG.md
index 2c58efa9d..55d130c1d 100644
--- a/src/ic-cdk/CHANGELOG.md
+++ b/src/ic-cdk/CHANGELOG.md
@@ -12,6 +12,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Attribute `#[on_low_wasm_memory]` for low-memory hook. (#528)
 - Support Threshold Schnorr signing management canister API. (#518)
 
 ## [0.16.0] - 2024-08-27
diff --git a/src/ic-cdk/README.md b/src/ic-cdk/README.md
index aaa8b802d..80d7ffcf2 100644
--- a/src/ic-cdk/README.md
+++ b/src/ic-cdk/README.md
@@ -58,6 +58,7 @@ These macros are directly related to the [Internet Computer Specification](https
 * [`post_upgrade`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.post_upgrade.html)
 * [`inspect_message`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.inspect_message.html)
 * [`heartbeat`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.heartbeat.html)
+* [`on_low_wasm_memory`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.on_low_wasm_memory.html)
 * [`update`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.update.html)
 * [`query`](https://docs.rs/ic-cdk/latest/ic_cdk/attr.query.html)
 
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() {}