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_preprocessor::{
10 Preprocessor, PreprocessorContext,
11 book::{Book, BookItem},
12};
13use 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 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 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 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}