Include indent fixes

crates/mdbook-driver/src/builtin_preprocessors/links.rs

@@ -1,6 +1,5 @@
 use self::take_lines::{
-    take_anchored_lines, take_lines, take_rustdoc_include_anchored_lines,
-    take_rustdoc_include_lines,
+    take_anchored_lines, take_lines, take_rustdoc_anchored_lines, take_rustdoc_lines,
 };
 use anyhow::{Context, Result};
 use mdbook_core::book::{Book, BookItem};
@@ -91,32 +90,41 @@ where
     let path = path.as_ref();
     let source = source.as_ref();
     let mut previous_end_index = 0;
-    let mut replaced = String::new();
+    let mut replaced = String::with_capacity(s.len());
 
     for link in find_links(s) {
-        replaced.push_str(&s[previous_end_index..link.start_index]);
+        // Text from end of last link to start of current link
+        let text = &s[previous_end_index..link.start_index];
 
-        match link.render_with_path(path, chapter_title) {
+        // Text from start of line to start of current link.
+        let prefix = text.rfind('\n').map_or(text, |i| &text[i + 1..]);
+
+        // Stable equivalent of trim_suffix()
+        replaced.push_str(text.strip_suffix(prefix).unwrap_or(text));
+
+        match link.render_with_path(path, chapter_title, prefix) {
             Ok(new_content) => {
-                if depth < MAX_LINK_NESTED_DEPTH {
-                    if let Some(rel_path) = link.link_type.relative_path(path) {
-                        replaced.push_str(&replace_all(
-                            &new_content,
-                            rel_path,
-                            source,
-                            depth + 1,
-                            chapter_title,
-                        ));
-                    } else {
-                        replaced.push_str(&new_content);
-                    }
-                } else {
+                previous_end_index = link.end_index;
+
+                if depth >= MAX_LINK_NESTED_DEPTH {
                     error!(
                         "Stack depth exceeded in {}. Check for cyclic includes",
                         source.display()
                     );
+                    continue;
+                }
+
+                if let Some(rel_path) = link.link_type.relative_path(path) {
+                    replaced.push_str(&replace_all(
+                        &new_content,
+                        rel_path,
+                        source,
+                        depth + 1,
+                        chapter_title,
+                    ));
+                } else {
+                    replaced.push_str(&new_content);
                 }
-                previous_end_index = link.end_index;
             }
             Err(e) => {
                 error!("Error updating \"{}\", {}", link.link_text, e);
@@ -286,22 +294,23 @@ impl<'a> Link<'a> {
                 Some(LinkType::Title(title.as_str()))
             }
             (_, Some(typ), Some(rest)) => {
-                let mut path_props = rest.as_str().split_whitespace();
-                let file_arg = path_props.next();
-                let props: Vec<&str> = path_props.collect();
+                let mut props = rest.as_str().split_whitespace();
+                let file_arg = props.next();
 
                 match (typ.as_str(), file_arg) {
-                    ("include", Some(pth)) => Some(parse_include_path(pth)),
-                    ("playground", Some(pth)) => Some(LinkType::Playground(pth.into(), props)),
-                    ("playpen", Some(pth)) => {
+                    ("include", Some(path)) => Some(parse_include_path(path)),
+                    ("playground", Some(path)) => {
+                        Some(LinkType::Playground(path.into(), props.collect()))
+                    }
+                    ("playpen", Some(path)) => {
                         warn!(
                             "the {{{{#playpen}}}} expression has been \
                             renamed to {{{{#playground}}}}, \
                             please update your book to use the new name"
                         );
-                        Some(LinkType::Playground(pth.into(), props))
+                        Some(LinkType::Playground(path.into(), props.collect()))
                     }
-                    ("rustdoc_include", Some(pth)) => Some(parse_rustdoc_include_path(pth)),
+                    ("rustdoc_include", Some(path)) => Some(parse_rustdoc_include_path(path)),
                     _ => None,
                 }
             }
@@ -325,68 +334,120 @@ impl<'a> Link<'a> {
         &self,
         base: P,
         chapter_title: &mut String,
+        prefix: &str,
     ) -> Result<String> {
+        use std::fmt::Write;
         let base = base.as_ref();
+        let mut out = String::new();
         match self.link_type {
             // omit the escape char
-            LinkType::Escaped => Ok(self.link_text[1..].to_owned()),
-            LinkType::Include(ref pat, ref range_or_anchor) => {
-                let target = base.join(pat);
-
-                fs::read_to_string(&target)
-                    .map(|s| match range_or_anchor {
-                        RangeOrAnchor::Range(range) => take_lines(&s, range.clone()),
-                        RangeOrAnchor::Anchor(anchor) => take_anchored_lines(&s, anchor),
-                    })
-                    .with_context(|| {
-                        format!(
-                            "Could not read file for link {} ({})",
-                            self.link_text,
-                            target.display(),
-                        )
-                    })
+            LinkType::Escaped => {
+                write!(out, "{prefix}{}", &self.link_text[1..]).expect("String writes don't fail");
+                Ok(out)
+            }
+
+            LinkType::Include(ref path, ref range_or_anchor) => {
+                use RangeOrAnchor::*;
+
+                let target = base.join(path);
+
+                let contents = fs::read_to_string(&target).with_context(|| {
+                    format!(
+                        "Could not read file for link {} ({})",
+                        self.link_text,
+                        target.display(),
+                    )
+                })?;
+
+                let lines: Vec<_> = match range_or_anchor {
+                    Range(range) => take_lines(&contents, range.clone()).collect(),
+                    Anchor(anchor) => take_anchored_lines(&contents, anchor).collect(),
+                };
+
+                // Count shared leading spaces
+                let trim = lines
+                    .iter()
+                    .map(|line| line.bytes().take_while(|&b| b == b' ').count())
+                    .fold(usize::MAX, std::cmp::min);
+
+                for line in lines {
+                    write!(out, "{prefix}{}\n", &line[trim..]).expect("String writes don't fail");
+                }
+
+                // Trim trailing new line
+                out.pop();
+                Ok(out)
             }
-            LinkType::RustdocInclude(ref pat, ref range_or_anchor) => {
-                let target = base.join(pat);
 
-                fs::read_to_string(&target)
-                    .map(|s| match range_or_anchor {
-                        RangeOrAnchor::Range(range) => {
-                            take_rustdoc_include_lines(&s, range.clone())
+            LinkType::RustdocInclude(ref path, ref range_or_anchor) => {
+                let target = base.join(path);
+
+                let contents = fs::read_to_string(&target).with_context(|| {
+                    format!(
+                        "Could not read file for link {} ({})",
+                        self.link_text,
+                        target.display(),
+                    )
+                })?;
+
+                match range_or_anchor {
+                    RangeOrAnchor::Range(range) => {
+                        for (line, show) in take_rustdoc_lines(&contents, range.clone()) {
+                            write!(
+                                out,
+                                "{prefix}{}{line}\n",
+                                show.then_some("").unwrap_or("# ")
+                            )
+                            .expect("String writes don't fail");
                         }
-                        RangeOrAnchor::Anchor(anchor) => {
-                            take_rustdoc_include_anchored_lines(&s, anchor)
+                    }
+                    RangeOrAnchor::Anchor(anchor) => {
+                        for (line, show) in take_rustdoc_anchored_lines(&contents, anchor) {
+                            write!(
+                                out,
+                                "{prefix}{}{line}\n",
+                                show.then_some("").unwrap_or("# ")
+                            )
+                            .expect("String writes don't fail");
                         }
-                    })
-                    .with_context(|| {
-                        format!(
-                            "Could not read file for link {} ({})",
-                            self.link_text,
-                            target.display(),
-                        )
-                    })
+                    }
+                }
+
+                // Trim trailing new line
+                out.pop();
+                Ok(out)
             }
-            LinkType::Playground(ref pat, ref attrs) => {
-                let target = base.join(pat);
 
-                let mut contents = fs::read_to_string(&target).with_context(|| {
+            LinkType::Playground(ref path, ref attrs) => {
+                let target = base.join(path);
+
+                let contents = fs::read_to_string(&target).with_context(|| {
                     format!(
                         "Could not read file for link {} ({})",
                         self.link_text,
                         target.display()
                     )
                 })?;
-                let ftype = if !attrs.is_empty() { "rust," } else { "rust" };
-                if !contents.ends_with('\n') {
-                    contents.push('\n');
+
+                let mut out = String::with_capacity(contents.len() + 11);
+
+                out.push_str(prefix);
+                out.push_str("```rust");
+                for s in attrs {
+                    out.push_str(",");
+                    out.push_str(s);
                 }
-                Ok(format!(
-                    "```{}{}\n{}```\n",
-                    ftype,
-                    attrs.join(","),
-                    contents
-                ))
+                out.push_str("\n");
+
+                for line in contents.lines() {
+                    write!(out, "{prefix}{line}\n").expect("String writes don't fail");
+                }
+
+                out.push_str("```");
+
+                Ok(out)
             }
+
             LinkType::Title(title) => {
                 *chapter_title = title.to_owned();
                 Ok(String::new())