From 9b2ab45eafebaf474f2265b1d77ae38adce2d4bb Mon Sep 17 00:00:00 2001
From: Weihang Lo <weihanglotw@gmail.com>
Date: Mon, 30 Apr 2018 23:47:38 +0800
Subject: [PATCH] Add `rewrite_to_dir` config to rewrite index file path

Example: `/path/to/index.html` would be rewritten to `/path/to`
---
 src/config.rs                                |  2 +
 src/renderer/html_handlebars/hbs_renderer.rs |  5 +-
 src/renderer/html_handlebars/helpers/toc.rs  | 48 ++++++++++++++++++--
 3 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 898ec6b7..291c8a63 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -424,6 +424,8 @@ pub struct HtmlConfig {
     pub livereload_url: Option<String>,
     /// Should section labels be rendered?
     pub no_section_label: bool,
+    /// Filenames listed here would be rewritten to directory index `/`.
+    pub rewrite_to_dir: Vec<String>,
     /// Search settings. If `None`, the default will be used.
     pub search: Option<Search>,
 }
diff --git a/src/renderer/html_handlebars/hbs_renderer.rs b/src/renderer/html_handlebars/hbs_renderer.rs
index fbaad523..b4d73854 100644
--- a/src/renderer/html_handlebars/hbs_renderer.rs
+++ b/src/renderer/html_handlebars/hbs_renderer.rs
@@ -219,7 +219,10 @@ impl HtmlHandlebars {
     }
 
     fn register_hbs_helpers(&self, handlebars: &mut Handlebars, html_config: &HtmlConfig) {
-        handlebars.register_helper("toc", Box::new(helpers::toc::RenderToc {no_section_label: html_config.no_section_label}));
+        handlebars.register_helper("toc", Box::new(helpers::toc::RenderToc {
+            no_section_label: html_config.no_section_label,
+            rewrite_to_dir: html_config.rewrite_to_dir.to_owned(),
+        }));
         handlebars.register_helper("previous", Box::new(helpers::navigation::previous));
         handlebars.register_helper("next", Box::new(helpers::navigation::next));
     }
diff --git a/src/renderer/html_handlebars/helpers/toc.rs b/src/renderer/html_handlebars/helpers/toc.rs
index f44ea2ac..2ade3e21 100644
--- a/src/renderer/html_handlebars/helpers/toc.rs
+++ b/src/renderer/html_handlebars/helpers/toc.rs
@@ -1,4 +1,4 @@
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::collections::BTreeMap;
 
 use serde_json;
@@ -6,9 +6,9 @@ use handlebars::{Handlebars, Helper, HelperDef, RenderContext, RenderError};
 use pulldown_cmark::{html, Event, Parser, Tag};
 
 // Handlebars helper to construct TOC
-#[derive(Clone, Copy)]
 pub struct RenderToc {
     pub no_section_label: bool,
+    pub rewrite_to_dir: Vec<String>,
 }
 
 impl HelperDef for RenderToc {
@@ -69,8 +69,11 @@ impl HelperDef for RenderToc {
                 if !path.is_empty() {
                     rc.writer.write_all(b"<a href=\"")?;
 
-                    let tmp = Path::new(item.get("path").expect("Error: path should be Some(_)"))
-                        .with_extension("html")
+                    let tmp = {
+                        // To be recognized by browsers, rewrite extenstion to `.html`.
+                        let path = Path::new(path).with_extension("html");
+                        self.rewrite_directory_index(&path)
+                    }
                         .to_str()
                         .unwrap()
                         // Hack for windows who tends to use `\` as separator instead of `/`
@@ -137,4 +140,41 @@ impl HelperDef for RenderToc {
         rc.writer.write_all(b"</ol>")?;
         Ok(())
     }
+
+}
+
+impl RenderToc {
+    // Rewrite filenames matches any in `rewrite_to_dir` to directory index.
+    fn rewrite_directory_index(&self, path: &Path) -> PathBuf {
+        for filename in self.rewrite_to_dir.iter() {
+            if filename.as_str() == path.file_name().unwrap_or_default() {
+                return path.with_file_name("");
+            }
+        }
+        return path.to_owned();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn rewrite_dir_success() {
+        let render = RenderToc {
+            no_section_label: true,
+            rewrite_to_dir: vec![
+                "index.html".to_owned(),
+                "index.md".to_owned(),
+            ],
+        };
+        let path = PathBuf::from("index.html");
+        assert_eq!(render.rewrite_directory_index(&path), PathBuf::from(""));
+
+        let path = PathBuf::from("index.md");
+        assert_eq!(render.rewrite_directory_index(&path), PathBuf::from(""));
+
+        let path = PathBuf::from("index.asp");
+        assert_eq!(render.rewrite_directory_index(&path), path);
+    }
 }