From b27c7a34c39ce1d3f4c475517ab30f549b39bb28 Mon Sep 17 00:00:00 2001
From: IvanKobzarev <ivan.kobzarev@gmail.com>
Date: Wed, 8 Jan 2025 17:34:11 +0100
Subject: [PATCH] guard_added_fast rendering in compilation metrics

---
 src/lib.rs       |  9 +++++++++
 src/parsers.rs   | 14 ++++++++++++++
 src/templates.rs | 13 +++++++++++++
 src/types.rs     | 24 ++++++++++++++++++++++++
 4 files changed, 60 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 4293ced..e584d86 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -218,6 +218,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
 
     let symbolic_shape_specialization_index: RefCell<SymbolicShapeSpecializationIndex> =
         RefCell::new(FxHashMap::default());
+    let guard_added_fast_index: RefCell<GuardAddedFastIndex> = RefCell::new(FxHashMap::default());
 
     // Store results in an output Vec<PathBuf, String>
     let mut output: Vec<(PathBuf, String)> = Vec::new();
@@ -385,6 +386,7 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
                     tt: &tt,
                     stack_index: &stack_index,
                     symbolic_shape_specialization_index: &symbolic_shape_specialization_index,
+                    guard_added_fast_index: &guard_added_fast_index,
                     output_files: &copied_directory,
                     compile_id_dir: &compile_id_dir,
                 });
@@ -466,6 +468,13 @@ pub fn parse_path(path: &PathBuf, config: ParseConfig) -> anyhow::Result<ParseOu
                 .or_default()
                 .push(specialization);
         }
+        if let Some(guard_added_fast) = e.guard_added_fast {
+            guard_added_fast_index
+                .borrow_mut()
+                .entry(e.compile_id.clone())
+                .or_default()
+                .push(guard_added_fast)
+        }
 
         if let Some(m) = e.dynamo_start {
             if let Some(mut stack) = m.stack {
diff --git a/src/parsers.rs b/src/parsers.rs
index f3cdda2..98a1397 100644
--- a/src/parsers.rs
+++ b/src/parsers.rs
@@ -346,6 +346,7 @@ pub struct CompilationMetricsParser<'t> {
     pub tt: &'t TinyTemplate<'t>,
     pub stack_index: &'t RefCell<StackIndex>,
     pub symbolic_shape_specialization_index: &'t RefCell<SymbolicShapeSpecializationIndex>,
+    pub guard_added_fast_index: &'t RefCell<GuardAddedFastIndex>,
     pub output_files: &'t Vec<OutputFile>,
     pub compile_id_dir: &'t PathBuf,
 }
@@ -409,6 +410,18 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                     stack_html: format_stack(&spec.stack.unwrap_or(Vec::new())),
                 })
                 .collect();
