Add test case and token finder to address 12790

Author: TopGunSnake

crates/ide-assists/src/handlers/extract_module.rs

@@ -19,7 +19,7 @@ use syntax::{
         make, HasName, HasVisibility,
     },
     match_ast, ted, AstNode, SourceFile,
-    SyntaxKind::WHITESPACE,
+    SyntaxKind::{self, WHITESPACE},
     SyntaxNode, TextRange,
 };
 
@@ -380,7 +380,24 @@ impl Module {
         }
 
         for (vis, syntax) in replacements {
-            add_change_vis(vis, syntax.first_child_or_token());
+            let item = syntax.children_with_tokens().find(|node_or_token| {
+                match node_or_token.kind() {
+                    // We're looking for the start of functions, impls, structs, traits, and other documentable/attribute
+                    // macroable items that would have pub(crate) in front of it
+                    SyntaxKind::FN_KW
+                    | SyntaxKind::IMPL_KW
+                    | SyntaxKind::STRUCT_KW
+                    | SyntaxKind::TRAIT_KW
+                    | SyntaxKind::TYPE_KW
+                    | SyntaxKind::MOD_KW => true,
+                    // If we didn't find a keyword, we want to cover the record fields
+                    SyntaxKind::NAME => true,
+                    // Otherwise, the token shouldn't have pub(crate) before it
+                    _ => false,
+                }
+            });
+
+            add_change_vis(vis, item);
         }
     }
 
@@ -1581,4 +1598,129 @@ mod modname {
         ",
         )
     }
+
+    #[test]
+    fn test_issue_12790() {
+        check_assist(
+            extract_module,
+            r"
+            $0/// A documented function
+            fn documented_fn() {}
+            
+            // A commented function with a #[] attribute macro
+            #[cfg(test)]
+            fn attribute_fn() {}
+            
+            // A normally commented function
+            fn normal_fn() {}
+            
+            /// A documented Struct
+            struct DocumentedStruct {
+                // Normal field
+                x: i32,
+            
+                /// Documented field
+                y: i32,
+            
+                // Macroed field
+                #[cfg(test)]
+                z: i32,
+            }
+            
+            // A macroed Struct
+            #[cfg(test)]
+            struct MacroedStruct {
+                // Normal field
+                x: i32,
+            
+                /// Documented field
+                y: i32,
+            
+                // Macroed field
+                #[cfg(test)]
+                z: i32,
+            }
+            
+            // A normal Struct
+            struct NormalStruct {
+                // Normal field
+                x: i32,
+            
+                /// Documented field
+                y: i32,
+            
+                // Macroed field
+                #[cfg(test)]
+                z: i32,
+            }
+
+            /// A documented type
+            type DocumentedType = i32;
+
+            // A macroed type
+            #[cfg(test)]
+            type MacroedType = i32;$0
+        ",
+            r"
+            mod modname {
+                /// A documented function
+                pub(crate) fn documented_fn() {}
+            
+                // A commented function with a #[] attribute macro
+                #[cfg(test)]
+                pub(crate) fn attribute_fn() {}
+            
+                // A normally commented function
+                pub(crate) fn normal_fn() {}
+            
+                /// A documented Struct
+                pub(crate) struct DocumentedStruct {
+                    // Normal field
+                    pub(crate) x: i32,
+            
+                    /// Documented field
+                    pub(crate) y: i32,
+            
+                    // Macroed field
+                    #[cfg(test)]
+                    pub(crate) z: i32,
+                }
+            
+                // A macroed Struct
+                #[cfg(test)]
+                pub(crate) struct MacroedStruct {
+                    // Normal field
+                    pub(crate) x: i32,
+            
+                    /// Documented field
+                    pub(crate) y: i32,
+            
+                    // Macroed field
+                    #[cfg(test)]
+                    pub(crate) z: i32,
+                }
+            
+                // A normal Struct
+                pub(crate) struct NormalStruct {
+                    // Normal field
+                    pub(crate) x: i32,
+            
+                    /// Documented field
+                    pub(crate) y: i32,
+            
+                    // Macroed field
+                    #[cfg(test)]
+                    pub(crate) z: i32,
+                }
+
+                /// A documented type
+                pub(crate) type DocumentedType = i32;
+
+                // A macroed type
+                #[cfg(test)]
+                pub(crate) type MacroedType = i32;
+            }
+        ",
+        )
+    }
 }