chore: Updates to canister-snapshots example for beryllium demo

rust/canister-snapshots/src/chat/src/lib.rs

@@ -3,12 +3,20 @@ use std::{cell::RefCell, collections::HashSet};
 
 thread_local! {
     static CHAT: RefCell<Vec<String>>  = Default::default();
+    static LENGTH: RefCell<u64> = Default::default();
 }
 
 /// Appends a new message to the chat database.
 #[ic_cdk_macros::update]
 fn append(message: String) {
+    // Ensure that the length of the chat is consistent with the actual chat
+    // contents.
+    assert_eq!(
+        LENGTH.with_borrow(|len| *len),
+        CHAT.with_borrow(|chat| chat.len() as u64)
+    );
     CHAT.with_borrow_mut(|chat| chat.push(message));
+    LENGTH.with_borrow_mut(|len| *len += 1);
 }
 
 /// Dumps all the chat messages.
@@ -36,6 +44,7 @@ fn remove_spam() -> u64 {
     for message in chat {
         if message.split(" ").any(|word| spam_keywords.contains(word)) {
             spam += 1;
+            ic_cdk::println!("Found spam message: {message}");
             new_chat.push("(removed spam message)".into());
         } else {
             new_chat.push(message);
@@ -47,5 +56,22 @@ fn remove_spam() -> u64 {
         ic_cdk::println!("Removed {spam} messages, updating the chat...");
         CHAT.set(new_chat);
     }
+    ic_cdk::println!("Filtered chat: {:?}", CHAT.with_borrow(|chat| chat.clone()));
     spam
 }
+
+#[ic_cdk_macros::pre_upgrade]
+fn pre_upgrade() {
+    let chat = CHAT.with_borrow(|chat| chat.clone());
+    let length = LENGTH.with_borrow(|len| *len);
+    ic_cdk::storage::stable_save((chat, length)).expect("failed to save stable state");
+}
+
+#[ic_cdk_macros::post_upgrade]
+fn post_upgrade() {
+    let (chat, length): (Vec<String>, u64) =
+        ic_cdk::storage::stable_restore().expect("failed to restore stable state");
+
+    CHAT.set(chat);
+    LENGTH.with_borrow_mut(|len| *len = length);
+}