Neater iterators based on code review feedback

jtojnar · Jan 26, 2021 · 2811feb · 2811feb
1 parent 44ea8a4
commit 2811feb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 41 deletions.
diff --git a/ast-checks/src/comment_finders.rs b/ast-checks/src/comment_finders.rs
@@ -3,17 +3,8 @@ use rnix::{NodeOrToken, SyntaxElement, SyntaxKind::*, SyntaxNode};
 
 pub fn find_comment_within(node: &SyntaxNode) -> Option<String> {
     // find the first TOKEN_COMMENT within
-    let tok = walk_kind(node, TOKEN_COMMENT).find_map(|item| match item {
-        NodeOrToken::Token(ref _tok) => Some(item),
-        NodeOrToken::Node(_) => {
-            unreachable!(
-                "`tok` should always be of kind==TOKEN_COMMENT,
-                which means that it's a SyntaxToken rather than
-                SyntaxNode"
-            );
-        }
-    })?;
-
+    let tok = walk_kind(node, TOKEN_COMMENT).next()?;
+    // iterate over sibling comments and concatenate them
     let node_iter = next_siblings(&tok);
     let doc = collect_comment_text(node_iter);
     Some(doc).filter(|it| !it.is_empty())
@@ -33,17 +24,11 @@ fn collect_comment_text(node_iter: impl Iterator<Item = SyntaxElement>) -> Strin
     // Note this would be more clearly written using map_while, but that
     // doesn't seem to be on rust stable yet.
     node_iter
-        .take_while(|node| match node {
-            NodeOrToken::Token(tok) if tok.kind() == TOKEN_COMMENT || tok.kind().is_trivia() => {
-                true
-            }
-            _ => false,
-        })
-        .filter_map(|node| match node {
-            NodeOrToken::Token(tok) => Some(tok.text().to_string()),
-            _ => None,
-        })
-        .map(|it| it.trim_start_matches('#').trim().to_string())
+        .filter_map(|node| node.into_token())
+        .take_while(|tok| tok.kind() == TOKEN_COMMENT || tok.kind().is_trivia())
+        .map(|tok| tok.text().to_string())
+        .map(|s| s.trim_start_matches('#').trim().to_string())
+        .filter(|s| !s.is_empty())
         .collect::<Vec<_>>()
         .join("\n")
 }
diff --git a/ast-checks/src/tree_utils.rs b/ast-checks/src/tree_utils.rs
@@ -1,6 +1,4 @@
-use rnix::{
-    types::*, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxKind::*, SyntaxNode, WalkEvent,
-};
+use rnix::{types::*, SyntaxElement, SyntaxKind, SyntaxKind::*, SyntaxNode, WalkEvent};
 use std::iter::{once, successors};
 
 pub fn walk_kind(node: &SyntaxNode, kind: SyntaxKind) -> impl Iterator<Item = SyntaxElement> {
@@ -15,22 +13,12 @@ pub fn walk_keyvalues_filter_key<'a>(
     node: &SyntaxNode,
     key: &'a str,
 ) -> impl Iterator<Item = KeyValue> + 'a {
-    walk_kind(node, NODE_KEY_VALUE).filter_map(move |element| match element {
-        NodeOrToken::Node(node) => {
-            let kv = KeyValue::cast(node).unwrap();
-            if kv.key().and_then(|e| Some(e.node().to_string())) == Some(key.to_string()) {
-                Some(kv)
-            } else {
-                None
-            }
-        }
-        NodeOrToken::Token(_tok) => {
-            unreachable!(
-                "walk_kind is iterating over kind=`NODE_KEY_VALUE`, which
-                          should only yield `SyntaxNode`s rather than `SytnaxToken`s"
-            );
-        }
-    })
+    walk_kind(node, NODE_KEY_VALUE)
+        .filter_map(|element| element.into_node())
+        .filter_map(|node| KeyValue::cast(node))
+        .filter(move |kv| {
+            kv.key().map_or(false, |e| e.node().to_string() == key.to_string())
+        })
 }
 
 pub fn next_siblings(element: &SyntaxElement) -> impl Iterator<Item = SyntaxElement> {