1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::{
  fs,
  marker::PhantomData,
  ops::Range,
  path::{Path, PathBuf},
};

use anyhow::Result;
use mdbook::{
  book::Book,
  preprocess::{Preprocessor, PreprocessorContext},
  BookItem,
};
use rayon::prelude::*;

#[derive(Copy, Clone)]
pub struct Asset {
  pub name: &'static str,
  pub contents: &'static [u8],
}

#[macro_export]
macro_rules! asset_generator {
  ($base:expr) => {
    macro_rules! make_asset {
      ($name:expr) => {
        $crate::Asset {
          name: $name,
          contents: include_bytes!(concat!($base, $name)),
        }
      };
    }
  };
}

pub trait SimplePreprocessor: Sized + Send + Sync {
  fn name() -> &'static str;
  fn build(ctx: &PreprocessorContext) -> Result<Self>;
  fn replacements(&self, chapter_dir: &Path, content: &str) -> Result<Vec<(Range<usize>, String)>>;
  fn linked_assets(&self) -> Vec<Asset>;
  fn all_assets(&self) -> Vec<Asset>;
  fn finish(self) {}
}

struct SimplePreprocessorDriverCtxt<P: SimplePreprocessor> {
  sp: P,
  src_dir: PathBuf,
}

impl<P: SimplePreprocessor> SimplePreprocessorDriverCtxt<P> {
  fn copy_assets(&self) -> Result<()> {
    // Rather than copying directly to the build directory, we instead copy to the book source
    // since mdBook will clean the build-dir after preprocessing. See mdBook#1087 for more.
    let dst_dir = self.src_dir.join(P::name());
    fs::create_dir_all(&dst_dir)?;

    for asset in self.sp.all_assets() {
      fs::write(dst_dir.join(asset.name), asset.contents)?;
    }

    Ok(())
  }

  fn process_chapter(&self, chapter_dir: &Path, content: &mut String) -> Result<()> {
    let replacements = self.sp.replacements(chapter_dir, content)?;
    if !replacements.is_empty() {
      for (range, html) in replacements.into_iter().rev() {
        content.replace_range(range, &html);
      }

      // If a chapter is located at foo/bar/the_chapter.md, then the generated source files
      // will be at foo/bar/the_chapter.html. So they need to reference preprocessor files
      // at ../../<preprocessor>/embed.js, i.e. we generate the right number of "..".
      let chapter_rel_path = chapter_dir.strip_prefix(&self.src_dir).unwrap();
      let depth = chapter_rel_path.components().count();
      let prefix = vec![".."; depth].into_iter().collect::<PathBuf>();

      // Ensure there's space between existing markdown and inserted HTML
      content.push_str("\n\n");

      for asset in self.sp.linked_assets() {
        let asset_rel = prefix.join(P::name()).join(asset.name);
        let asset_str = asset_rel.display().to_string();
        let link = match &*asset_rel.extension().unwrap().to_string_lossy() {
          "js" => format!(r#"<script type="text/javascript" src="{asset_str}"></script>"#),
          "mjs" => format!(r#"<script type="module" src="{asset_str}"></script>"#),
          "css" => format!(r#"<link rel="stylesheet" type="text/css" href="{asset_str}">"#),
          _ => continue,
        };
        content.push_str(&link);
      }
    }
    Ok(())
  }
}

pub(crate) struct SimplePreprocessorDriver<P: SimplePreprocessor>(PhantomData<P>);

impl<P: SimplePreprocessor> SimplePreprocessorDriver<P> {
  pub fn new() -> Self {
    SimplePreprocessorDriver(PhantomData)
  }
}

impl<P: SimplePreprocessor> Preprocessor for SimplePreprocessorDriver<P> {
  fn name(&self) -> &str {
    P::name()
  }

  fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
    let src_dir = ctx.root.join(&ctx.config.book.src);
    let sp = P::build(ctx)?;
    let ctxt = SimplePreprocessorDriverCtxt { sp, src_dir };
    ctxt.copy_assets()?;

    fn for_each_mut<'a, P: SimplePreprocessor>(
      ctxt: &SimplePreprocessorDriverCtxt<P>,
      chapters: &mut Vec<(PathBuf, &'a mut String)>,
      items: impl IntoIterator<Item = &'a mut BookItem>,
    ) {
      for item in items {
        if let BookItem::Chapter(chapter) = item {
          if chapter.path.is_some() {
            let chapter_path_abs = ctxt.src_dir.join(chapter.path.as_ref().unwrap());
            let chapter_dir = chapter_path_abs.parent().unwrap().to_path_buf();
            chapters.push((chapter_dir, &mut chapter.content));

            for_each_mut(ctxt, chapters, &mut chapter.sub_items);
          }
        }
      }
    }

    let mut chapters = Vec::new();
    for_each_mut(&ctxt, &mut chapters, &mut book.sections);

    chapters
      .into_par_iter()
      .map(|(chapter_dir, content)| ctxt.process_chapter(&chapter_dir, content))
      .collect::<Result<Vec<_>>>()?;

    ctxt.sp.finish();

    Ok(book)
  }
}