Merge pull request #927 from eagr/forbid-keys-remove

Lint `ForbidKeysRemoveCall`

build.rs

@@ -60,6 +60,7 @@ fn main() {
         };
 
         track_lint(ForbidAsPrimitiveConversion::lint(&parsed_file));
+        track_lint(ForbidKeysRemoveCall::lint(&parsed_file));
         track_lint(RequireFreezeStruct::lint(&parsed_file));
         track_lint(RequireExplicitPalletIndex::lint(&parsed_file));
     });

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_as_primitive.rs

@@ -45,10 +45,11 @@ fn is_as_primitive(ident: &Ident) -> bool {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use quote::quote;
 
-    fn lint(input: &str) -> Result {
-        let expr: ExprMethodCall = syn::parse_str(input).expect("should only use on a method call");
+    fn lint(input: proc_macro2::TokenStream) -> Result {
         let mut visitor = AsPrimitiveVisitor::default();
+        let expr: ExprMethodCall = syn::parse2(input).expect("should be a valid method call");
         visitor.visit_expr_method_call(&expr);
         if !visitor.errors.is_empty() {
             return Err(visitor.errors);
@@ -58,21 +59,21 @@ mod tests {
 
     #[test]
     fn test_as_primitives() {
-        let input = r#"x.as_u32()"#;
+        let input = quote! {x.as_u32() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_u64()"#;
+        let input = quote! {x.as_u64() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_u128()"#;
+        let input = quote! {x.as_u128() };
         assert!(lint(input).is_err());
-        let input = r#"x.as_usize()"#;
+        let input = quote! {x.as_usize() };
         assert!(lint(input).is_err());
     }
 
     #[test]
     fn test_non_as_primitives() {
-        let input = r#"x.as_ref()"#;
+        let input = quote! {x.as_ref() };
         assert!(lint(input).is_ok());
-        let input = r#"x.as_slice()"#;
+        let input = quote! {x.as_slice() };
         assert!(lint(input).is_ok());
     }
 }

support/linting/src/forbid_keys_remove.rs

@@ -0,0 +1,119 @@
+use super::*;
+use syn::{
+    punctuated::Punctuated, spanned::Spanned, token::Comma, visit::Visit, Expr, ExprCall, ExprPath,
+    File, Path,
+};
+
+pub struct ForbidKeysRemoveCall;
+
+impl Lint for ForbidKeysRemoveCall {
+    fn lint(source: &File) -> Result {
+        let mut visitor = KeysRemoveVisitor::default();
+        visitor.visit_file(source);
+
+        if visitor.errors.is_empty() {
+            Ok(())
+        } else {
+            Err(visitor.errors)
+        }
+    }
+}
+
+#[derive(Default)]
+struct KeysRemoveVisitor {
+    errors: Vec<syn::Error>,
+}
+
+impl<'ast> Visit<'ast> for KeysRemoveVisitor {
+    fn visit_expr_call(&mut self, node: &'ast syn::ExprCall) {
+        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 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: Path { segments: func, .. },
+        ..
+    }) = func
+    else {
+        return false;
+    };
+
+    func.len() == 2
+        && args.len() == 2
+        && func[0].ident == "Keys"
+        && !func[0].arguments.is_none()
+        && func[1].ident == "remove"
+        && func[1].arguments.is_none()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use quote::quote;
+
+    fn lint(input: proc_macro2::TokenStream) -> Result {
+        let mut visitor = KeysRemoveVisitor::default();
+        let expr: syn::ExprCall = syn::parse2(input).expect("should be a valid function call");
+        visitor.visit_expr_call(&expr);
+
+        if visitor.errors.is_empty() {
+            Ok(())
+        } else {
+            Err(visitor.errors)
+        }
+    }
+
+    #[test]
+    fn test_keys_remove_forbidden() {
+        let input = quote! { Keys::<T>::remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_err());
+        let input = quote! { Keys::<U>::remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_err());
+        let input = quote! { Keys::<U>::remove(1, "2".parse().unwrap(),) };
+        assert!(lint(input).is_err());
+    }
+
+    #[test]
+    fn test_non_keys_remove_not_forbidden() {
+        let input = quote! { remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_ok());
+        let input = quote! { Keys::remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_ok());
+        let input = quote! { Keys::<T>::remove::<U>(netuid, uid_to_replace) };
+        assert!(lint(input).is_ok());
+        let input = quote! { Keys::<T>::remove(netuid, uid_to_replace, third_wheel) };
+        assert!(lint(input).is_ok());
+        let input = quote! { ParentKeys::remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_ok());
+        let input = quote! { ChildKeys::<T>::remove(netuid, uid_to_replace) };
+        assert!(lint(input).is_ok());
+    }
+
+    #[test]
+    fn test_keys_remove_allowed() {
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<T>::remove(netuid, uid_to_replace)
+        };
+        assert!(lint(input).is_ok());
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<U>::remove(netuid, uid_to_replace)
+        };
+        assert!(lint(input).is_ok());
+        let input = quote! {
+            #[allow(unknown_lints)]
+            Keys::<U>::remove(1, "2".parse().unwrap(),)
+        };
+        assert!(lint(input).is_ok());
+    }
+}

support/linting/src/lib.rs

@@ -2,9 +2,11 @@ pub mod lint;
 pub use lint::*;
 
 mod forbid_as_primitive;
+mod forbid_keys_remove;
 mod pallet_index;
 mod require_freeze_struct;
 
 pub use forbid_as_primitive::ForbidAsPrimitiveConversion;
+pub use forbid_keys_remove::ForbidKeysRemoveCall;
 pub use pallet_index::RequireExplicitPalletIndex;
 pub use require_freeze_struct::RequireFreezeStruct;

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"
+            })
+    })
+}