From 2811febfcb8035baf66390e0c3f2935ac9fa90ad Mon Sep 17 00:00:00 2001
From: "Robert T. McGibbon" <rmcgibbo@gmail.com>
Date: Mon, 25 Jan 2021 14:29:43 -0500
Subject: [PATCH] Neater iterators based on code review feedback

---
 ast-checks/src/comment_finders.rs | 29 +++++++----------------------
 ast-checks/src/tree_utils.rs      | 26 +++++++-------------------
 2 files changed, 14 insertions(+), 41 deletions(-)

diff --git a/ast-checks/src/comment_finders.rs b/ast-checks/src/comment_finders.rs
index bcec10c..ab38004 100644
--- 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
index ce8e376..bc8f9df 100644
--- 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> {