Cache embedding regexes (#730)

* Run cargo fmt and apply clippy lints from rust 1.51.0

* Rename can_perform_single_line_analisys to try_perform_single_line_analisys and streamline body

* Remove ops::Add impl for CodeStats; only AddAssign is meaningful and used

* Add embedding module to cache heuristic regexes used to detect embedded languages

* Add separate embedding tests; Not all cases can be execised with a single test

* Simplify html embedding test file_triggeringprincipal_frame_1.html until quoting issues are resolved

build.rs

@@ -69,7 +69,18 @@ fn generate_tests(out_dir: &OsStr) -> Result<(), Box<dyn error::Error>> {
     const INITIAL_BUFFER_SIZE: usize = 989 * 130;
     let mut string = String::with_capacity(INITIAL_BUFFER_SIZE);
 
-    let walker = Walk::new("./tests/data/").filter(|p| match p {
+    generate_tests_batch("./tests/data", None, &mut string)?;
+    generate_tests_batch("./tests/embedding", Some("embedding"), &mut string)?;
+
+    Ok(fs::write(Path::new(&out_dir).join("tests.rs"), string)?)
+}
+
+fn generate_tests_batch(
+    src_dir: &str,
+    test_module: Option<&str>,
+    string: &mut String,
+) -> Result<(), Box<dyn error::Error>> {
+    let walker = Walk::new(src_dir).filter(|p| match p {
         Ok(ref p) => {
             if let Ok(ref p) = p.metadata() {
                 p.is_file()
@@ -80,6 +91,17 @@ fn generate_tests(out_dir: &OsStr) -> Result<(), Box<dyn error::Error>> {
         _ => false,
     });
 
+    if let Some(test_module) = test_module {
+        string.push_str(&format!(
+            r####"
+#[cfg(test)]
+mod {0} {{
+use super::*;
+        "####,
+            test_module
+        ));
+    }
+
     for path in walker {
         let path = path?;
         let path = path.path();
@@ -129,10 +151,14 @@ fn generate_tests(out_dir: &OsStr) -> Result<(), Box<dyn error::Error>> {
         }}
         "####,
             name,
-            path.display(),
+            path.to_string_lossy().replace('\\', "/"),
             std::fs::canonicalize(root.join(path)).unwrap().display(),
         ));
     }
 
-    Ok(fs::write(Path::new(&out_dir).join("tests.rs"), string)?)
+    if test_module.is_some() {
+        string.push_str("\n}");
+    }
+
+    Ok(())
 }

src/cli.rs

@@ -169,10 +169,10 @@ impl<'a> Cli<'a> {
         crate::cli_utils::setup_logger(verbose);
 
         let cli = Cli {
+            matches,
             columns,
             files,
             hidden,
-            matches,
             no_ignore,
             no_ignore_parent,
             no_ignore_dot,
@@ -182,9 +182,9 @@ impl<'a> Cli<'a> {
             print_languages,
             sort,
             types,
-            verbose,
-            number_format,
             compact,
+            number_format,
+            verbose,
         };
 
         debug!("CLI Config: {:#?}", cli);

src/input.rs

@@ -1,5 +1,5 @@
-use std::{collections::BTreeMap, error::Error, str::FromStr};
 use serde_json::{json, Map};
+use std::{collections::BTreeMap, error::Error, str::FromStr};
 
 use tokei::{Language, LanguageType, Languages};
 

src/language/embedding.rs

@@ -0,0 +1,199 @@
+#![allow(clippy::trivial_regex)]
+
+use crate::LanguageType;
+use once_cell::sync::Lazy;
+use regex::bytes::Regex;
+
+pub static START_SCRIPT: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r#"<script(?:.*type="(.*)")?.*?>"#).unwrap());
+pub static END_SCRIPT: Lazy<Regex> = Lazy::new(|| Regex::new(r#"</script>"#).unwrap());
+
+pub static START_STYLE: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r#"<style(?:.*lang="(.*)")?.*?>"#).unwrap());
+pub static END_STYLE: Lazy<Regex> = Lazy::new(|| Regex::new(r#"</style>"#).unwrap());
+
+pub static START_TEMPLATE: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r#"<template(?:.*lang="(.*)")?.*?>"#).unwrap());
+pub static END_TEMPLATE: Lazy<Regex> = Lazy::new(|| Regex::new(r#"</template>"#).unwrap());
+
+pub static STARTING_MARKDOWN_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r#"```\S+\s"#).unwrap());
+pub static ENDING_MARKDOWN_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r#"```\s?"#).unwrap());
+
+/// A memory of a regex matched.
+/// The values provided by `Self::start` and `Self::end` are in the same space as the
+/// start value supplied to `RegexCache::build`
+pub struct Capture<'a> {
+    start: usize,
+    text: &'a [u8],
+}
+
+impl Capture<'_> {
+    #[inline(always)]
+    fn start(&self) -> usize {
+        self.start
+    }
+    #[inline(always)]
+    pub fn end(&self) -> usize {
+        self.start + self.text.len()
+    }
+    #[inline(always)]
+    pub fn as_bytes(&self) -> &[u8] {
+        self.text
+    }
+}
+
+impl<'a> std::fmt::Debug for Capture<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Capture")
+            .field("start", &self.start)
+            .field("end", &self.end())
+            .field("text", &String::from_utf8_lossy(&self.text))
+            .finish()
+    }
+}
+
+pub(crate) struct RegexCache<'a> {
+    inner: Option<RegexFamily<'a>>,
+}
+
+/// Embedding regexes are similar between different sets of languages.
+/// `RegexFamily` records both which family the language belongs to,
+/// as well as the actual matches
+pub(crate) enum RegexFamily<'a> {
+    HtmlLike(HtmlLike<'a>),
+    Markdown(Markdown<'a>),
+    Rust,
+}
+
+pub(crate) struct HtmlLike<'a> {
+    start_script: Option<Box<[Capture<'a>]>>,
+    start_style: Option<Box<[Capture<'a>]>>,
+    start_template: Option<Box<[Capture<'a>]>>,
+}
+
+pub(crate) struct Markdown<'a> {
+    starts: Option<Box<[Capture<'a>]>>,
+}
+
+impl<'a> HtmlLike<'a> {
+    pub fn start_script_in_range(
+        &'a self,
+        start: usize,
+        end: usize,
+    ) -> Option<impl Iterator<Item = &'a Capture<'a>>> {
+        filter_range(self.start_script.as_ref()?, start, end)
+    }
+
+    pub fn start_style_in_range(
+        &'a self,
+        start: usize,
+        end: usize,
+    ) -> Option<impl Iterator<Item = &'a Capture<'a>>> {
+        filter_range(self.start_style.as_ref()?, start, end)
+    }
+
+    pub fn start_template_in_range(
+        &'a self,
+        start: usize,
+        end: usize,
+    ) -> Option<impl Iterator<Item = &'a Capture<'a>>> {
+        filter_range(self.start_template.as_ref()?, start, end)
+    }
+}
+
+impl<'a> Markdown<'a> {
+    pub fn starts_in_range(&'a self, start: usize, end: usize) -> Option<&Capture<'a>> {
+        filter_range(self.starts.as_ref()?, start, end).and_then(|mut it| it.next())
+    }
+}
+
+fn filter_range<'a>(
+    dataset: &'a [Capture<'a>],
+    start: usize,
+    end: usize,
+) -> Option<impl Iterator<Item = &'a Capture<'a>>> {
+    let pos = dataset
+        .binary_search_by_key(&start, |cap| cap.start())
+        .ok()?;
+
+    if pos >= dataset.len() || dataset[pos].end() > end {
+        None
+    } else {
+        Some(
+            dataset[pos..]
+                .iter()
+                .take_while(move |cap| cap.end() <= end),
+        )
+    }
+}
+
+impl<'a> RegexCache<'a> {
+    /// Returns the language family for which regexes were matched, if any
+    pub(crate) fn family(&self) -> Option<&RegexFamily> {
+        self.inner.as_ref()
+    }
+
+    /// Tries to memoize any matches of embedding regexes that occur within lines[start..end]
+    /// for the given language. Any `Capture` values eventually recovered will use the same
+    /// zero for their start as the given `start` argument.
+    pub(crate) fn build(lang: LanguageType, lines: &'a [u8], start: usize, end: usize) -> Self {
+        let inner = match lang {
+            LanguageType::Markdown | LanguageType::UnrealDeveloperMarkdown => {
+                let markdown = Markdown {
+                    starts: save_captures(&STARTING_MARKDOWN_REGEX, lines, start, end),
+                };
+
+                if markdown.starts.is_some() {
+                    Some(RegexFamily::Markdown(markdown))
+                } else {
+                    None
+                }
+            }
+            LanguageType::Rust => Some(RegexFamily::Rust),
+            LanguageType::Html
+            | LanguageType::RubyHtml
+            | LanguageType::Svelte
+            | LanguageType::Vue => {
+                let html = HtmlLike {
+                    start_script: save_captures(&START_SCRIPT, lines, start, end),
+                    start_style: save_captures(&START_STYLE, lines, start, end),
+                    start_template: save_captures(&START_TEMPLATE, lines, start, end),
+                };
+
+                if html.start_script.is_some()
+                    || html.start_style.is_some()
+                    || html.start_template.is_some()
+                {
+                    Some(RegexFamily::HtmlLike(html))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        };
+        Self { inner }
+    }
+}
+
+fn save_captures<'a>(
+    regex: &Regex,
+    lines: &'a [u8],
+    start: usize,
+    end: usize,
+) -> Option<Box<[Capture<'a>]>> {
+    let v: Vec<_> = regex
+        .captures(&lines[start..end])?
+        .iter()
+        .flatten()
+        .map(|cap| Capture {
+            start: start + cap.start(),
+            text: cap.as_bytes(),
+        })
+        .collect();
+
+    if v.is_empty() {
+        None
+    } else {
+        Some(v.into())
+    }
+}

src/language/language_type.rs

@@ -142,7 +142,7 @@ impl LanguageType {
             };
             trace!("{}", String::from_utf8_lossy(line));
 
-            if syntax.can_perform_single_line_analysis(line, &mut stats) {
+            if syntax.try_perform_single_line_analysis(line, &mut stats) {
                 continue;
             }
 

src/language/language_type.tera.rs

@@ -7,6 +7,7 @@ use arbitrary::Arbitrary;
 #[derive(Deserialize, Serialize)]
 #[derive(Arbitrary, Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
 #[non_exhaustive]
+#[allow(clippy::upper_case_acronyms)]
 pub enum LanguageType {
     {% for key, _ in languages -%}
         #[allow(missing_docs)] {{key}},

src/language/mod.rs

@@ -1,3 +1,4 @@
+mod embedding;
 pub mod language_type;
 pub mod languages;
 mod syntax;

src/language/syntax.rs

@@ -5,8 +5,8 @@ use dashmap::DashMap;
 use grep_searcher::LineStep;
 use log::Level::Trace;
 use once_cell::sync::Lazy;
-use regex::bytes::Regex;
 
+use super::embedding::*;
 use crate::{stats::CodeStats, utils::ext::SliceExt, Config, LanguageType};
 
 /// Tracks the syntax of the language as well as the current state in the file.
@@ -168,38 +168,38 @@ impl SyntaxCounter {
 
     /// Try to see if we can determine what a line is from examining the whole
     /// line at once. Returns `true` if successful.
-    pub(crate) fn can_perform_single_line_analysis(
+    pub(crate) fn try_perform_single_line_analysis(
         &self,
         line: &[u8],
         stats: &mut crate::stats::CodeStats,
     ) -> bool {
-        if self.is_plain_mode() {
-            if line.trim().is_empty() {
-                stats.blanks += 1;
-                trace!("Blank No.{}", stats.blanks);
-                return true;
-            } else if !self.shared.important_syntax.is_match(line) {
-                trace!("^ Skippable");
+        if !self.is_plain_mode() {
+            false
+        } else if line.trim().is_empty() {
+            stats.blanks += 1;
+            trace!("Blank No.{}", stats.blanks);
+            true
+        } else if !self.shared.important_syntax.is_match(line) {
+            trace!("^ Skippable");
 
-                if self.shared.is_literate
-                    || self
-                        .shared
-                        .line_comments
-                        .iter()
-                        .any(|c| line.starts_with(c.as_bytes()))
-                {
-                    stats.comments += 1;
-                    trace!("Comment No.{}", stats.comments);
-                } else {
-                    stats.code += 1;
-                    trace!("Code No.{}", stats.code);
-                }
-
-                return true;
+            if self.shared.is_literate
+                || self
+                    .shared
+                    .line_comments
+                    .iter()
+                    .any(|c| line.starts_with(c.as_bytes()))
+            {
+                stats.comments += 1;
+                trace!("Comment No.{}", stats.comments);
+            } else {
+                stats.code += 1;
+                trace!("Code No.{}", stats.code);
             }
-        }
 
-        false
+            true
+        } else {
+            false
+        }
     }
 
     pub(crate) fn perform_multi_line_analysis(
@@ -217,6 +217,8 @@ impl SyntaxCounter {
             }};
         }
 
+        let regex_cache = RegexCache::build(self.shared.language, lines, start, end);
+
         for i in start..end {
             if skip != 0 {
                 skip -= 1;
@@ -242,7 +244,7 @@ impl SyntaxCounter {
                 continue;
             }
 
-            if let Some(child) = self.parse_context(lines, i, end, config) {
+            if let Some(child) = self.parse_context(lines, i, end, config, &regex_cache) {
                 return AnalysisReport::ChildLanguage(child);
             }
 
@@ -348,6 +350,7 @@ impl SyntaxCounter {
         start: usize,
         end: usize,
         config: &Config,
+        regex_cache: &RegexCache,
     ) -> Option<FileContext> {
         use std::str::FromStr;
 
@@ -356,19 +359,14 @@ impl SyntaxCounter {
             return None;
         }
 
-        match self.shared.language {
-            LanguageType::Markdown | LanguageType::UnrealDeveloperMarkdown => {
-                static STARTING_MARKDOWN_REGEX: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"^```\S+\s"#).unwrap());
-                static ENDING_MARKDOWN_REGEX: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"```\s?"#).unwrap());
-
+        match regex_cache.family()? {
+            RegexFamily::Markdown(md) => {
                 if !lines[start..end].contains_slice(b"```") {
                     return None;
                 }
 
-                let opening_fence = STARTING_MARKDOWN_REGEX.find(&lines[start..end])?;
-                let start_of_code = start + opening_fence.end();
+                let opening_fence = md.starts_in_range(start, end)?;
+                let start_of_code = opening_fence.end();
                 let closing_fence = ENDING_MARKDOWN_REGEX.find(&lines[start_of_code..]);
                 if let Some(m) = &closing_fence {
                     trace!("{:?}", String::from_utf8_lossy(m.as_bytes()))
@@ -400,7 +398,7 @@ impl SyntaxCounter {
                     stats,
                 ))
             }
-            LanguageType::Rust => {
+            RegexFamily::Rust => {
                 let rest = &lines[start..];
                 let comment_syntax = if rest.trim_start().starts_with(b"///") {
                     b"///"
@@ -436,28 +434,13 @@ impl SyntaxCounter {
                     doc_block,
                 ))
             }
-            #[allow(clippy::trivial_regex)]
-            LanguageType::Html
-            | LanguageType::RubyHtml
-            | LanguageType::Svelte
-            | LanguageType::Vue => {
-                static START_SCRIPT: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"^<script(?:.*type="(.*)")?.*?>"#).unwrap());
-                static END_SCRIPT: Lazy<Regex> = Lazy::new(|| Regex::new(r#"</script>"#).unwrap());
-                static START_STYLE: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"^<style(?:.*lang="(.*)")?.*?>"#).unwrap());
-                static END_STYLE: Lazy<Regex> = Lazy::new(|| Regex::new(r#"</style>"#).unwrap());
-                static START_TEMPLATE: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"^<template(?:.*lang="(.*)")?.*?>"#).unwrap());
-                static END_TEMPLATE: Lazy<Regex> =
-                    Lazy::new(|| Regex::new(r#"</template>"#).unwrap());
-
-                if let Some(captures) = START_SCRIPT.captures(&lines[start..end]) {
-                    let start_of_code = start + captures.get(0).unwrap().end();
+            RegexFamily::HtmlLike(html) => {
+                if let Some(mut captures) = html.start_script_in_range(start, end) {
+                    let start_of_code = captures.next().unwrap().end();
                     let closing_tag = END_SCRIPT.find(&lines[start_of_code..])?;
                     let end_of_code = start_of_code + closing_tag.start();
                     let language = captures
-                        .get(1)
+                        .next()
                         .and_then(|m| {
                             LanguageType::from_mime(&String::from_utf8_lossy(m.as_bytes().trim()))
                         })
@@ -476,12 +459,12 @@ impl SyntaxCounter {
                         end_of_code,
                         stats,
                     ))
-                } else if let Some(captures) = START_STYLE.captures(&lines[start..end]) {
-                    let start_of_code = start + captures.get(0).unwrap().end();
+                } else if let Some(mut captures) = html.start_style_in_range(start, end) {
+                    let start_of_code = captures.next().unwrap().end();
                     let closing_tag = END_STYLE.find(&lines[start_of_code..])?;
                     let end_of_code = start_of_code + closing_tag.start();
                     let language = captures
-                        .get(1)
+                        .next()
                         .and_then(|m| {
                             LanguageType::from_str(
                                 &String::from_utf8_lossy(m.as_bytes().trim()).to_lowercase(),
@@ -503,12 +486,12 @@ impl SyntaxCounter {
                         end_of_code,
                         stats,
                     ))
-                } else if let Some(captures) = START_TEMPLATE.captures(&lines[start..end]) {
-                    let start_of_code = start + captures.get(0).unwrap().end();
+                } else if let Some(mut captures) = html.start_template_in_range(start, end) {
+                    let start_of_code = captures.next().unwrap().end();
                     let closing_tag = END_TEMPLATE.find(&lines[start_of_code..])?;
                     let end_of_code = start_of_code + closing_tag.start();
                     let language = captures
-                        .get(1)
+                        .next()
                         .and_then(|m| {
                             LanguageType::from_str(
                                 &String::from_utf8_lossy(m.as_bytes().trim()).to_lowercase(),
@@ -534,7 +517,6 @@ impl SyntaxCounter {
                     None
                 }
             }
-            _ => None,
         }
     }
 

src/stats.rs

@@ -44,15 +44,6 @@ impl CodeStats {
     }
 }
 
-impl ops::Add for CodeStats {
-    type Output = Self;
-
-    fn add(mut self, rhs: Self) -> Self::Output {
-        self += rhs;
-        self
-    }
-}
-
 impl ops::AddAssign for CodeStats {
     fn add_assign(&mut self, rhs: Self) {
         self.blanks += rhs.blanks;

src/utils/ext.rs

@@ -1,22 +1,22 @@
 //! Various extensions to Rust std types.
 
 pub(crate) trait AsciiExt {
-    fn is_whitespace(self) -> bool;
-    fn is_not_line_ending_whitespace(self) -> bool;
-    fn is_line_ending_whitespace(self) -> bool;
+    fn is_whitespace(&self) -> bool;
+    fn is_not_line_ending_whitespace(&self) -> bool;
+    fn is_line_ending_whitespace(&self) -> bool;
 }
 
 impl AsciiExt for u8 {
-    fn is_whitespace(self) -> bool {
-        self == b' ' || (b'\x09'..=b'\x0d').contains(&self)
+    fn is_whitespace(&self) -> bool {
+        *self == b' ' || (b'\x09'..=b'\x0d').contains(&self)
     }
 
-    fn is_not_line_ending_whitespace(self) -> bool {
+    fn is_not_line_ending_whitespace(&self) -> bool {
         self.is_whitespace() && !self.is_line_ending_whitespace()
     }
 
-    fn is_line_ending_whitespace(self) -> bool {
-        self == b'\n'
+    fn is_line_ending_whitespace(&self) -> bool {
+        *self == b'\n'
     }
 }
 

tests/embedding/file_triggeringprincipal_frame_1.html

@@ -0,0 +1,27 @@
+<!-- 27 lines 20 code 5 comments 2 blanks -->
+< 0!DOCTYPE HTML>
+<html>
+<head><meta charset="utf-8"></head>
+<body>
+<b>Frame 1</b><br/>
+
+<script type="application/javascript">
+  // make sure to set document.domain to the same domain as the subframe
+  window.onload = function() {
+    document.domain = "mochi.test";
+  };
+  window.addEventListener("message", receiveMessage);
+  function receiveMessage(event) {
+    // make sure to get the right start command, otherwise
+    // let the parent know and fail the test
+    if (event.data.start !== "startTest") {
+      window.removeEventListener("message", receiveMessage);
+      window.parent.postMessage({triggeringPrincipalURI: "false"}, "*");
+    }
+    // click the link to navigate the subframe
+    document.getElementById("testlink").click();
+  }
+</script>
+
+</body>
+</html>