From 96aca0c26dec13f1e47d1b9077459c0e3fd999e8 Mon Sep 17 00:00:00 2001
From: Jeeyong Um <conr2d@proton.me>
Date: Mon, 28 Oct 2024 12:10:15 +0900
Subject: [PATCH] feat: Add clearing addresses on killed account (#87)

---
 frame/babel/src/extensions/unify_account.rs | 22 ++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/frame/babel/src/extensions/unify_account.rs b/frame/babel/src/extensions/unify_account.rs
index 46a5ee69..d6d0d6df 100644
--- a/frame/babel/src/extensions/unify_account.rs
+++ b/frame/babel/src/extensions/unify_account.rs
@@ -18,7 +18,10 @@
 
 use crate::traits::AccountIdProvider;
 use core::marker::PhantomData;
-use frame_support::traits::tokens::{fungible, Fortitude, Preservation};
+use frame_support::traits::{
+	tokens::{fungible, Fortitude, Preservation},
+	OnKilledAccount as OnKilledAccountT,
+};
 use np_babel::VarAddress;
 use pallet_multimap::traits::UniqueMultimap;
 use parity_scale_codec::{Decode, Encode, MaxEncodedLen};
@@ -182,6 +185,14 @@ where
 	}
 }
 
+pub struct OnKilledAccount<T: Config>(PhantomData<T>);
+
+impl<T: Config> OnKilledAccountT<T::AccountId> for OnKilledAccount<T> {
+	fn on_killed_account(who: &T::AccountId) {
+		T::AddressMap::remove_all(who);
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -222,4 +233,13 @@ mod tests {
 			assert_eq!(<MockConfig as Config>::AddressMap::find_key(nostr), Some(who));
 		}
 	}
+
+	#[test]
+	fn on_killed_account_works() {
+		let who = AccountId::from(dev_public());
+		let _ = UnifyAccount::<MockConfig>::unify_ecdsa(&who);
+		assert!(!<MockConfig as Config>::AddressMap::get(&who).is_empty());
+		OnKilledAccount::<MockConfig>::on_killed_account(&who);
+		assert!(<MockConfig as Config>::AddressMap::get(&who).is_empty());
+	}
 }