Skip to main content

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_preprocessor::{
10  Preprocessor, PreprocessorContext,
11  book::{Book, BookItem},
12};
13// use mdbook::{
14//   book::Book,
15//   preprocess::{Preprocessor, PreprocessorContext},
16//   BookItem,
17// };
18use rayon::prelude::*;
19
20#[derive(Copy, Clone)]
21pub struct Asset {
22  pub name: &'static str,
23  pub contents: &'static [u8],
24}
25
26#[macro_export]
27macro_rules! asset_generator {
28  ($base:expr) => {
29    macro_rules! make_asset {
30      ($name:expr) => {
31        $crate::Asset {
32          name: $name,
33          contents: include_bytes!(concat!($base, $name)),
34        }
35      };
36    }
37  };
38}
39
40pub trait SimplePreprocessor: Sized + Send + Sync {
41  type Args: clap::CommandFactory;
42  fn name() -> &'static str;
43  fn build(ctx: &PreprocessorContext) -> Result<Self>;
44  fn replacements(&self, chapter_dir: &Path, content: &str) -> Result<Vec<(Range<usize>, String)>>;
45  fn linked_assets(&self) -> Vec<Asset>;
46  fn all_assets(&self) -> Vec<Asset>;
47  fn finish(self) {}
48}
49
50struct SimplePreprocessorDriverCtxt<P: SimplePreprocessor> {
51  sp: P,
52  src_dir: PathBuf,
53}
54
55impl<P: SimplePreprocessor> SimplePreprocessorDriverCtxt<P> {
56  fn copy_assets(&self) -> Result<()> {
57    // Rather than copying directly to the build directory, we instead copy to the book source
58    // since mdBook will clean the build-dir after preprocessing. See mdBook#1087 for more.
59    let dst_dir = self.src_dir.join(P::name());
60    fs::create_dir_all(&dst_dir)?;
61
62    for asset in self.sp.all_assets() {
63      fs::write(dst_dir.join(asset.name), asset.contents)?;
64    }
65
66    Ok(())
67  }
68
69  fn process_chapter(&self, chapter_dir: &Path, content: &mut String) -> Result<()> {
70    let mut replacements = self.sp.replacements(chapter_dir, content)?;
71    if !replacements.is_empty() {
72      replacements.sort_by_key(|(range, _)| range.start);
73
74      for (range, html) in replacements.into_iter().rev() {
75        content.replace_range(range, &html);
76      }
77
78      // If a chapter is located at foo/bar/the_chapter.md, then the generated source files
79      // will be at foo/bar/the_chapter.html. So they need to reference preprocessor files
80      // at ../../<preprocessor>/embed.js, i.e. we generate the right number of "..".
81      let chapter_rel_path = chapter_dir.strip_prefix(&self.src_dir).unwrap();
82      let depth = chapter_rel_path.components().count();
83      let prefix = vec![".."; depth].into_iter().collect::<PathBuf>();
84
85      // Ensure there's space between existing markdown and inserted HTML
86      content.push_str("\n\n");
87
88      for asset in self.sp.linked_assets() {
89        let asset_rel = prefix.join(P::name()).join(asset.name);
90        let asset_str = asset_rel.display().to_string();
91        let link = match &*asset_rel.extension().unwrap().to_string_lossy() {
92          "js" => format!(r#"<script type="text/javascript" src="{asset_str}"></script>"#),
93          "mjs" => format!(r#"<script type="module" src="{asset_str}"></script>"#),
94          "css" => format!(r#"<link rel="stylesheet" type="text/css" href="{asset_str}">"#),
95          _ => continue,
96        };
97        content.push_str(&link);
98      }
99    }
100    Ok(())
101  }
102}
103
104pub(crate) struct SimplePreprocessorDriver<P: SimplePreprocessor>(PhantomData<P>);
105
106impl<P: SimplePreprocessor> SimplePreprocessorDriver<P> {
107  pub fn new() -> Self {
108    SimplePreprocessorDriver(PhantomData)
109  }
110}
111
112impl<P: SimplePreprocessor> Preprocessor for SimplePreprocessorDriver<P> {
113  fn name(&self) -> &str {
114    P::name()
115  }
116
117  fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
118    let src_dir = ctx.root.join(&ctx.config.book.src);
119    let sp = P::build(ctx)?;
120    let ctxt = SimplePreprocessorDriverCtxt { sp, src_dir };
121    ctxt.copy_assets()?;
122
123    fn for_each_mut<'a, P: SimplePreprocessor>(
124      ctxt: &SimplePreprocessorDriverCtxt<P>,
125      chapters: &mut Vec<(PathBuf, &'a mut String)>,
126      items: impl IntoIterator<Item = &'a mut BookItem>,
127    ) {
128      for item in items {
129        if let BookItem::Chapter(chapter) = item
130          && let Some(path) = &chapter.path
131        {
132          let chapter_path_abs = ctxt.src_dir.join(path);
133          let chapter_dir = chapter_path_abs.parent().unwrap().to_path_buf();
134          chapters.push((chapter_dir, &mut chapter.content));
135
136          for_each_mut(ctxt, chapters, &mut chapter.sub_items);
137        }
138      }
139    }
140
141    let mut chapters = Vec::new();
142    for_each_mut(&ctxt, &mut chapters, &mut book.items);
143
144    chapters
145      .into_par_iter()
146      .map(|(chapter_dir, content)| ctxt.process_chapter(&chapter_dir, content))
147      .collect::<Result<Vec<_>>>()?;
148
149    ctxt.sp.finish();
150
151    Ok(book)
152  }
153}