+            let guards_added_fast = self
+                .guard_added_fast_index
+                .borrow_mut()
+                .remove(&cid)
+                .unwrap_or(Vec::new())
+                .drain(..)
+                .map(|guard| GuardAddedFastContext {
+                    expr: guard.expr.unwrap_or("".to_string()),
+                    user_stack_html: format_stack(&guard.user_stack.unwrap_or(Vec::new())),
+                    stack_html: format_stack(&guard.stack.unwrap_or(Vec::new())),
+                })
+                .collect();
             let remove_prefix = |x: &String| -> String {
                 // url is X_Y_Z/<rest>. Get the rest of the string for the link
                 // on compilation metrics page
@@ -433,6 +446,7 @@ impl StructuredLogParser for CompilationMetricsParser<'_> {
                 stack_html: stack_html,
                 mini_stack_html: mini_stack_html,
                 symbolic_shape_specializations: specializations,
+                guards_added_fast: guards_added_fast,
                 output_files: &output_files,
                 compile_id_dir: &self.compile_id_dir,
                 qps: TEMPLATE_QUERY_PARAM_SCRIPT,
diff --git a/src/templates.rs b/src/templates.rs
index 5c1881e..40f0dd8 100644
--- a/src/templates.rs
+++ b/src/templates.rs
@@ -302,6 +302,19 @@ pub static TEMPLATE_COMPILATION_METRICS: &str = r#"
     </tr>
     {{ endfor }}
     </table>
+    <h2>Guards added fast</h2>
+    <table>
+    <tr>
+        <th>Expr</th> <th>User stack</th> <th>Framework stack</th>
+    </tr>
+    {{ for g in guards_added_fast }}
+    <tr>
+        <td>{g.expr}</td>
+        <td>{g.user_stack_html | format_unescaped}</td>
+        <td>{g.stack_html | format_unescaped}</td>
+    </tr>
+    {{ endfor }}
+    </table>
     {qps | format_unescaped}
 </body>
 </html>
diff --git a/src/types.rs b/src/types.rs
index fe02bb4..bdfe93f 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -18,6 +18,8 @@ pub type CompilationMetricsIndex = FxIndexMap<Option<CompileId>, Vec<Compilation
 pub type StackIndex = FxHashMap<Option<CompileId>, StackSummary>; // NB: attempt is always 0 here
 pub type SymbolicShapeSpecializationIndex =
     FxHashMap<Option<CompileId>, Vec<SymbolicShapeSpecializationMetadata>>;
+pub type GuardAddedFastIndex =
+    FxHashMap<Option<CompileId>, Vec<GuardAddedFastMetadata>>;
 
 pub type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<FxHasher>>;
 
@@ -390,12 +392,18 @@ pub struct CompilationMetricsContext<'e> {
     pub compile_id: String,
     pub stack_html: String,
     pub symbolic_shape_specializations: Vec<SymbolicShapeSpecializationContext>,
+    pub guards_added_fast: Vec<GuardAddedFastContext>,
     pub output_files: &'e Vec<OutputFile>,
     pub compile_id_dir: &'e PathBuf,
     pub mini_stack_html: String,
     pub qps: &'static str,
 }
 
+#[derive(Debug, Serialize)]
+pub struct GuardsAddedFastContext {
+    pub guards: Vec<GuardAddedFastContext>,
+}
+
 #[derive(Debug, Serialize)]
 pub enum FailureReason {
     Failure((String, String, String, u32)), // (failure type, failure reason, user frame filename, user frame lineno)
@@ -452,6 +460,7 @@ pub enum Metadata<'e> {
     BwdCompilationMetrics(&'e BwdCompilationMetricsMetadata),
     Artifact(&'e ArtifactMetadata),
     DumpFile(&'e DumpFileMetadata),
+    GuardAddedFast(&'e GuardAddedFastMetadata),
 }
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -459,6 +468,13 @@ pub struct DumpFileMetadata {
     pub name: String,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+pub struct GuardAddedFastMetadata {
+    pub expr: Option<String>,
+    pub stack: Option<StackSummary>,
+    pub user_stack: Option<StackSummary>,
+}
+
 #[derive(Debug, Deserialize)]
 pub struct Envelope {
     pub rank: Option<u32>,
@@ -496,6 +512,7 @@ pub struct Envelope {
     pub describe_source: Option<SourceDesc>,
     pub dump_file: Option<DumpFileMetadata>,
     pub chromium_event: Option<EmptyMetadata>,
+    pub guard_added_fast: Option<GuardAddedFastMetadata>,
     #[serde(flatten)]
     pub _other: FxHashMap<String, Value>,
 }
@@ -618,3 +635,10 @@ pub struct SymbolicShapeSpecializationContext {
     pub user_stack_html: String,
     pub stack_html: String,
 }
+
+#[derive(Debug, Serialize)]
+pub struct GuardAddedFastContext {
+    pub expr: String,
+    pub user_stack_html: String,
+    pub stack_html: String,
+}