From 6cbcd7b25ca9a1dae6d69b2989d2a3144e7e2e4f Mon Sep 17 00:00:00 2001
From: Raymond Li <hi@raymond.li>
Date: Tue, 19 Nov 2024 23:44:40 -0800
Subject: [PATCH] Add JS into every HTML page to preserve query params (#79)

Preserve query params (`path/to/index.html-->?foo=bar&baz<--`) for all
relative links on all HTML pages

Skips:
- http(s?):// links
- \# links to parts of current page
- empty links

Testing:
<img width="535" alt="image"
src="https://github.com/user-attachments/assets/0212fc8a-96bc-447a-97bb-497d0b7cbc7f">
<img width="697" alt="image"
src="https://github.com/user-attachments/assets/efd5cc7d-bbf4-41b4-b41e-786d09d6a915">
---
 src/lib.rs       |  2 ++
 src/parsers.rs   | 13 +++++++++++--
 src/templates.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 src/types.rs     |  6 ++++++
 4 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index fef375e..94017d1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -244,6 +244,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
     let mut breaks = RestartsAndFailuresContext {
         css: TEMPLATE_FAILURES_CSS,
         failures: Vec::new(),
+        qps: TEMPLATE_QUERY_PARAM_SCRIPT,
     };
 
     // NB: Sometimes, the log output we get from Logarithm stutters with a blank line.
@@ -515,6 +516,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
         has_unknown_stack_trie: !unknown_stack_trie.is_empty(),
         num_breaks: breaks.failures.len(),
         has_chromium_events: !chromium_events.is_empty(),
+        qps: TEMPLATE_QUERY_PARAM_SCRIPT,
     };
     output.push((
         PathBuf::from("index.html"),
diff --git a/src/parsers.rs b/src/parsers.rs
index a9e582c..cae918f 100644
--- a/src/parsers.rs
+++ b/src/parsers.rs
@@ -1,3 +1,4 @@
+use crate::templates::TEMPLATE_QUERY_PARAM_SCRIPT;
 use crate::{types::*, ParseConfig};
 use html_escape::encode_text;
 use std::cell::RefCell;
@@ -186,7 +187,10 @@ impl StructuredLogParser for DynamoGuardParser<'_> {
     ) -> anyhow::Result<ParserResults> {
         let filename = format!("{}.html", self.name());
         let guards = serde_json::from_str::<Vec<DynamoGuard>>(payload)?;
-        let guards_context = DynamoGuardsContext { guards };
+        let guards_context = DynamoGuardsContext {
+            guards,
+            qps: TEMPLATE_QUERY_PARAM_SCRIPT,
+        };
         let output = self.tt.render(&filename, &guards_context)?;
         simple_file_output(&filename, lineno, compile_id, &output)
     }
@@ -435,6 +439,7 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                 symbolic_shape_specializations: specializations,
                 output_files: &output_files,
                 compile_id_dir: &self.compile_id_dir,
+                qps: TEMPLATE_QUERY_PARAM_SCRIPT,
             };
             let output = self.tt.render(&filename, &context)?;
             simple_file_output(&filename, lineno, compile_id, &output)
@@ -473,6 +478,7 @@ impl StructuredLogParser for AOTAutogradBackwardCompilationMetricsParser<'_> {
                 css: crate::CSS,
                 m: &m,
                 compile_id: id,
+                qps: TEMPLATE_QUERY_PARAM_SCRIPT,
             };
             let output = self.tt.render(&filename, &context)?;
             simple_file_output(&filename, lineno, compile_id, &output)
@@ -513,6 +519,7 @@ impl StructuredLogParser for BwdCompilationMetricsParser<'_> {
                 css: crate::CSS,
                 m: &m,
                 compile_id: id,
+                qps: TEMPLATE_QUERY_PARAM_SCRIPT,
             };
             let output = self.tt.render(&filename, &context)?;
             simple_file_output(&filename, lineno, compile_id, &output)
@@ -599,7 +606,9 @@ pub fn anchor_source(text: &str) -> String {
         ));
     }
 
-    html.push_str("</pre></body></html>");
+    html.push_str(&format!(
+        "</pre>{TEMPLATE_QUERY_PARAM_SCRIPT}</body></html>"
+    ));
     html
 }
 
diff --git a/src/templates.rs b/src/templates.rs
index 686ef7e..62bc214 100644
--- a/src/templates.rs
+++ b/src/templates.rs
@@ -44,6 +44,7 @@ pub static TEMPLATE_DYNAMO_GUARDS: &str = r#"
     <li><code>{guard.code}</code></li>
 {{ endfor }}
 </ul>
+{qps | format_unescaped}
 </body>
 </html>
 "#;
