mdbook_preprocessor_utils/
processor.rs

1use 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    // Rather than copying directly to the build directory, we instead copy to the book source
54    // since mdBook will clean the build-dir after preprocessing. See mdBook#1087 for more.
55    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      // If a chapter is located at foo/bar/the_chapter.md, then the generated source files
75      // will be at foo/bar/the_chapter.html. So they need to reference preprocessor files
76      // at ../../<preprocessor>/embed.js, i.e. we generate the right number of "..".
77      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      // Ensure there's space between existing markdown and inserted HTML
82      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}