1use crate::error::{MdBookLintError, Result};
2use comrak::nodes::{AstNode, NodeValue};
3use comrak::{Arena, ComrakOptions, parse_document};
4use std::path::PathBuf;
5
6#[derive(Debug)]
8pub struct Document {
9 pub content: String,
11 pub path: PathBuf,
13 pub lines: Vec<String>,
15 pub book_src_dir: Option<PathBuf>,
17}
18
19impl Document {
20 pub fn new(content: String, path: PathBuf) -> Result<Self> {
22 Self::with_book_src_dir(content, path, None)
23 }
24
25 pub fn with_book_src_dir(
32 content: String,
33 path: PathBuf,
34 book_src_dir: Option<PathBuf>,
35 ) -> Result<Self> {
36 let lines: Vec<String> = content.lines().map(|s| s.to_owned()).collect();
41
42 Ok(Document {
43 content,
44 path,
45 lines,
46 book_src_dir,
47 })
48 }
49
50 pub fn parse_ast<'a>(&self, arena: &'a Arena<AstNode<'a>>) -> &'a AstNode<'a> {
52 let mut options = ComrakOptions::default();
54 options.extension.strikethrough = true;
55 options.extension.tagfilter = false;
56 options.extension.table = true;
57 options.extension.autolink = true;
58 options.extension.tasklist = true;
59 options.extension.superscript = false;
60 options.extension.header_ids = None;
61 options.extension.footnotes = true;
62 options.extension.description_lists = true;
63 options.extension.front_matter_delimiter = Some("---".to_owned());
64 options.parse.smart = false;
65 options.parse.default_info_string = None;
66 options.parse.relaxed_tasklist_matching = false;
67 options.parse.relaxed_autolinks = false;
68
69 parse_document(arena, &self.content, &options)
70 }
71
72 pub fn parse_ast_with_context<'a>(
74 &self,
75 arena: &'a Arena<AstNode<'a>>,
76 ) -> Result<&'a AstNode<'a>> {
77 let ast = self.parse_ast(arena);
79
80 if ast.children().count() == 0 && !self.content.trim().is_empty() {
82 return Err(MdBookLintError::Document(
83 "Failed to parse document AST - no content nodes found".to_string(),
84 ));
85 }
86
87 Ok(ast)
88 }
89
90 pub fn line_number_at_offset(&self, offset: usize) -> usize {
92 let mut current_offset = 0;
93 for (line_idx, line) in self.lines.iter().enumerate() {
94 if current_offset + line.len() >= offset {
95 return line_idx + 1; }
97 current_offset += line.len() + 1; }
99 self.lines.len() }
101
102 pub fn column_number_at_offset(&self, offset: usize) -> usize {
104 let mut current_offset = 0;
105 for line in &self.lines {
106 if current_offset + line.len() >= offset {
107 return offset - current_offset + 1; }
109 current_offset += line.len() + 1; }
111 1 }
113
114 pub fn line_ending(&self, line: usize) -> Option<&str> {
120 if line == 0 {
121 return None;
122 }
123
124 let source_line = self.content.split_inclusive('\n').nth(line - 1)?;
125 if source_line.ends_with("\r\n") {
126 Some("\r\n")
127 } else if source_line.ends_with('\n') {
128 Some("\n")
129 } else {
130 None
131 }
132 }
133
134 pub fn headings<'a>(&self, ast: &'a AstNode<'a>) -> Vec<&'a AstNode<'a>> {
136 let mut headings = Vec::new();
137 self.collect_headings(ast, &mut headings);
138 headings
139 }
140
141 pub fn headings_with_context<'a>(&self, ast: &'a AstNode<'a>) -> Result<Vec<&'a AstNode<'a>>> {
143 let headings = self.headings(ast);
144 Ok(headings)
145 }
146
147 pub fn code_blocks<'a>(&self, ast: &'a AstNode<'a>) -> Vec<&'a AstNode<'a>> {
149 let mut code_blocks = Vec::new();
150 self.collect_code_blocks(ast, &mut code_blocks);
151 code_blocks
152 }
153
154 pub fn code_blocks_with_context<'a>(
156 &self,
157 ast: &'a AstNode<'a>,
158 ) -> Result<Vec<&'a AstNode<'a>>> {
159 let code_blocks = self.code_blocks(ast);
160 Ok(code_blocks)
161 }
162
163 #[allow(clippy::only_used_in_recursion)]
165 fn collect_headings<'a>(&self, node: &'a AstNode<'a>, result: &mut Vec<&'a AstNode<'a>>) {
166 if let NodeValue::Heading(..) = &node.data.borrow().value {
167 result.push(node)
168 }
169
170 for child in node.children() {
172 self.collect_headings(child, result);
173 }
174 }
175
176 #[allow(clippy::only_used_in_recursion)]
178 fn collect_code_blocks<'a>(&self, node: &'a AstNode<'a>, result: &mut Vec<&'a AstNode<'a>>) {
179 if let NodeValue::CodeBlock(..) = &node.data.borrow().value {
180 result.push(node)
181 }
182
183 for child in node.children() {
185 self.collect_code_blocks(child, result);
186 }
187 }
188
189 pub fn heading_level<'a>(node: &'a AstNode<'a>) -> Option<u32> {
191 match &node.data.borrow().value {
192 NodeValue::Heading(heading) => Some(heading.level.into()),
193 _ => None,
194 }
195 }
196
197 pub fn node_text<'a>(&self, node: &'a AstNode<'a>) -> String {
199 let mut text = String::new();
200 self.collect_text(node, &mut text);
201 text
202 }
203
204 #[allow(clippy::only_used_in_recursion)]
206 fn collect_text<'a>(&self, node: &'a AstNode<'a>, text: &mut String) {
207 match &node.data.borrow().value {
208 NodeValue::Text(t) => text.push_str(t),
209 NodeValue::Code(code) => text.push_str(&code.literal),
210 _ => {
211 for child in node.children() {
212 self.collect_text(child, text);
213 }
214 }
215 }
216 }
217
218 pub fn node_position<'a>(&self, node: &'a AstNode<'a>) -> Option<(usize, usize)> {
220 let sourcepos = node.data.borrow().sourcepos;
221 if sourcepos.start.line > 0 {
222 Some((sourcepos.start.line, sourcepos.start.column))
223 } else {
224 None
225 }
226 }
227
228 pub fn frontmatter_line_range(&self) -> Option<(usize, usize)> {
241 if self.lines.first().map(String::as_str) != Some("---") {
244 return None;
245 }
246
247 self.lines
249 .iter()
250 .enumerate()
251 .skip(1)
252 .find(|(_, line)| line.as_str() == "---")
253 .map(|(idx, _)| (1, idx + 1))
254 }
255
256 pub fn frontmatter_ast_offset<'a>(&self, ast: &'a AstNode<'a>) -> usize {
265 ast.children()
266 .find_map(|child| match &child.data.borrow().value {
267 NodeValue::FrontMatter(text) => Some(text.matches('\n').count()),
268 _ => None,
269 })
270 .unwrap_or(0)
271 }
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277 use comrak::Arena;
278 use std::path::PathBuf;
279
280 #[test]
281 fn test_document_creation() {
282 let content = "# Test\n\nThis is a test.".to_string();
283 let path = PathBuf::from("test.md");
284
285 let doc = Document::new(content, path).expect("Failed to create document");
286
287 assert_eq!(doc.lines.len(), 3);
288 assert_eq!(doc.lines[0], "# Test");
289 assert_eq!(doc.lines[1], "");
290 assert_eq!(doc.lines[2], "This is a test.");
291 }
292
293 #[test]
294 fn test_empty_document_allowed() {
295 let content = "".to_string();
296 let path = PathBuf::from("empty.md");
297
298 let result = Document::new(content, path);
299 assert!(result.is_ok());
300
301 let document = result.unwrap();
302 assert_eq!(document.content, "");
303 assert_eq!(document.lines.len(), 0);
304 assert_eq!(document.path, PathBuf::from("empty.md"));
305 }
306
307 #[test]
308 fn test_whitespace_only_document_allowed() {
309 let content = " \n \n ".to_string();
310 let path = PathBuf::from("whitespace.md");
311
312 let result = Document::new(content, path);
313 assert!(result.is_ok());
314
315 let document = result.unwrap();
316 assert_eq!(document.content, " \n \n ");
317 assert_eq!(document.lines.len(), 3);
318 assert_eq!(document.path, PathBuf::from("whitespace.md"));
319 }
320
321 #[test]
322 fn test_line_number_calculation() {
323 let content = "Line 1\nLine 2\nLine 3".to_string();
324 let path = PathBuf::from("test.md");
325
326 let doc = Document::new(content, path).expect("Failed to create document");
327
328 assert_eq!(doc.line_number_at_offset(0), 1); assert_eq!(doc.line_number_at_offset(7), 2); assert_eq!(doc.line_number_at_offset(14), 3); }
332
333 #[test]
334 fn test_heading_extraction() {
335 let content = "# H1\n## H2\n### H3\nText".to_string();
336 let path = PathBuf::from("test.md");
337
338 let doc = Document::new(content, path).expect("Failed to create document");
339 let arena = Arena::new();
340 let ast = doc.parse_ast(&arena);
341 let headings = doc.headings(ast);
342
343 assert_eq!(headings.len(), 3);
344
345 assert_eq!(Document::heading_level(headings[0]), Some(1));
346 assert_eq!(Document::heading_level(headings[1]), Some(2));
347 assert_eq!(Document::heading_level(headings[2]), Some(3));
348 }
349
350 #[test]
351 fn test_document_with_unicode() {
352 let content = "# 标题\n\n这是一个测试。\n\n```rust\nfn main() {\n println!(\"Hello, 世界!\");\n}\n```".to_string();
353 let path = PathBuf::from("unicode.md");
354
355 let doc = Document::new(content, path).expect("Failed to create document");
356 assert!(doc.content.contains("标题"));
357 assert!(doc.content.contains("世界"));
358 assert_eq!(doc.lines.len(), 9);
359 }
360
361 #[test]
362 fn test_document_with_very_long_lines() {
363 let long_line = "a".repeat(10000);
364 let content = format!("# Test\n\n{long_line}\n\nEnd");
365 let path = PathBuf::from("long.md");
366
367 let doc = Document::new(content, path).expect("Failed to create document");
368 assert_eq!(doc.lines.len(), 5);
369 assert_eq!(doc.lines[2].len(), 10000);
370 }
371
372 #[test]
373 fn test_document_with_mixed_line_endings() {
374 let content = "Line 1\r\nLine 2\nLine 3\r\nLine 4".to_string();
375 let path = PathBuf::from("mixed.md");
376
377 let doc = Document::new(content, path).expect("Failed to create document");
378 assert_eq!(doc.lines.len(), 4);
380 assert_eq!(doc.lines[0], "Line 1");
381 assert_eq!(doc.lines[1], "Line 2");
382 assert_eq!(doc.lines[2], "Line 3");
383 assert_eq!(doc.lines[3], "Line 4");
384 assert_eq!(doc.line_ending(1), Some("\r\n"));
385 assert_eq!(doc.line_ending(2), Some("\n"));
386 assert_eq!(doc.line_ending(3), Some("\r\n"));
387 assert_eq!(doc.line_ending(4), None);
388 assert_eq!(doc.line_ending(0), None);
389 }
390
391 #[test]
392 fn test_document_with_only_newlines() {
393 let content = "\n\n\n\n".to_string();
394 let path = PathBuf::from("newlines.md");
395
396 let doc = Document::new(content, path).expect("Failed to create document");
397 assert_eq!(doc.lines.len(), 4);
398 for line in &doc.lines {
399 assert_eq!(line, "");
400 }
401 }
402
403 #[test]
404 fn test_node_position_edge_cases() {
405 let content = "# Test\n\n```rust\ncode\n```\n\n- Item".to_string();
406 let path = PathBuf::from("test.md");
407
408 let doc = Document::new(content, path).expect("Failed to create document");
409 let arena = Arena::new();
410 let ast = doc.parse_ast(&arena);
411
412 let position = doc.node_position(ast);
414 assert!(position.is_some());
415 let (line, col) = position.unwrap();
416 assert!(line >= 1);
417 assert!(col >= 1);
418 }
419
420 #[test]
421 fn test_code_blocks_extraction() {
422 let content = r#"# Test
423
424```rust
425fn main() {}
426```
427
428Some text.
429
430```bash
431echo "hello"
432```
433
434 // Indented code block
435 let x = 5;
436
437```
438No language
439```
440"#
441 .to_string();
442 let path = PathBuf::from("test.md");
443
444 let doc = Document::new(content, path).expect("Failed to create document");
445 let arena = Arena::new();
446 let ast = doc.parse_ast(&arena);
447 let code_blocks = doc.code_blocks(ast);
448
449 assert!(code_blocks.len() >= 3);
451 }
452
453 #[test]
454 fn test_links_extraction() {
455 let content = r#"# Test
456
457[Link 1](http://example.com)
458
459[Link 2](./relative.md)
460
461<https://autolink.com>
462
463[Reference link][ref]
464
465[ref]: http://reference.com
466"#
467 .to_string();
468 let path = PathBuf::from("test.md");
469
470 let doc = Document::new(content, path).expect("Failed to create document");
471 let arena = Arena::new();
472 let ast = doc.parse_ast(&arena);
473
474 assert!(!doc.content.is_empty());
476 assert!(!doc.lines.is_empty());
477
478 let mut has_heading = false;
480 for node in ast.descendants() {
481 if matches!(
482 node.data.borrow().value,
483 comrak::nodes::NodeValue::Heading(_)
484 ) {
485 has_heading = true;
486 break;
487 }
488 }
489 assert!(has_heading, "Expected to find heading in parsed AST");
490 }
491
492 #[test]
493 fn test_line_number_at_offset_edge_cases() {
494 let content = "a\nb\nc".to_string();
495 let path = PathBuf::from("test.md");
496
497 let doc = Document::new(content, path).expect("Failed to create document");
498
499 let line = doc.line_number_at_offset(100);
501 assert!(line >= 1);
502
503 assert_eq!(doc.line_number_at_offset(0), 1); assert_eq!(doc.line_number_at_offset(2), 2); assert_eq!(doc.line_number_at_offset(4), 3); }
508
509 #[test]
510 fn test_ast_parsing_with_extensions() {
511 let content = r#"# Test
512
513| Table | Header |
514|-------|--------|
515| Cell | Data |
516
517~~Strikethrough~~
518
519- [x] Task done
520- [ ] Task pending
521
522^Super^script
523
524[^footnote]: Footnote content
525
526Front matter:
527---
528title: Test
529---
530"#
531 .to_string();
532 let path = PathBuf::from("test.md");
533
534 let doc = Document::new(content, path).expect("Failed to create document");
535 let arena = Arena::new();
536 let ast = doc.parse_ast(&arena);
537
538 assert!(doc.content.contains("~~Strikethrough~~"));
540 assert!(doc.content.contains("| Table | Header |"));
541 assert!(doc.content.contains("- [x] Task done"));
542 assert!(doc.content.contains("title: Test"));
543
544 let node_count = ast.descendants().count();
546 assert!(
547 node_count > 5,
548 "Expected AST to contain multiple nodes, got {node_count}"
549 );
550 }
551
552 #[test]
553 fn test_frontmatter_line_range_detected() {
554 let content = "---\ntitle: My Document\nstatus: accepted\n---\n\n# My Document\n";
555 let doc = Document::new(content.to_string(), PathBuf::from("test.md")).unwrap();
556 assert_eq!(doc.frontmatter_line_range(), Some((1, 4)));
557 }
558
559 #[test]
560 fn test_frontmatter_line_range_none_without_frontmatter() {
561 let doc = Document::new("# Heading\n\nBody.\n".to_string(), PathBuf::from("t.md")).unwrap();
562 assert_eq!(doc.frontmatter_line_range(), None);
563 }
564
565 #[test]
566 fn test_frontmatter_line_range_requires_first_line() {
567 let doc =
569 Document::new("\n---\ntitle: x\n---\n".to_string(), PathBuf::from("t.md")).unwrap();
570 assert_eq!(doc.frontmatter_line_range(), None);
571 }
572
573 #[test]
574 fn test_frontmatter_line_range_requires_exact_delimiter() {
575 let doc =
577 Document::new("----\ntitle: x\n----\n".to_string(), PathBuf::from("t.md")).unwrap();
578 assert_eq!(doc.frontmatter_line_range(), None);
579 }
580
581 #[test]
582 fn test_frontmatter_line_range_requires_closing() {
583 let doc = Document::new("---\ntitle: x\n# H\n".to_string(), PathBuf::from("t.md")).unwrap();
584 assert_eq!(doc.frontmatter_line_range(), None);
585 }
586
587 #[test]
588 fn test_frontmatter_ast_offset_matches_comrak() {
589 let content = "---\ntitle: x\nstatus: y\n---\n\n# My Document\n";
592 let doc = Document::new(content.to_string(), PathBuf::from("test.md")).unwrap();
593 let arena = Arena::new();
594 let ast = doc.parse_ast(&arena);
595 let offset = doc.frontmatter_ast_offset(ast);
596 assert_eq!(offset, 5);
597
598 let heading = doc.headings(ast)[0];
599 let (line, _) = doc.node_position(heading).unwrap();
600 assert_eq!(line + offset, 6);
601 }
602
603 #[test]
604 fn test_frontmatter_ast_offset_zero_without_frontmatter() {
605 let doc = Document::new("# H\n\nBody.\n".to_string(), PathBuf::from("t.md")).unwrap();
606 let arena = Arena::new();
607 let ast = doc.parse_ast(&arena);
608 assert_eq!(doc.frontmatter_ast_offset(ast), 0);
609 }
610
611 #[test]
612 fn test_empty_path() {
613 let content = "# Test".to_string();
614 let path = PathBuf::new();
615
616 let doc = Document::new(content, path).expect("Failed to create document");
617 assert_eq!(doc.path, PathBuf::new());
618 }
619
620 #[test]
621 fn test_complex_nested_structure() {
622 let content = r#"# Main Title
623
624## Section 1
625
626### Subsection
627
628Some text with **bold** and *italic*.
629
630> Blockquote with `code` inside.
631>
632> > Nested blockquote.
633
6341. Ordered list
635 - Nested unordered
636 - Another item
6372. Second ordered item
638
639```rust
640// Code with comments
641fn complex_function() {
642 println!("Complex: {}", "test");
643}
644```
645
646Final paragraph.
647"#
648 .to_string();
649 let path = PathBuf::from("complex.md");
650
651 let doc = Document::new(content, path).expect("Failed to create document");
652 let arena = Arena::new();
653 let ast = doc.parse_ast(&arena);
654
655 let headings = doc.headings(ast);
657 assert!(headings.len() >= 3);
658
659 let code_blocks = doc.code_blocks(ast);
660 assert!(!code_blocks.is_empty());
661
662 assert!(doc.content.contains("# Main Title"));
664 assert!(doc.content.contains("## Section 1"));
665 assert!(doc.content.contains("### Subsection"));
666 assert!(doc.content.contains("```rust"));
667 assert!(doc.content.contains("Final paragraph"));
668
669 assert!(
671 doc.lines.len() > 10,
672 "Expected multiple lines in complex document"
673 );
674 }
675}