move shared code to query_testing

cli/src/query_testing.rs

@@ -0,0 +1,153 @@
+use crate::error;
+use crate::error::Result;
+use lazy_static::lazy_static;
+use regex::Regex;
+use std::collections::hash_map::HashMap;
+use std::fs;
+use tree_sitter::{Language, Parser, Point};
+
+lazy_static! {
+    static ref HIGHLIGHT_NAME_REGEX: Regex = Regex::new("[\\w_\\-.]+").unwrap();
+}
+
+#[derive(Debug, Eq, PartialEq)]
+pub struct CaptureInfo {
+    pub name: String,
+    pub position: Point,
+}
+
+pub struct Assertion {
+    pub position: Point,
+    pub expected: String,
+}
+
+/// Parse the given source code, finding all of the comments that contain
+/// highlighting assertions. Return a vector of (position, expected highlight name)
+/// pairs.
+pub fn parse_highlight_test(
+    parser: &mut Parser,
+    language: Language,
+    source: &[u8],
+) -> Result<Vec<Assertion>> {
+    let mut result = Vec::new();
+    let mut assertion_ranges = Vec::new();
+
+    // Parse the code.
+    parser.set_included_ranges(&[]).unwrap();
+    parser.set_language(language).unwrap();
+    let tree = parser.parse(source, None).unwrap();
+
+    // Walk the tree, finding comment nodes that contain assertions.
+    let mut ascending = false;
+    let mut cursor = tree.root_node().walk();
+    loop {
+        if ascending {
+            let node = cursor.node();
+
+            // Find every comment node.
+            if node.kind().contains("comment") {
+                if let Ok(text) = node.utf8_text(source) {
+                    let mut position = node.start_position();
+                    if position.row == 0 {
+                        continue;
+                    }
+
+                    // Find the arrow character ("^" or '<-") in the comment. A left arrow
+                    // refers to the column where the comment node starts. An up arrow refers
+                    // to its own column.
+                    let mut has_left_caret = false;
+                    let mut has_arrow = false;
+                    let mut arrow_end = 0;
+                    for (i, c) in text.char_indices() {
+                        arrow_end = i + 1;
+                        if c == '-' && has_left_caret {
+                            has_arrow = true;
+                            break;
+                        }
+                        if c == '^' {
+                            has_arrow = true;
+                            position.column += i;
+                            break;
+                        }
+                        has_left_caret = c == '<';
+                    }
+
+                    // If the comment node contains an arrow and a highlight name, record the
+                    // highlight name and the position.
+                    if let (true, Some(mat)) =
+                        (has_arrow, HIGHLIGHT_NAME_REGEX.find(&text[arrow_end..]))
+                    {
+                        assertion_ranges.push((node.start_position(), node.end_position()));
+                        result.push(Assertion {
+                            position: position,
+                            expected: mat.as_str().to_string(),
+                        });
+                    }
+                }
+            }
+
+            // Continue walking the tree.
+            if cursor.goto_next_sibling() {
+                ascending = false;
+            } else if !cursor.goto_parent() {
+                break;
+            }
+        } else if !cursor.goto_first_child() {
+            ascending = true;
+        }
+    }
+
+    // Adjust the row number in each assertion's position to refer to the line of
+    // code *above* the assertion. There can be multiple lines of assertion comments,
+    // so the positions may have to be decremented by more than one row.
+    let mut i = 0;
+    for assertion in result.iter_mut() {
+        loop {
+            let on_assertion_line = assertion_ranges[i..]
+                .iter()
+                .any(|(start, _)| start.row == assertion.position.row);
+            if on_assertion_line {
+                assertion.position.row -= 1;
+            } else {
+                while i < assertion_ranges.len()
+                    && assertion_ranges[i].0.row < assertion.position.row
+                {
+                    i += 1;
+                }
+                break;
+            }
+        }
+    }
+
+    // The assertions can end up out of order due to the line adjustments.
+    result.sort_unstable_by_key(|a| a.position);
+
+    Ok(result)
+}
+
+pub fn assert_expected_captures(
+    infos: Vec<CaptureInfo>,
+    path: String,
+    parser: &mut Parser,
+    language: Language,
+) -> Result<()> {
+    let contents = fs::read_to_string(path)?;
+    let pairs = parse_highlight_test(parser, language, contents.as_bytes())?;
+
+    let per_position_index: HashMap<Point, &String> =
+        pairs.iter().map(|a| (a.position, &a.expected)).collect();
+
+    for info in &infos {
+        if !per_position_index.contains_key(&info.position) {
+            continue;
+        }
+        let found = per_position_index.get(&info.position).unwrap();
+        if **found != info.name && info.name != "name" {
+            Err(error::Error::new(format!(
+                "Assertion failed: at {}, found {}, expected {}",
+                info.position, found, info.name
+            )))?
+        }
+    }
+    Ok(())
+}