simplify strip attrs

5 files changed, 35 insertions, 29 deletions

diff --git a/crates/ra_assists/src/assists/add_missing_impl_members.rs b/crates/ra_assists/src/assists/add_missing_impl_members.rs
index 682455bce..3fce4a5b7 100644
--- a/crates/ra_assists/src/assists/add_missing_impl_members.rs
+++ b/crates/ra_assists/src/assists/add_missing_impl_members.rs
@@ -1,6 +1,6 @@
 use hir::{db::HirDatabase, HasSource};
 use ra_syntax::{
-    ast::{self, make, AstNode, NameOwner},
+    ast::{self, edit, make, AstNode, NameOwner},
     SmolStr,
 };
 
@@ -76,8 +76,8 @@ fn add_missing_impl_members_inner(
     ctx.add_action(AssistId(assist_id), label, |edit| {
         let n_existing_items = impl_item_list.impl_items().count();
         let items = missing_items.into_iter().map(|it| match it {
-            ast::ImplItem::FnDef(def) => strip_docstring(add_body(def).into()),
-            _ => strip_docstring(it),
+            ast::ImplItem::FnDef(def) => edit::strip_attrs_and_docs(add_body(def).into()),
+            _ => edit::strip_attrs_and_docs(it),
         });
         let mut ast_editor = AstEditor::new(impl_item_list);
 
@@ -93,12 +93,6 @@ fn add_missing_impl_members_inner(
     ctx.build()
 }
 
-fn strip_docstring(item: ast::ImplItem) -> ast::ImplItem {
-    let mut ast_editor = AstEditor::new(item);
-    ast_editor.strip_attrs_and_docs();
-    ast_editor.ast().to_owned()
-}
-
 fn add_body(fn_def: ast::FnDef) -> ast::FnDef {
     if fn_def.body().is_none() {
         fn_def.with_body(make::block_from_expr(make::expr_unimplemented()))
diff --git a/crates/ra_assists/src/ast_editor.rs b/crates/ra_assists/src/ast_editor.rs
index 72c8c478a..60b8923e1 100644
--- a/crates/ra_assists/src/ast_editor.rs
+++ b/crates/ra_assists/src/ast_editor.rs
@@ -212,23 +212,6 @@ impl AstEditor<ast::ItemList> {
     }
 }
 
-impl AstEditor<ast::ImplItem> {
-    pub fn strip_attrs_and_docs(&mut self) {
-        while let Some(start) = self
-            .ast()
-            .syntax()
-            .children_with_tokens()
-            .find(|it| it.kind() == ATTR || it.kind() == COMMENT)
-        {
-            let end = match &start.next_sibling_or_token() {
-                Some(el) if el.kind() == WHITESPACE => el.clone(),
-                Some(_) | None => start.clone(),
-            };
-            self.ast = self.replace_children(RangeInclusive::new(start, end), iter::empty());
-        }
-    }
-}
-
 impl AstEditor<ast::TypeParam> {
     pub fn remove_bounds(&mut self) -> &mut Self {
         let colon = match self.ast.colon_token() {
diff --git a/crates/ra_syntax/src/ast.rs b/crates/ra_syntax/src/ast.rs
index fdffd8cb1..1b2ce921a 100644
--- a/crates/ra_syntax/src/ast.rs
+++ b/crates/ra_syntax/src/ast.rs
@@ -5,7 +5,7 @@ mod traits;
 mod tokens;
 mod extensions;
 mod expr_extensions;
-mod edit;
+pub mod edit;
 pub mod make;
 
 use std::marker::PhantomData;
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index c65899812..7013cc9b5 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -1,14 +1,16 @@
 //! This module contains functions for editing syntax trees. As the trees are
 //! immutable, all function here return a fresh copy of the tree, instead of
 //! doing an in-place modification.
+use std::{iter, ops::RangeInclusive};
 
 use arrayvec::ArrayVec;
-use std::ops::RangeInclusive;
 
 use crate::{
     algo,
     ast::{self, make, AstNode},
     InsertPosition, SyntaxElement,
+    SyntaxKind::{ATTR, COMMENT, WHITESPACE},
+    SyntaxNode,
 };
 
 impl ast::FnDef {
@@ -31,6 +33,23 @@ impl ast::FnDef {
     }
 }
 
+pub fn strip_attrs_and_docs<N: ast::AttrsOwner>(node: N) -> N {
+    N::cast(strip_attrs_and_docs_inner(node.syntax().clone())).unwrap()
+}
+
+fn strip_attrs_and_docs_inner(mut node: SyntaxNode) -> SyntaxNode {
+    while let Some(start) =
+        node.children_with_tokens().find(|it| it.kind() == ATTR || it.kind() == COMMENT)
+    {
+        let end = match &start.next_sibling_or_token() {
+            Some(el) if el.kind() == WHITESPACE => el.clone(),
+            Some(_) | None => start.clone(),
+        };
+        node = algo::replace_children(&node, RangeInclusive::new(start, end), &mut iter::empty());
+    }
+    node
+}
+
 #[must_use]
 fn insert_children<N: AstNode>(
     parent: &N,
diff --git a/crates/ra_syntax/src/ast/extensions.rs b/crates/ra_syntax/src/ast/extensions.rs
index 0433edb84..8c5ece65d 100644
--- a/crates/ra_syntax/src/ast/extensions.rs
+++ b/crates/ra_syntax/src/ast/extensions.rs
@@ -4,7 +4,7 @@
 use itertools::Itertools;
 
 use crate::{
-    ast::{self, child_opt, children, AstNode, SyntaxNode},
+    ast::{self, child_opt, children, AstChildren, AstNode, SyntaxNode},
     SmolStr, SyntaxElement,
     SyntaxKind::*,
     SyntaxToken, T,
@@ -203,6 +203,16 @@ impl ast::ImplBlock {
     }
 }
 
+impl ast::AttrsOwner for ast::ImplItem {
+    fn attrs(&self) -> AstChildren<ast::Attr> {
+        match self {
+            ast::ImplItem::FnDef(it) => it.attrs(),
+            ast::ImplItem::TypeAliasDef(it) => it.attrs(),
+            ast::ImplItem::ConstDef(it) => it.attrs(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum StructKind {
     Tuple(ast::TupleFieldDefList),