[WIP] Initial try at LOL HTML rewriter

jyn514 · Aug 3, 2020 · b3ab913 · b3ab913
1 parent c3cfbc8
commit b3ab913
Show file tree

Hide file tree

Showing 5 changed files with 199 additions and 73 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -27,7 +27,6 @@ backtrace = "0.3"
 failure = { version = "0.1.3", features = ["backtrace"] }
 comrak = { version = "0.8", default-features = false }
 toml = "0.5"
-kuchiki = "0.8"
 schemamama = "0.3"
 schemamama_postgres = "0.3"
 systemstat = "0.1.4"
@@ -41,6 +40,7 @@ path-slash = "0.1.3"
 once_cell = { version = "1.4.0", features = ["parking_lot"] }
 base64 = "0.12.1"
 strum = { version = "0.18.0", features = ["derive"] }
+lol_html = "0.2"
 
 # Async
 tokio = { version = "0.2.22", features = ["rt-threaded"] }
@@ -81,6 +81,7 @@ procfs = "0.7"
 
 [dev-dependencies]
 criterion = "0.3"
+kuchiki = "0.8"
 rand = "0.7.3"
 
 [[bench]]

diff --git a/src/utils/html.rs b/src/utils/html.rs
@@ -1,40 +1,64 @@
-use crate::error::Result;
-use failure::err_msg;
-use kuchiki::traits::TendrilSink;
-use kuchiki::NodeRef;
-
-/// Extracts the contents of the `<head>` and `<body>` tags from an HTML document, as well as the
-/// classes on the `<body>` tag, if any.
-pub fn extract_head_and_body(html: &str) -> Result<(String, String, String)> {
-    let dom = kuchiki::parse_html().one(html);
-
-    let head = dom
-        .select_first("head")
-        .map_err(|_| err_msg("couldn't find <head> tag in rustdoc output"))?;
-    let body = dom
-        .select_first("body")
-        .map_err(|_| err_msg("couldn't find <body> tag in rustdoc output"))?;
-
-    let class = body
-        .attributes
-        .borrow()
-        .get("class")
-        .map(|v| v.to_owned())
-        .unwrap_or_default();
-
-    Ok((serialize(head.as_node()), serialize(body.as_node()), class))
-}
+use crate::web::page::TemplateData;
+use lol_html::errors::RewritingError;
+use tera::Context;
 
-fn serialize(v: &NodeRef) -> String {
-    let mut contents = Vec::new();
-    for child in v.children() {
-        child
-            .serialize(&mut contents)
-            .expect("serialization failed");
-    }
-    String::from_utf8(contents).expect("non utf-8 html")
+pub(crate) fn rewrite_lol(
+    html: &str,
+    ctx: Context,
+    templates: &TemplateData,
+) -> Result<String, RewritingError> {
+    use lol_html::html_content::{ContentType, Element};
+    use lol_html::{ElementContentHandlers, RewriteStrSettings};
+
+    let templates = templates.templates.load();
+    let tera_head = templates.render("rustdoc/head.html", &ctx).unwrap();
+    let tera_body = templates.render("rustdoc/body.html", &ctx).unwrap();
+
+    let head_handler = |head: &mut Element| {
+        head.append(&tera_head, ContentType::Html);
+        Ok(())
+    };
+    // Before: <body> ... rustdoc content ... </body>
+    // After:
+    // ```html
+    // <div id="rustdoc_body_wrapper" class="{{ rustdoc_body_class }}" tabindex="-1">
+    //      ... rustdoc content ...
+    // </div>
+    // ```
+    let body_handler = |rustdoc_body_class: &mut Element| {
+        // Add the `rustdoc` classes to the html body
+        rustdoc_body_class.set_attribute("container-rustdoc", "")?;
+        rustdoc_body_class.set_attribute("id", "rustdoc_body_wrapper")?;
+        rustdoc_body_class.set_attribute("tabindex", "-1")?;
+        // Change the `body` to a `div`
+        rustdoc_body_class.set_tag_name("div")?;
+        // Prepend the tera content
+        rustdoc_body_class.prepend(&tera_body, ContentType::Html);
+        // Now, make this a full <body> tag
+        rustdoc_body_class.before("<body>", ContentType::Html);
+        rustdoc_body_class.after("</body>", ContentType::Html);
+
+        Ok(())
+    };
+
+    let (head_selector, body_selector) = ("head".parse().unwrap(), "body".parse().unwrap());
+    let head = (
+        &head_selector,
+        ElementContentHandlers::default().element(head_handler),
+    );
+    let body = (
+        &body_selector,
+        ElementContentHandlers::default().element(body_handler),
+    );
+    let settings = RewriteStrSettings {
+        element_content_handlers: vec![head, body],
+        ..RewriteStrSettings::default()
+    };
+
+    lol_html::rewrite_str(html, settings)
 }
 
+/*
 #[cfg(test)]
 mod test {
     #[test]
@@ -60,3 +84,4 @@ mod test {
         assert_eq!(class, "rustdoc struct");
     }
 }
+*/
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
@@ -4,7 +4,7 @@ pub(crate) use self::cargo_metadata::{CargoMetadata, Package as MetadataPackage}
 pub(crate) use self::copy::copy_doc_dir;
 pub use self::daemon::start_daemon;
 pub use self::github_updater::GithubUpdater;
-pub use self::html::extract_head_and_body;
+pub(crate) use self::html::rewrite_lol;
 pub use self::queue::{get_crate_priority, remove_crate_priority, set_crate_priority};
 pub use self::queue_builder::queue_builder;
 pub use self::release_activity_updater::update_release_activity;