use super::comment_patterns::CommentPattern;
use crate::core::types::FileStats;
use std::io::{self, BufRead};
const MAX_RETAINED_LINE_CAPACITY: usize = 1 << 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct LineTally {
pub total_lines: usize,
pub code_lines: usize,
pub comment_lines: usize,
pub doc_lines: usize,
pub blank_lines: usize,
}
impl LineTally {
pub fn into_file_stats(self, file_size: u64) -> FileStats {
FileStats {
total_lines: self.total_lines,
code_lines: self.code_lines,
comment_lines: self.comment_lines,
blank_lines: self.blank_lines,
file_size,
doc_lines: self.doc_lines,
}
}
pub fn is_partitioned(&self) -> bool {
self.code_lines + self.comment_lines + self.doc_lines + self.blank_lines == self.total_lines
}
}
fn trim_ascii(mut bytes: &[u8]) -> &[u8] {
while let [first, rest @ ..] = bytes {
if first.is_ascii_whitespace() {
bytes = rest;
} else {
break;
}
}
while let [rest @ .., last] = bytes {
if last.is_ascii_whitespace() {
bytes = rest;
} else {
break;
}
}
bytes
}
const BOM: &[u8] = b"\xef\xbb\xbf";
fn strip_bom(bytes: &[u8]) -> &[u8] {
bytes.strip_prefix(BOM).unwrap_or(bytes)
}
fn is_shebang(line: &[u8]) -> bool {
line.starts_with(b"#!")
}
fn find_sub(haystack: &[u8], needle: &[u8]) -> Option<usize> {
memchr::memmem::find(haystack, needle)
}
fn contains_sub(haystack: &[u8], needle: &[u8]) -> bool {
find_sub(haystack, needle).is_some()
}
fn starts_with_any(line: &[u8], prefixes: &[&str]) -> bool {
prefixes
.iter()
.any(|prefix| line.starts_with(prefix.as_bytes()))
}
fn contains_any(line: &[u8], needles: &[&str]) -> bool {
needles
.iter()
.any(|needle| contains_sub(line, needle.as_bytes()))
}
struct BlockOpener {
finder: memchr::memmem::Finder<'static>,
opener_len: usize,
end: &'static str,
symmetric: bool,
}
impl BlockOpener {
fn prepare(pattern: CommentPattern) -> Vec<Self> {
pattern
.multi_line_start
.iter()
.enumerate()
.map(|(index, opener)| {
let end = pattern.multi_line_end.get(index).copied().unwrap_or(opener);
Self {
finder: memchr::memmem::Finder::new(opener.as_bytes()),
opener_len: opener.len(),
end,
symmetric: end == *opener,
}
})
.collect()
}
}
fn next_line<R: BufRead>(reader: &mut R, buf: &mut Vec<u8>) -> io::Result<bool> {
if buf.capacity() > MAX_RETAINED_LINE_CAPACITY {
*buf = Vec::with_capacity(256);
} else {
buf.clear();
}
Ok(reader.read_until(b'\n', buf)? != 0)
}
fn split_lines(bytes: &[u8]) -> impl Iterator<Item = &[u8]> {
let mut rest = bytes;
std::iter::from_fn(move || {
if rest.is_empty() {
return None;
}
let line = match memchr::memchr(b'\n', rest) {
Some(at) => {
let (line, tail) = rest.split_at(at + 1);
rest = tail;
line
}
None => std::mem::take(&mut rest),
};
Some(line)
})
}
struct Classifier {
pattern: CommentPattern,
openers: Vec<BlockOpener>,
doc_marker_is_ambiguous: bool,
in_block: bool,
in_doc_block: bool,
block_end: &'static str,
tally: LineTally,
}
impl Classifier {
fn new(pattern: CommentPattern) -> Self {
Self {
pattern,
openers: BlockOpener::prepare(pattern),
doc_marker_is_ambiguous: pattern
.doc_patterns
.iter()
.any(|doc| pattern.single_line.contains(doc)),
in_block: false,
in_doc_block: false,
block_end: "",
tally: LineTally::default(),
}
}
fn feed(&mut self, bytes: &[u8]) {
self.tally.total_lines += 1;
let first_line = self.tally.total_lines == 1;
let raw: &[u8] = if first_line { strip_bom(bytes) } else { bytes };
let indented = raw.first().is_some_and(u8::is_ascii_whitespace);
let line = trim_ascii(raw);
if line.is_empty() {
self.tally.blank_lines += 1;
return;
}
if first_line && is_shebang(line) {
self.tally.code_lines += 1;
return;
}
if self.in_block {
let is_doc = self.in_doc_block;
if contains_sub(line, self.block_end.as_bytes()) {
self.in_block = false;
self.in_doc_block = false;
}
self.count_comment(is_doc);
return;
}
let opened = self
.openers
.iter()
.filter_map(|opener| {
let at = opener.finder.find(line)?;
(!opener.symmetric || at == 0).then_some((opener, at))
})
.min_by_key(|&(_, at)| at);
if let Some((opener, at)) = opened {
let is_doc = contains_any(line, self.pattern.doc_patterns);
if !contains_sub(&line[at + opener.opener_len..], opener.end.as_bytes()) {
self.in_block = true;
self.in_doc_block = is_doc;
self.block_end = opener.end;
}
self.count_comment(is_doc);
} else if starts_with_any(line, self.pattern.single_line) {
let is_doc = starts_with_any(line, self.pattern.doc_patterns)
&& !(self.doc_marker_is_ambiguous && indented);
self.count_comment(is_doc);
} else {
self.tally.code_lines += 1;
}
}
fn count_comment(&mut self, is_doc: bool) {
if is_doc {
self.tally.doc_lines += 1;
} else {
self.tally.comment_lines += 1;
}
}
fn finish(self) -> LineTally {
debug_assert!(self.tally.is_partitioned());
self.tally
}
}
pub fn classify<R: BufRead>(reader: &mut R, pattern: CommentPattern) -> io::Result<LineTally> {
let mut classifier = Classifier::new(pattern);
let mut buf = Vec::with_capacity(256);
while next_line(reader, &mut buf)? {
classifier.feed(&buf);
}
Ok(classifier.finish())
}
pub fn classify_bytes(bytes: &[u8], pattern: CommentPattern) -> LineTally {
let mut classifier = Classifier::new(pattern);
for line in split_lines(bytes) {
classifier.feed(line);
}
classifier.finish()
}
pub fn classify_markdown<R: BufRead>(reader: &mut R) -> io::Result<LineTally> {
let mut classifier = MarkdownClassifier::default();
let mut buf = Vec::with_capacity(256);
while next_line(reader, &mut buf)? {
classifier.feed(&buf);
}
Ok(classifier.finish())
}
pub fn classify_markdown_bytes(bytes: &[u8]) -> LineTally {
let mut classifier = MarkdownClassifier::default();
for line in split_lines(bytes) {
classifier.feed(line);
}
classifier.finish()
}
#[derive(Default)]
struct MarkdownClassifier {
in_code_block: bool,
in_html_comment: bool,
tally: LineTally,
}
impl MarkdownClassifier {
fn feed(&mut self, bytes: &[u8]) {
self.tally.total_lines += 1;
let bytes: &[u8] = if self.tally.total_lines == 1 {
strip_bom(bytes)
} else {
bytes
};
let raw = trim_ascii_end(bytes);
let line = trim_ascii(bytes);
if line.is_empty() {
self.tally.blank_lines += 1;
return;
}
if line.starts_with(b"<!--") {
self.in_html_comment = true;
}
if self.in_html_comment {
self.tally.comment_lines += 1;
if line.ends_with(b"-->") {
self.in_html_comment = false;
}
return;
}
if line.starts_with(b"```") || line.starts_with(b"~~~") {
self.in_code_block = !self.in_code_block;
self.tally.code_lines += 1;
return;
}
if self.in_code_block || raw.starts_with(b" ") || raw.starts_with(b"\t") {
self.tally.code_lines += 1;
} else {
self.tally.doc_lines += 1;
}
}
fn finish(self) -> LineTally {
debug_assert!(self.tally.is_partitioned());
self.tally
}
}
fn trim_ascii_end(mut bytes: &[u8]) -> &[u8] {
while let [rest @ .., last] = bytes {
if *last == b'\n' || *last == b'\r' {
bytes = rest;
} else {
break;
}
}
bytes
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::counter::comment_patterns;
fn scan(source: &str, ext: &str) -> LineTally {
scan_bytes(source.as_bytes(), ext)
}
fn scan_bytes(source: &[u8], ext: &str) -> LineTally {
let pattern = comment_patterns::lookup_or_empty(ext);
let streamed = classify(&mut &source[..], pattern).unwrap();
let in_memory = classify_bytes(source, pattern);
assert_eq!(
streamed,
in_memory,
"the streaming and in-memory classifiers disagreed on {:?}",
String::from_utf8_lossy(source)
);
streamed
}
#[test]
fn trim_ascii_strips_both_ends_and_terminators() {
assert_eq!(trim_ascii(b" hi \t\r\n"), b"hi");
assert_eq!(trim_ascii(b"\r\n"), b"");
assert_eq!(trim_ascii(b""), b"");
assert_eq!(trim_ascii(b"x"), b"x");
}
#[test]
fn contains_sub_matches_std_semantics() {
let cases: &[(&str, &str)] = &[
("hello world", "world"),
("hello world", "hello"),
("hello", "hello"),
("hello", "helloo"),
("aaa", "aa"),
("abc", "d"),
("", "a"),
("abc", ""),
("/* doc */", "*/"),
];
for (hay, needle) in cases {
assert_eq!(
contains_sub(hay.as_bytes(), needle.as_bytes()),
hay.contains(needle),
"contains_sub disagreed with str::contains for {hay:?} / {needle:?}"
);
}
}
#[test]
fn empty_input_yields_zero_lines() {
assert_eq!(scan("", "rs"), LineTally::default());
}
#[test]
fn line_counts_match_bufread_lines() {
for source in [
"",
"a",
"a\n",
"a\nb",
"a\nb\n",
"\n",
"\n\n",
"a\n\nb\n",
"a\r\nb\r\n",
] {
let expected = source.as_bytes().lines().count();
let tally = scan(source, "rs");
assert_eq!(
tally.total_lines, expected,
"total_lines disagreed with BufRead::lines for {source:?}"
);
}
}
#[test]
fn blank_lines_are_counted_separately() {
let tally = scan("\n\n\n\n", "rs");
assert_eq!(tally.total_lines, 4);
assert_eq!(tally.blank_lines, 4);
assert_eq!(tally.code_lines, 0);
}
#[test]
fn whitespace_only_lines_are_blank() {
let tally = scan(" \n\t\n \t \n", "rs");
assert_eq!(tally.blank_lines, 3);
assert_eq!(tally.code_lines, 0);
}
#[test]
fn rust_comments_docs_and_code_are_distinguished() {
let tally = scan(
"// plain comment\n/// doc comment\n//! inner doc\nfn main() {}\n\n",
"rs",
);
assert_eq!(tally.total_lines, 5);
assert_eq!(tally.comment_lines, 1);
assert_eq!(tally.doc_lines, 2);
assert_eq!(tally.code_lines, 1);
assert_eq!(tally.blank_lines, 1);
}
#[test]
fn block_comments_span_lines_and_terminate() {
let tally = scan("/*\nstill comment\n*/\nfn main() {}\n", "c");
assert_eq!(tally.comment_lines, 3);
assert_eq!(tally.code_lines, 1);
}
#[test]
fn unterminated_block_comment_consumes_the_rest() {
let tally = scan("/*\na\nb\nc\n", "c");
assert_eq!(tally.total_lines, 4);
assert_eq!(tally.comment_lines, 4);
assert_eq!(tally.code_lines, 0);
assert!(tally.is_partitioned());
}
#[test]
fn crlf_is_handled_identically_to_lf() {
let lf = scan("// c\nfn main() {}\n\n", "rs");
let crlf = scan("// c\r\nfn main() {}\r\n\r\n", "rs");
assert_eq!(lf, crlf, "CRLF input produced a different tally than LF");
}
#[test]
fn invalid_utf8_is_counted_not_rejected() {
let source = b"fn main() {}\n\xff\xfe not utf8 \x80\n// comment\n";
let tally = scan_bytes(source, "rs");
assert_eq!(tally.total_lines, 3);
assert_eq!(tally.comment_lines, 1);
assert_eq!(tally.code_lines, 2);
assert!(tally.is_partitioned());
}
#[test]
fn a_line_with_no_terminator_still_counts() {
let tally = scan("fn a() {}\nfn b() {}", "rs");
assert_eq!(tally.total_lines, 2);
assert_eq!(tally.code_lines, 2);
}
#[test]
fn very_long_lines_are_classified_correctly() {
let long = format!("// {}\n", "x".repeat(200_000));
let tally = scan(&long, "rs");
assert_eq!(tally.total_lines, 1);
assert_eq!(tally.comment_lines, 1);
}
#[test]
fn oversized_line_buffer_is_released() {
let mut buf = Vec::with_capacity(MAX_RETAINED_LINE_CAPACITY * 2);
let mut reader = &b"short\n"[..];
next_line(&mut reader, &mut buf).unwrap();
assert!(
buf.capacity() < MAX_RETAINED_LINE_CAPACITY,
"oversized line buffer was retained"
);
}
#[test]
fn python_docstrings_count_as_documentation() {
let tally = scan(
"def f():\n \"\"\"Doc line\n more doc\n \"\"\"\n # comment\n return 1\n",
"py",
);
assert!(tally.doc_lines >= 3, "docstring lines: {tally:?}");
assert_eq!(tally.comment_lines, 1);
assert!(tally.code_lines >= 2);
}
#[test]
fn go_doc_comments_are_distinguished_from_body_comments() {
let tally = scan(
"// Package p does things.\npackage p\n\nfunc f() {\n\t// step one\n\treturn\n}\n",
"go",
);
assert_eq!(tally.doc_lines, 1, "{tally:?}");
assert_eq!(tally.comment_lines, 1, "{tally:?}");
assert_eq!(tally.code_lines, 4);
assert!(tally.is_partitioned());
}
#[test]
fn indented_rust_doc_comments_remain_documentation() {
let tally = scan("impl T {\n /// Doc.\n fn f() {}\n}\n", "rs");
assert_eq!(tally.doc_lines, 1, "{tally:?}");
assert_eq!(tally.comment_lines, 0);
}
#[test]
fn single_line_docstring_closes_itself() {
let tally = scan("def f():\n \"\"\"One liner.\"\"\"\n return 1\n", "py");
assert_eq!(tally.doc_lines, 1);
assert_eq!(tally.code_lines, 2);
}
#[test]
fn multiline_string_body_is_code_not_documentation() {
let tally = scan(
"query = '''SELECT 1\nFROM t\nWHERE x = 2'''\nrun(query)\n",
"py",
);
assert_eq!(tally.doc_lines, 0, "{tally:?}");
assert_eq!(tally.code_lines, 4);
}
#[test]
fn unterminated_block_comment_stays_a_comment() {
let tally = scan("code();\n/* opened\nstill inside\nnever closed\n", "rs");
assert_eq!(tally.code_lines, 1);
assert_eq!(tally.comment_lines, 3);
assert!(tally.is_partitioned());
}
#[test]
fn languages_without_comment_syntax_treat_everything_as_code() {
let tally = scan("a\nb\n\n", "unknown-extension");
assert_eq!(tally.code_lines, 2);
assert_eq!(tally.blank_lines, 1);
assert_eq!(tally.comment_lines, 0);
}
#[test]
fn markdown_separates_prose_code_and_comments() {
let source = "# Title\n\nProse here.\n\n```rust\nfn main() {}\n```\n\n<!-- note -->\n";
let tally = classify_markdown(&mut source.as_bytes()).unwrap();
assert_eq!(tally.total_lines, 9);
assert_eq!(tally.blank_lines, 3);
assert_eq!(tally.comment_lines, 1);
assert_eq!(tally.code_lines, 3, "two fences plus one body line");
assert_eq!(tally.doc_lines, 2);
assert!(tally.is_partitioned());
}
#[test]
fn markdown_indented_code_is_code() {
let tally = classify_markdown(&mut "prose\n\n indented code\n".as_bytes()).unwrap();
assert_eq!(tally.doc_lines, 1);
assert_eq!(tally.code_lines, 1);
assert_eq!(tally.blank_lines, 1);
}
#[test]
fn categories_always_partition_the_total() {
let sources: &[&str] = &[
"",
"\n",
"// a\n/* b\nc */\nd\n",
"/*\n*/\n/*\n*/\n",
"\"\"\"\na\n\"\"\"\n",
"#\n##\n###\n",
"-- a\n{- b -}\n",
"<!-- a -->\nb\n",
"a\r\n\r\n\tb\n",
];
for ext in ["rs", "py", "c", "hs", "html", "md", "sh", "unknown"] {
for source in sources {
let tally = scan(source, ext);
assert!(
tally.is_partitioned(),
"categories did not partition total for ext={ext} source={source:?} -> {tally:?}"
);
}
}
}
}