mdbook_preprocessor_utils/
processor.rs1use std::{
2 fs,
3 marker::PhantomData,
4 ops::Range,
5 path::{Path, PathBuf},
6};
7
8use anyhow::Result;
9use mdbook::{
10 book::Book,
11 preprocess::{Preprocessor, PreprocessorContext},
12 BookItem,
13};
14use rayon::prelude::*;
15
16#[derive(Copy, Clone)]
17pub struct Asset {
18 pub name: &'static str,
19 pub contents: &'static [u8],
20}
21
22#[macro_export]
23macro_rules! asset_generator {
24 ($base:expr) => {
25 macro_rules! make_asset {
26 ($name:expr) => {
27 $crate::Asset {
28 name: $name,
29 contents: include_bytes!(concat!($base, $name)),
30 }
31 };
32 }
33 };
34}
35
36pub trait SimplePreprocessor: Sized + Send + Sync {
37 type Args: clap::CommandFactory;
38 fn name() -> &'static str;
39 fn build(ctx: &PreprocessorContext) -> Result<Self>;
40 fn replacements(&self, chapter_dir: &Path, content: &str) -> Result<Vec<(Range<usize>, String)>>;
41 fn linked_assets(&self) -> Vec<Asset>;
42 fn all_assets(&self) -> Vec<Asset>;
43 fn finish(self) {}
44}
45
46struct SimplePreprocessorDriverCtxt<P: SimplePreprocessor> {
47 sp: P,
48 src_dir: PathBuf,
49}
50
51impl<P: SimplePreprocessor> SimplePreprocessorDriverCtxt<P> {
52 fn copy_assets(&self) -> Result<()> {
53 let dst_dir = self.src_dir.join(P::name());
56 fs::create_dir_all(&dst_dir)?;
57
58 for asset in self.sp.all_assets() {
59 fs::write(dst_dir.join(asset.name), asset.contents)?;
60 }
61
62 Ok(())
63 }
64
65 fn process_chapter(&self, chapter_dir: &Path, content: &mut String) -> Result<()> {
66 let mut replacements = self.sp.replacements(chapter_dir, content)?;
67 if !replacements.is_empty() {
68 replacements.sort_by_key(|(range, _)| range.start);
69
70 for (range, html) in replacements.into_iter().rev() {
71 content.replace_range(range, &html);
72 }
73
74 let chapter_rel_path = chapter_dir.strip_prefix(&self.src_dir).unwrap();
78 let depth = chapter_rel_path.components().count();
79 let prefix = vec![".."; depth].into_iter().collect::<PathBuf>();
80
81 content.push_str("\n\n");
83
84 for asset in self.sp.linked_assets() {
85 let asset_rel = prefix.join(P::name()).join(asset.name);
86 let asset_str = asset_rel.display().to_string();
87 let link = match &*asset_rel.extension().unwrap().to_string_lossy() {
88 "js" => format!(r#"<script type="text/javascript" src="{asset_str}"></script>"#),
89 "mjs" => format!(r#"<script type="module" src="{asset_str}"></script>"#),
90 "css" => format!(r#"<link rel="stylesheet" type="text/css" href="{asset_str}">"#),
91 _ => continue,
92 };
93 content.push_str(&link);
94 }
95 }
96 Ok(())
97 }
98}
99
100pub(crate) struct SimplePreprocessorDriver<P: SimplePreprocessor>(PhantomData<P>);
101
102impl<P: SimplePreprocessor> SimplePreprocessorDriver<P> {
103 pub fn new() -> Self {
104 SimplePreprocessorDriver(PhantomData)
105 }
106}
107
108impl<P: SimplePreprocessor> Preprocessor for SimplePreprocessorDriver<P> {
109 fn name(&self) -> &str {
110 P::name()
111 }
112
113 fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
114 let src_dir = ctx.root.join(&ctx.config.book.src);
115 let sp = P::build(ctx)?;
116 let ctxt = SimplePreprocessorDriverCtxt { sp, src_dir };
117 ctxt.copy_assets()?;
118
119 fn for_each_mut<'a, P: SimplePreprocessor>(
120 ctxt: &SimplePreprocessorDriverCtxt<P>,
121 chapters: &mut Vec<(PathBuf, &'a mut String)>,
122 items: impl IntoIterator<Item = &'a mut BookItem>,
123 ) {
124 for item in items {
125 if let BookItem::Chapter(chapter) = item {
126 if chapter.path.is_some() {
127 let chapter_path_abs = ctxt.src_dir.join(chapter.path.as_ref().unwrap());
128 let chapter_dir = chapter_path_abs.parent().unwrap().to_path_buf();
129 chapters.push((chapter_dir, &mut chapter.content));
130
131 for_each_mut(ctxt, chapters, &mut chapter.sub_items);
132 }
133 }
134 }
135 }
136
137 let mut chapters = Vec::new();
138 for_each_mut(&ctxt, &mut chapters, &mut book.sections);
139
140 chapters
141 .into_par_iter()
142 .map(|(chapter_dir, content)| ctxt.process_chapter(&chapter_dir, content))
143 .collect::<Result<Vec<_>>>()?;
144
145 ctxt.sp.finish();
146
147 Ok(book)
148 }
149}