allow Keys::<T>::remove() in replace_neuron()

Author: eagr

pallets/subtensor/src/subnets/uids.rs

@@ -29,6 +29,7 @@ impl<T: Config> Pallet<T> {
         // 2. Remove previous set memberships.
         Uids::<T>::remove(netuid, old_hotkey.clone());
         IsNetworkMember::<T>::remove(old_hotkey.clone(), netuid);
+        #[allow(unknown_lints)]
         Keys::<T>::remove(netuid, uid_to_replace);
 
         // 2a. Check if the uid is registered in any other subnetworks.

support/linting/src/forbid_keys_remove.rs

@@ -1,7 +1,7 @@
 use super::*;
 use syn::{
     punctuated::Punctuated, spanned::Spanned, token::Comma, visit::Visit, Expr, ExprCall, ExprPath,
-    File,
+    File, Path,
 };
 
 pub struct ForbidKeysRemoveCall;
@@ -26,25 +26,30 @@ struct KeysRemoveVisitor {
 
 impl<'ast> Visit<'ast> for KeysRemoveVisitor {
     fn visit_expr_call(&mut self, node: &'ast syn::ExprCall) {
-        let ExprCall { func, args, .. } = node;
-        if is_keys_remove_call(func, args) {
+        let ExprCall {
+            func, args, attrs, ..
+        } = node;
+
+        if is_keys_remove_call(func, args) && !is_allowed(attrs) {
             let msg = "Keys::<T>::remove()` is banned to prevent accidentally breaking \
-                the neuron sequence. If you need to replace neuron, try `SubtensorModule::replace_neuron()`";
+                the neuron sequence. If you need to replace neurons, try `SubtensorModule::replace_neuron()`";
             self.errors.push(syn::Error::new(node.func.span(), msg));
         }
     }
 }
 
 fn is_keys_remove_call(func: &Expr, args: &Punctuated<Expr, Comma>) -> bool {
-    let Expr::Path(ExprPath { path, .. }) = func else {
+    let Expr::Path(ExprPath {
+        path: Path { segments: func, .. },
+        ..
+    }) = func
+    else {
         return false;
     };
-    let func = &path.segments;
-    if func.len() != 2 || args.len() != 2 {
-        return false;
-    }
 
-    func[0].ident == "Keys"
+    func.len() == 2
+        && args.len() == 2
+        && func[0].ident == "Keys"
         && !func[0].arguments.is_none()
         && func[1].ident == "remove"
         && func[1].arguments.is_none()
@@ -91,4 +96,23 @@ mod tests {
         let input = r#"ChildKeys::<T>::remove(netuid, uid_to_replace)"#;
         assert!(lint(input).is_ok());
     }
+
+    #[test]
+    fn test_keys_remove_allowed() {
+        let input = r#"
+        #[allow(unknown_lints)]
+        Keys::<T>::remove(netuid, uid_to_replace)
+        "#;
+        assert!(lint(input).is_ok());
+        let input = r#"
+        #[allow(unknown_lints)]
+        Keys::<U>::remove(netuid, uid_to_replace)
+        "#;
+        assert!(lint(input).is_ok());
+        let input = r#"
+        #[allow(unknown_lints)]
+        Keys::<U>::remove(1, "2".parse().unwrap(),)
+        "#;
+        assert!(lint(input).is_ok());
+    }
 }

support/linting/src/lint.rs

@@ -1,4 +1,5 @@
-use syn::File;
+use proc_macro2::TokenTree;
+use syn::{Attribute, File, Meta, MetaList, Path};
 
 pub type Result = core::result::Result<(), Vec<syn::Error>>;
 
@@ -11,3 +12,29 @@ pub trait Lint: Send + Sync {
     /// Lints the given Rust source file, returning a compile error if any issues are found.
     fn lint(source: &File) -> Result;
 }
+
+pub fn is_allowed(attibutes: &[Attribute]) -> bool {
+    attibutes.iter().any(|attribute| {
+        let Attribute {
+            meta:
+                Meta::List(MetaList {
+                    path: Path { segments: attr, .. },
+                    tokens: attr_args,
+                    ..
+                }),
+            ..
+        } = attribute
+        else {
+            return false;
+        };
+
+        attr.len() == 1
+            && attr[0].ident == "allow"
+            && attr_args.clone().into_iter().any(|arg| {
+                let TokenTree::Ident(ref id) = arg else {
+                    return false;
+                };
+                id == "unknown_lints"
+            })
+    })
+}