@@ -160,6 +161,7 @@ Build products below:
 {unknown_stack_trie_html | format_unescaped}
 </div>
 {{ endif }}
+{qps | format_unescaped}
 </body>
 </html>
 "#;
@@ -205,6 +207,7 @@ pub static TEMPLATE_FAILURES_AND_RESTARTS: &str = r#"
     {{ for failure in failures }}
     <tr> <td> {failure.0 | format_unescaped} </td>{failure.1 | format_unescaped}</tr>
     {{ endfor }}
+    {qps | format_unescaped}
 </body>
 </html>
 "#;
@@ -290,6 +293,7 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"
     </tr>
     {{ endfor }}
     </table>
+    {qps | format_unescaped}
 </body>
 </html>
 "#;
@@ -311,6 +315,7 @@ pub static TEMPLATE_AOT_AUTOGRAD_BACKWARD_COMPILATION_METRICS: &str = r#"
     {{ else }}
     <p> No failures! </p>
     {{ endif }}
+    {qps | format_unescaped}
 </body>
 </html>
 "#;
@@ -339,6 +344,47 @@ pub static TEMPLATE_BWD_COMPILATION_METRICS: &str = r#"
     {{ else }}
     <p> No failures! </p>
     {{ endif }}
+    {qps | format_unescaped}
 </body>
 </html>
 "#;
+
+// NB: Invariant for generated HTML: all links must show up in the initial HTML for this to be applied.
+//     Links dynamically generated/added after document load (i.e. using JS) will not get this applied.
+pub static TEMPLATE_QUERY_PARAM_SCRIPT: &str = r#"
+    <script>
+    document.addEventListener('DOMContentLoaded', function() {
+
+        // Append the current URL's query parameters to all relative links on the page
+        const queryParams = new URLSearchParams(window.location.search);
+        if (queryParams.size === 0) return url; // No query params, return original URL
+
+        function appendQueryParams(url) {
+            const parsedUrl = url;
+            const noOrigParams = parsedUrl.lastIndexOf('?') == -1;
+            const newSearchParams = new URLSearchParams();
+
+            // Append query parameters
+            for (const [key, value] of queryParams) {
+                newSearchParams.set(key, value);
+            }
+
+            let updatedURL = url;
+            if (noOrigParams) { // Original URL has no params
+                return parsedUrl + "?" + newSearchParams.toString();
+            }
+
+            // Original URL has params - append to end
+            return parsedUrl + "&" + newSearchParams.toString();
+        }
+
+        // Select all relative links on the page
+        const relativeLinks = document.querySelectorAll('a[href]:not([href^="http://"]):not([href^="https://"]):not([href^="\#"])');
+
+        // Append query parameters to each relative link
+        relativeLinks.forEach((link) => {
+            link.setAttribute("href", appendQueryParams(link.getAttribute("href")))
+        });
+    });
+    </script>
+    "#;
diff --git a/src/types.rs b/src/types.rs
index 5b476da..8baa23b 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -335,6 +335,7 @@ pub struct BwdCompilationMetricsContext<'e> {
     pub m: &'e BwdCompilationMetricsMetadata,
     pub css: &'static str,
     pub compile_id: String,
+    pub qps: &'static str,
 }
 
 #[derive(Debug, Serialize)]
@@ -342,6 +343,7 @@ pub struct AOTAutogradBackwardCompilationMetricsContext<'e> {
     pub m: &'e AOTAutogradBackwardCompilationMetricsMetadata,
     pub css: &'static str,
     pub compile_id: String,
+    pub qps: &'static str,
 }
 
 #[derive(Clone, Debug, Serialize)]
@@ -362,6 +364,7 @@ pub struct CompilationMetricsContext<'e> {
     pub output_files: &'e Vec<OutputFile>,
     pub compile_id_dir: &'e PathBuf,
     pub mini_stack_html: String,
+    pub qps: &'static str,
 }
 
 #[derive(Debug, Serialize)]
@@ -402,6 +405,7 @@ pub struct RestartsAndFailuresContext {
     // Serialized versions of (CompileId, FailureReason)
     pub failures: Vec<(String, String)>,
     pub css: &'static str,
+    pub qps: &'static str,
 }
 
 #[derive(Debug)]
@@ -560,6 +564,7 @@ pub struct DynamoGuard {
 #[derive(Debug, Serialize)]
 pub struct DynamoGuardsContext {
     pub guards: Vec<DynamoGuard>,
+    pub qps: &'static str,
 }
 
 #[derive(Debug, Serialize)]
@@ -573,6 +578,7 @@ pub struct IndexContext {
     pub num_breaks: usize,
     pub custom_header_html: String,
     pub has_chromium_events: bool,
+    pub qps: &'static str,
 }
 
 #[derive(Debug, Serialize)]