Improved the viewer

Cargo.toml

@@ -34,6 +34,7 @@ object = { version = "0.32", features = ["wasm"] }
 gimli = "0.28"
 tempfile = "3.10"
 pretty_assertions = "1.4"
+base64 = "0.22"
 
 [profile.release]
 opt-level = 3

bin/stylus-trace-studio/src/main.rs

@@ -356,7 +356,7 @@ fn handle_view(tx_or_path: &str, rpc: &str) -> Result<()> {
         info!("Opening existing profile: {}", path.display());
         let profile = read_profile(&path).context("Failed to read profile JSON")?;
         let viewer_path = path.with_extension("html");
-        generate_viewer(&profile, &viewer_path)?;
+        generate_viewer(&profile, None, &viewer_path)?;
         open_browser(&viewer_path)?;
     } else if tx_or_path.starts_with("0x") && tx_or_path.len() == 66 {
         info!("Capturing and viewing transaction: {}", tx_or_path);

crates/stylus-trace-core/Cargo.toml

@@ -33,3 +33,4 @@ addr2line.workspace = true
 object.workspace = true
 gimli.workspace = true
 tempfile.workspace = true
+base64.workspace = true

crates/stylus-trace-core/src/aggregator/metrics.rs

@@ -4,7 +4,7 @@
 //! These are the primary targets for optimization.
 
 use super::stack_builder::CollapsedStack;
-use crate::parser::schema::HotPath;
+use crate::parser::schema::{GasCategory, HotPath};
 use log::debug;
 
 /// Calculate hot paths from collapsed stacks
@@ -49,10 +49,16 @@ pub fn create_hot_path(stack: &CollapsedStack, denominator: u64) -> HotPath {
         0.0
     };
 
+    // Determine the category from the leaf node of the stack.
+    // E.g., for "root;call;storage_flush_cache" the leaf is "storage_flush_cache".
+    let leaf = stack.stack.split(';').next_back().unwrap_or(&stack.stack);
+    let category = categorize_stack_leaf(leaf);
+
     HotPath {
         stack: stack.stack.clone(),
         gas: stack.weight,
         percentage,
+        category,
         source_hint: stack.last_pc.map(|pc| crate::parser::schema::SourceHint {
             file: "unknown".to_string(),
             line: None,
@@ -62,6 +68,45 @@ pub fn create_hot_path(stack: &CollapsedStack, denominator: u64) -> HotPath {
     }
 }
 
+/// Categorize a hot-path leaf node into a [`GasCategory`].
+///
+/// This mirrors (and supersedes) the frontend `getCategory()` heuristic in
+/// `viewer.js`, but is authoritative because it lives in the Rust backend
+/// where `HostIoType` semantics are well-defined.
+pub fn categorize_stack_leaf(leaf: &str) -> GasCategory {
+    let n = leaf.to_lowercase();
+
+    if n == "root" {
+        return GasCategory::Root;
+    }
+    // Expensive storage writes
+    if n.contains("storage_store") || n.contains("storage_flush") {
+        return GasCategory::StorageExpensive;
+    }
+    // Cheaper storage reads
+    if n.contains("storage_load") || n.contains("storage_cache") {
+        return GasCategory::StorageNormal;
+    }
+    // Cryptographic ops
+    if n.contains("keccak") {
+        return GasCategory::Crypto;
+    }
+    // Memory / ABI
+    if n.contains("memory") || n == "read_args" || n == "write_result" {
+        return GasCategory::Memory;
+    }
+    // External calls / creates
+    if n == "call" || n == "staticcall" || n == "delegatecall" || n.contains("create") {
+        return GasCategory::Call;
+    }
+    // System / context
+    if n.contains("msg_") || n.contains("block_") || n == "account_balance" {
+        return GasCategory::System;
+    }
+    // Catch-all: user code
+    GasCategory::UserCode
+}
+
 /// Calculate gas distribution statistics
 ///
 /// **Public** - provides summary statistics

crates/stylus-trace-core/src/commands/capture.rs

@@ -191,7 +191,11 @@ pub fn execute_capture(args: CaptureArgs) -> Result<()> {
             Some(stacks.to_vec()),
             mapper.as_ref(),
         );
-        crate::output::viewer::generate_viewer(&profile, &viewer_path)?;
+        // Generate SVG for the flamegraph tab in the viewer.
+        // We attempt this even if --output-svg was not requested; failure is non-fatal.
+        let viewer_svg =
+            generate_flamegraph(&stacks, args.flamegraph_config.as_ref(), mapper.as_ref()).ok();
+        crate::output::viewer::generate_viewer(&profile, viewer_svg.as_deref(), &viewer_path)?;
         info!("✓ Viewer generated at: {}", viewer_path.display());
         crate::output::viewer::open_browser(&viewer_path)?;
     }

crates/stylus-trace-core/src/commands/diff.rs

@@ -133,10 +133,20 @@ pub fn execute_diff(args: DiffArgs) -> Result<()> {
             .with_extension("html");
 
         let report_json = serde_json::to_value(&report)?;
+
+        // Attempt to generate diff flamegraph SVG for the multi-view tab.
+        // If full stacks are not available (older capture), viewer still works without it.
+        let diff_svg = baseline
+            .all_stacks
+            .as_ref()
+            .zip(target.all_stacks.as_ref())
+            .and_then(|(b, t)| crate::flamegraph::generate_diff_flamegraph(b, t, None).ok());
+
         crate::output::viewer::generate_diff_viewer(
             &baseline,
             &target,
             &report_json,
+            diff_svg.as_deref(),
             &viewer_path,
         )?;
         info!("✓ Diff viewer generated at: {}", viewer_path.display());

crates/stylus-trace-core/src/diff/output.rs

@@ -145,12 +145,12 @@ fn render_hot_path_comparison_table(report: &DiffReport) -> String {
             format!("{:<38}", display_stack)
         };
 
-        // Scale to Gas (ink / 10,000)
-        let baseline_gas = hp.baseline_gas / 10_000;
-        let target_gas = hp.target_gas / 10_000;
+        // Scale to Gas (ink / 10,000) with float precision
+        let baseline_gas = hp.baseline_gas as f64 / 10_000.0;
+        let target_gas = hp.target_gas as f64 / 10_000.0;
 
         out.push_str(&format!(
-            "  ┃ {} ┃ {:>12} ┃ {:>12} ┃ {}{:>9.2}%{} ┃\n",
+            "  ┃ {} ┃ {:>12.1} ┃ {:>12.1} ┃ {}{:>9.2}%{} ┃\n",
             display_stack_fixed, baseline_gas, target_gas, delta_color, hp.percent_change, reset
         ));
     }

crates/stylus-trace-core/src/output/viewer.rs

@@ -2,69 +2,90 @@
 
 use crate::parser::schema::Profile;
 use anyhow::{Context, Result};
+use base64::engine::general_purpose::STANDARD as BASE64;
+use base64::Engine;
 use std::fs;
 use std::path::Path;
 
 const HTML_TEMPLATE: &str = include_str!("viewer/index.html");
 const CSS_TEMPLATE: &str = include_str!("viewer/viewer.css");
 const JS_TEMPLATE: &str = include_str!("viewer/viewer.js");
 
-/// Generate a self-contained HTML viewer for a profile
-pub fn generate_viewer(profile: &Profile, output_path: &Path) -> Result<()> {
-    let profile_json = serde_json::to_string(profile)?;
+/// Encode a JSON string as Base64 for safe inline HTML injection.
+/// This prevents any `</script>` sequences in the JSON from breaking the page.
+fn encode_json(json: &str) -> String {
+    BASE64.encode(json.as_bytes())
+}
 
-    // In a real implementation with more time, we'd use a proper template engine or
-    // at least a better replacement strategy. For this "best effort" we'll do simple replacement.
+/// Build the final, self-contained HTML from the template.
+///
+/// Inlines CSS and JS, and injects the base64 JSON blobs + optional SVG.
+fn build_html(
+    profile_a_b64: &str,
+    profile_b_b64: Option<&str>,
+    diff_b64: Option<&str>,
+    flamegraph_svg: Option<&str>,
+) -> String {
     let mut html = HTML_TEMPLATE.to_string();
 
-    // Inject data
-    html = html.replace("/* PROFILE_DATA_JSON */", &profile_json);
+    // Inject data placeholders (base64-encoded JSON)
+    html = html.replace("/* PROFILE_DATA_B64 */", profile_a_b64);
+    html = html.replace(
+        "/* PROFILE_B_DATA_B64 */",
+        profile_b_b64.unwrap_or_default(),
+    );
+    html = html.replace("/* DIFF_DATA_B64 */", diff_b64.unwrap_or_default());
+    html = html.replace("/* FLAMEGRAPH_SVG */", flamegraph_svg.unwrap_or_default());
 
-    // Inline CSS and JS for "self-contained" requirement
+    // Inline CSS
     html = html.replace(
         "<link rel=\"stylesheet\" href=\"viewer.css\">",
         &format!("<style>{}</style>", CSS_TEMPLATE),
     );
+
+    // Inline JS
     html = html.replace(
         "<script src=\"viewer.js\"></script>",
         &format!("<script>{}</script>", JS_TEMPLATE),
     );
 
-    fs::write(output_path, html).context("Failed to write viewer HTML")?;
+    html
+}
+
+/// Generate a self-contained HTML viewer for a single profile.
+pub fn generate_viewer(
+    profile: &Profile,
+    flamegraph_svg: Option<&str>,
+    output_path: &Path,
+) -> Result<()> {
+    let profile_json = serde_json::to_string(profile)?;
+    let profile_b64 = encode_json(&profile_json);
+
+    let html = build_html(&profile_b64, None, None, flamegraph_svg);
 
+    fs::write(output_path, html).context("Failed to write viewer HTML")?;
     Ok(())
 }
 
-/// Generate a self-contained HTML viewer for a diff
+/// Generate a self-contained HTML viewer for a diff (baseline vs target).
 pub fn generate_diff_viewer(
     profile_a: &Profile,
     profile_b: &Profile,
     diff_report: &serde_json::Value,
+    flamegraph_svg: Option<&str>,
     output_path: &Path,
 ) -> Result<()> {
     let profile_a_json = serde_json::to_string(profile_a)?;
     let profile_b_json = serde_json::to_string(profile_b)?;
     let diff_json = serde_json::to_string(diff_report)?;
 
-    let mut html = HTML_TEMPLATE.to_string();
-
-    // Inject data
-    html = html.replace("/* PROFILE_DATA_JSON */", &profile_a_json);
-    html = html.replace("/* PROFILE_B_DATA_JSON */", &profile_b_json);
-    html = html.replace("/* DIFF_DATA_JSON */", &diff_json);
+    let a_b64 = encode_json(&profile_a_json);
+    let b_b64 = encode_json(&profile_b_json);
+    let diff_b64 = encode_json(&diff_json);
 
-    // Inline CSS and JS
-    html = html.replace(
-        "<link rel=\"stylesheet\" href=\"viewer.css\">",
-        &format!("<style>{}</style>", CSS_TEMPLATE),
-    );
-    html = html.replace(
-        "<script src=\"viewer.js\"></script>",
-        &format!("<script>{}</script>", JS_TEMPLATE),
-    );
+    let html = build_html(&a_b64, Some(&b_b64), Some(&diff_b64), flamegraph_svg);
 
     fs::write(output_path, html).context("Failed to write diff viewer HTML")?;
-
     Ok(())
 }
 

crates/stylus-trace-core/src/output/viewer/index.html

@@ -11,6 +11,14 @@
 </head>
 <body>
     <div class="crt-overlay"></div>
+    <div id="drop-zone-container" class="hidden">
+        <div class="drop-zone">
+            <div class="glitch" data-text="AWAITING // TRACE_DATA">AWAITING // TRACE_DATA</div>
+            <p>DRAG &amp; DROP REPORT.JSON</p>
+            <p class="subtext">Drop 1 file for Profiling | Drop 2 files for Diffing</p>
+            <input type="file" id="file-input" class="hidden" accept=".json" multiple>
+        </div>
+    </div>
     <div id="app">
         <header class="terminal-header">
             <div class="logo">
@@ -26,9 +34,16 @@
                 <button id="zoom-in">[ + ]</button>
                 <button id="zoom-out">[ - ]</button>
                 <button id="reset-view">[ RESET ]</button>
+                <button id="theme-toggle">[ LIGHT ]</button>
             </div>
         </header>
 
+        <!-- View Tab Switcher -->
+        <div id="view-tabs">
+            <button id="tab-pie" class="tab-btn active">[ PIE_CHART ]</button>
+            <button id="tab-flame" class="tab-btn">[ FLAMEGRAPH ]</button>
+        </div>
+
         <main>
             <aside id="sidebar">
                 <section class="tx-info panel">
@@ -59,36 +74,46 @@ <h2>:: HOT_PATHS ::</h2>
                         <!-- Populated by JS -->
                     </ul>
                 </section>
+
+                <section id="category-stats" class="panel">
+                    <h2>:: GAS_BY_CATEGORY ::</h2>
+                    <ul id="category-stats-list">
+                        <!-- Populated by JS -->
+                    </ul>
+                </section>
             </aside>
 
+            <!-- Pie Chart View -->
             <div id="viewer-container">
                 <div id="chart-a" class="chart-canvas">
                     <canvas id="canvas-a"></canvas>
-                    <div class="label">> BASELINE_DATA</div>
+                    <div class="label">&gt; BASELINE_DATA</div>
                 </div>
                 <div id="chart-b" class="chart-canvas hidden">
                     <canvas id="canvas-b"></canvas>
-                    <div class="label">> TARGET_DATA</div>
+                    <div class="label">&gt; TARGET_DATA</div>
                 </div>
                 <div id="tooltip"></div>
             </div>
+
+            <!-- Flamegraph View -->
+            <div id="flamegraph-view" class="hidden">
+                <div id="flamegraph-container"></div>
+                <div id="flamegraph-empty" class="flamegraph-empty hidden">&gt; NO_FLAMEGRAPH_DATA // Profile with --output-svg flag to generate one</div>
+            </div>
         </main>
 
         <footer class="terminal-footer">
-            <div id="status-line">> SYSTEM: ONLINE | PROFILE: <span id="profile-name">none</span></div>
+            <div id="status-line">&gt; SYSTEM: ONLINE | PROFILE: <span id="profile-name">none</span></div>
         </footer>
     </div>
 
-    <!-- Data Injection Point -->
-    <script id="profile-data" type="application/json">
-        /* PROFILE_DATA_JSON */
-    </script>
-    <script id="profile-b-data" type="application/json">
-        /* PROFILE_B_DATA_JSON */
-    </script>
-    <script id="diff-data" type="application/json">
-        /* DIFF_DATA_JSON */
-    </script>
+    <!-- Data Injection Point (base64-encoded JSON for safe injection) -->
+    <script id="profile-data" type="application/json">/* PROFILE_DATA_B64 */</script>
+    <script id="profile-b-data" type="application/json">/* PROFILE_B_DATA_B64 */</script>
+    <script id="diff-data" type="application/json">/* DIFF_DATA_B64 */</script>
+    <!-- Flamegraph SVG (injected as raw text) -->
+    <script id="flamegraph-svg-data" type="text/plain">/* FLAMEGRAPH_SVG */</script>
 
     <script src="viewer.js"></script>
 </body>