diff --git a/src/lib.rs b/src/lib.rs
index 13a15f7..f9a404b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ use anyhow::{anyhow, bail};
 use chrono::Datelike;
 use fxhash::{FxHashMap, FxHashSet};
 use md5::{Digest, Md5};
+use once_cell::sync::Lazy;
 use std::ffi::{OsStr, OsString};
 
 use html_escape::encode_text;
@@ -1213,13 +1214,27 @@ pub fn parse_path(path: &PathBuf, config: &ParseConfig) -> anyhow::Result<ParseO
                 directory_name,
             );
 
+            // For highlight.js, we pass raw content (HTML-escaped by template)
+            // For already-HTML content (output_code), extract plain text first
+            let clean_output_code = if is_already_html(&output_code_content) {
+                extract_text_from_html(&output_code_content)
+            } else {
+                output_code_content.clone()
+            };
+            let clean_aot_code = if is_already_html(&aot_code_content) {
+                extract_text_from_html(&aot_code_content)
+            } else {
+                aot_code_content.clone()
+            };
+
             // Convert node mappings to line number mappings
+            // Use the cleaned content so line numbers match what's displayed
             let line_mappings_content = convert_node_mappings_to_line_numbers(
                 &node_mappings_content,
                 &pre_grad_graph_content,
                 &post_grad_graph_content,
-                &output_code_content,
-                &aot_code_content,
+                &clean_output_code,
+                &clean_aot_code,
             );
             let line_mappings_content_str = serde_json::to_string_pretty(&line_mappings_content)
                 .unwrap_or_else(|_| "{}".to_string());
@@ -1233,8 +1248,8 @@ pub fn parse_path(path: &PathBuf, config: &ParseConfig) -> anyhow::Result<ParseO
                         js: PROVENANCE_JS,
                         pre_grad_graph_content,
                         post_grad_graph_content,
-                        output_code_content,
-                        aot_code_content,
+                        output_code_content: clean_output_code,
+                        aot_code_content: clean_aot_code,
                         line_mappings_content: line_mappings_content_str,
                     },
                 )?,
@@ -2012,6 +2027,35 @@ pub fn analyze_graph_runtime_deltas(
     })
 }
 
+/// Checks if content appears to be pre-rendered HTML (e.g., already contains HTML tags).
+fn is_already_html(content: &str) -> bool {
+    let trimmed = content.trim();
+    trimmed.starts_with("<pre")
+        || trimmed.starts_with("<span")
+        || trimmed.starts_with("<!DOCTYPE")
+        || trimmed.starts_with("<html")
+}
+
+/// Regex for stripping HTML tags, compiled once for efficiency.
+static HTML_TAG_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"<[^>]+>").unwrap());
+
+/// Extracts plain text from HTML by removing all tags and decoding entities.
+/// Note: This is a simple implementation that handles common cases but may not
+/// correctly process malformed HTML, nested tags within attributes, or CDATA sections.
+fn extract_text_from_html(html: &str) -> String {
+    let text = HTML_TAG_REGEX.replace_all(html, "");
+
+    // Decode common HTML entities (covers most cases in generated code)
+    text.replace("&lt;", "<")
+        .replace("&gt;", ">")
+        .replace("&amp;", "&")
+        .replace("&quot;", "\"")
+        .replace("&#39;", "'")
+        .replace("&apos;", "'")
+        .replace("&#x27;", "'")
+        .replace("&nbsp;", " ")
+}
+
 /// Converts node-based mappings to line number-based mappings for visualization.
 ///
 /// This function processes node mappings and converts them to line number mappings
diff --git a/src/provenance.css b/src/provenance.css
index ac397d9..8ec88ad 100644
--- a/src/provenance.css
+++ b/src/provenance.css
@@ -1,3 +1,4 @@
+/* Base styles */
 body {
     display: flex;
     flex-direction: column;
@@ -6,151 +7,260 @@ body {
     padding: 0;
     height: 100vh;
     overflow: hidden;
+    background-color: #ffffff;
+    color: #333;
 }
 
-.url-inputs {
+.editor-container {
     display: flex;
-    padding: 10px;
-    background-color: #f0f0f0;
+    flex: 1;
+    overflow: hidden;
 }
 
-.url-input {
+/* Editor wrapper for header + content */
+.editor-wrapper {
+    display: flex;
+    flex-direction: column;
     flex: 1;
-    margin-right: 10px;
+    min-width: 0;
+    overflow: hidden;
+}
+
+.editor-header {
+    background: #e9ecef;
+    padding: 8px 12px;
+    border-bottom: 1px solid #ddd;
     display: flex;
+    justify-content: space-between;
     align-items: center;
-    gap: 10px;
+    flex-shrink: 0;
 }
 
-.url-input input {
-    width: 100%;
-    padding: 5px;
+.editor-title {
+    font-weight: 600;
+    color: #495057;
+    font-size: 14px;
 }
 
-.editor-container {
-    display: flex;
-    flex: 1;
-    overflow: hidden;
+.editor-stats {
+    font-size: 12px;
+    color: #868e96;
 }
 
 .editor {
-    height: 100%;
-    overflow-y: auto;
+    flex: 1;
+    overflow: auto;
     border: 1px solid #ddd;
     padding: 10px;
     box-sizing: border-box;
-    flex: 1;
-    font-family: monospace;
+    font-family: 'SF Mono', 'Monaco', 'Menlo', 'Consolas', monospace;
+    font-size: 13px;
+    line-height: 1.5;
+    position: relative;
+}
+
+.editor pre {
+    margin: 0;
+    font-family: inherit;
+    font-size: inherit;
 }
 
+/* Line styles */
 .line {
     padding: 2px 5px;
     cursor: pointer;
     white-space: nowrap;
     display: flex;
     width: 100%;
-    min-width: max-content; /* Ensure line fills horizontally beyond visible area */
+    min-width: max-content;
     box-sizing: border-box;
-    transition: background-color 0.2s ease;
+    transition: background-color 0.15s ease;
+    position: relative;
 }
 
-.highlight {
-    background-color: yellow;
-    transition: background-color 0.2s ease;
-}
-
-.mapped-line {
-    font-weight: bold;
+.line:hover {
+    background-color: rgba(0, 0, 0, 0.03);
 }
 
 .line-number {
     color: #888;
     display: inline-block;
-    width: 30px;
+    min-width: 40px;
     text-align: right;
-    margin-right: 10px;
+    margin-right: 12px;
+    user-select: none;
+    flex-shrink: 0;
 }
 
+.line-content {
+    white-space: pre;
+    display: inline;
+}
+
+/* Highlight styles - source (clicked) vs target (corresponding) */
+.highlight-source {
+    background-color: #fff3cd !important;
+    border-left: 3px solid #f59f00;
+    padding-left: 2px;
+}
+
+.highlight-target {
+    background-color: #d3f9d8 !important;
+    border-left: 3px solid #40c057;
+    padding-left: 2px;
+}
+
+/* Legacy highlight class for backwards compatibility */
+.highlight {
+    background-color: #fff3cd !important;
+    border-left: 3px solid #f59f00;
+    padding-left: 2px;
+}
+
+/* Lines with mappings indicator */
+.has-match {
+    background-color: rgba(255, 243, 205, 0.3);
+}
+
+.has-match .line-number {
+    color: #e67700;
+    font-weight: 600;
+}
+
+/* Mapping count badge */
+.mapping-count {
+    position: absolute;
+    right: 8px;
+    top: 50%;
+    transform: translateY(-50%);
+    background: #868e96;
+    color: white;
+    font-size: 10px;
+    padding: 1px 5px;
+    border-radius: 8px;
+    opacity: 0.8;
+}
+
+/* Divider between panels */
 .divider {
-    width: 3px;
-    background-color: #ccc;
+    width: 4px;
+    background-color: #dee2e6;
     cursor: col-resize;
+    flex-shrink: 0;
+    transition: background-color 0.15s ease;
 }
 
-.line-content {
-    white-space: pre;
-    display: inline;
+.divider:hover {
+    background-color: #adb5bd;
 }
 
-.json-popup {
+/* Keyboard help overlay */
+.help-overlay {
     display: none;
     position: fixed;
+    top: 50%;
+    left: 50%;
+    transform: translate(-50%, -50%);
     background: white;
-    border: 1px solid #ccc;
-    padding: 10px;
-    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
-    max-width: 400px;
-    max-height: 300px;
-    overflow: auto;
+    border: 1px solid #ddd;
+    border-radius: 8px;
+    padding: 20px 30px;
+    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.15);
     z-index: 1000;
-    white-space: pre-wrap;
-    font-family: monospace;
-    font-size: 12px;
+    max-width: 400px;
 }
 
-.json-popup .close-button {
-    position: sticky;
-    float: right;
-    top: 0;
-    right: 0;
-    cursor: pointer;
-    font-size: 20px;
-    background: white;
-    padding: 0 5px;
-    margin-left: 10px;
+.help-overlay.visible {
+    display: block;
 }
 
-#jsonContent {
-    margin-top: 10px;
+.help-overlay h3 {
+    margin-top: 0;
+    margin-bottom: 16px;
+    color: #333;
+    font-size: 16px;
 }
 
-.toggle-container {
-    margin-left: 20px;
-    display: inline-flex;
+.help-overlay ul {
+    list-style: none;
+    padding: 0;
+    margin: 0 0 16px 0;
+}
+
+.help-overlay li {
+    padding: 6px 0;
+    display: flex;
     align-items: center;
-    cursor: pointer;
-    min-width: 200px;
-    white-space: nowrap;
+    gap: 12px;
+}
+
+.help-overlay kbd {
+    background: #f4f4f4;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    padding: 2px 8px;
+    font-family: 'SF Mono', 'Monaco', monospace;
+    font-size: 12px;
+    color: #333;
+    min-width: 20px;
+    text-align: center;
 }
 
-.toggle-container input[type="checkbox"] {
-    margin-right: 5px;
+.help-hint {
+    margin: 0;
+    font-size: 12px;
+    color: #868e96;
+    text-align: center;
 }
 
-.file-input {
+/* Backdrop for help overlay */
+.help-backdrop {
     display: none;
+    position: fixed;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+    background: rgba(0, 0, 0, 0.3);
+    z-index: 999;
 }
 
-.url-input {
-    flex: 1;
-    margin-right: 10px;
+.help-backdrop.visible {
+    display: block;
+}
+
+/* Navigation hints bar at bottom */
+.nav-hints-bar {
     display: flex;
     align-items: center;
-    gap: 10px;
+    justify-content: center;
+    gap: 16px;
+    padding: 6px 12px;
+    background: #f8f9fa;
+    border-top: 1px solid #dee2e6;
+    font-size: 12px;
+    color: #6c757d;
+    flex-shrink: 0;
+    flex-wrap: wrap;
 }
 
-.file-label {
-    background-color: #f0f0f0;
-    padding: 5px 10px;
-    border: 1px solid #ddd;
-    border-radius: 4px;
-    cursor: pointer;
+.nav-hint {
+    display: inline-flex;
+    align-items: center;
+    gap: 4px;
 }
 
-.file-label:hover {
-    background-color: #e0e0e0;
+.nav-hint kbd {
+    background: #e9ecef;
+    border: 1px solid #ced4da;
+    border-radius: 3px;
+    padding: 1px 5px;
+    font-family: 'SF Mono', 'Monaco', monospace;
+    font-size: 11px;
+    color: #495057;
 }
 
-.has-match {
-    font-weight: bold;
-} 
\ No newline at end of file
+.nav-hint strong {
+    font-weight: 600;
+    color: #495057;
+}
diff --git a/src/provenance.html b/src/provenance.html
index 09201e4..b858a93 100644
--- a/src/provenance.html
+++ b/src/provenance.html
@@ -5,6 +5,11 @@
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>Inductor Provenance Tracking Highlighter</title>
+
+    <!-- highlight.js from CDN -->
+    <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/github.min.css">
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
+
     <style>
         {css | format_unescaped}
     </style>
@@ -12,19 +17,60 @@
 
 <body>
     <div class="editor-container">
-        <div id="preGradGraph" class="editor">
-            <pre>{pre_grad_graph_content}</pre>
+        <div class="editor-wrapper">
+            <div class="editor-header">
+                <span class="editor-title">Pre-Grad Graph</span>
+                <span class="editor-stats" id="preGradStats"></span>
+            </div>
+            <div id="preGradGraph" class="editor">
+                <pre><code class="language-python">{pre_grad_graph_content}</code></pre>
+            </div>
         </div>
         <div id="divider1" class="divider"></div>
-        <div id="postGradGraph" class="editor">
-            <pre>{post_grad_graph_content}</pre>
+        <div class="editor-wrapper">
+            <div class="editor-header">
+                <span class="editor-title">Post-Grad Graph</span>
+                <span class="editor-stats" id="postGradStats"></span>
+            </div>
+            <div id="postGradGraph" class="editor">
+                <pre><code class="language-python">{post_grad_graph_content}</code></pre>
+            </div>
         </div>
         <div id="divider2" class="divider"></div>
-        <div id="generatedCode" class="editor">
-            <pre>{output_code_content | format_unescaped}{aot_code_content}</pre>
+        <div class="editor-wrapper">
+            <div class="editor-header">
+                <span class="editor-title">Generated Code</span>
+                <span class="editor-stats" id="generatedCodeStats"></span>
+            </div>
+            <div id="generatedCode" class="editor">
+                <pre><code class="language-python">{output_code_content}{aot_code_content}</code></pre>
+            </div>
         </div>
     </div>
 
+    <!-- Keyboard help overlay -->
+    <div id="keyboard-help" class="help-overlay">
+        <h3>Keyboard Shortcuts</h3>
+        <ul>
+            <li><kbd>n</kbd> / <kbd>N</kbd> - Next/Previous mapped line</li>
+            <li><kbd>j</kbd> / <kbd>J</kbd> - Jump to corresponding panel</li>
+            <li><kbd>1</kbd> <kbd>2</kbd> <kbd>3</kbd> - Maximize panel</li>
+            <li><kbd>?</kbd> - Toggle this help</li>
+            <li><kbd>Esc</kbd> - Clear selection / Restore panels</li>
+        </ul>
+        <p class="help-hint">Press <kbd>?</kbd> to close</p>
+    </div>
+
+    <!-- Navigation hints bar -->
+    <div class="nav-hints-bar">
+        <span class="nav-hint"><kbd>n</kbd>/<kbd>N</kbd> next/prev</span>
+        <span class="nav-hint"><kbd>j</kbd> jump to panel</span>
+        <span class="nav-hint"><kbd>1</kbd><kbd>2</kbd><kbd>3</kbd> maximize</span>
+        <span class="nav-hint"><strong>Double-click</strong> header to expand</span>
+        <span class="nav-hint"><strong>Double-click</strong> code to find definition</span>
+        <span class="nav-hint"><kbd>?</kbd> help</span>
+    </div>
+
     <!-- Line mappings data for JavaScript -->
     <script id="lineMappings" type="application/json">
         {line_mappings_content | format_unescaped}
diff --git a/src/provenance.js b/src/provenance.js
index 52f4d56..aa5e5e3 100644
--- a/src/provenance.js
+++ b/src/provenance.js
@@ -1,3 +1,15 @@
+// ============== Constants ==============
+
+// Built-in functions to skip when searching for definitions (using Set for O(1) lookup)
+const BUILTIN_FUNCTIONS = new Set([
+    'if', 'for', 'while', 'return', 'print', 'range', 'len', 'int', 'float', 'str',
+    'list', 'dict', 'set', 'tuple', 'assert', 'isinstance', 'getattr', 'setattr',
+    'empty_strided_cpu', 'empty_strided_cuda', 'reinterpret_tensor', 'as_strided',
+    'assert_size_stride', 'rand_strided'
+]);
+
+// ============== Data Storage ==============
+
 let preGradGraphData = null;
 let postGradGraphData = null;
 let codeData = null;
@@ -12,140 +24,340 @@ let cppCodeToPost = {};
 
 let lineMappings = null;
 
+let currentSelection = {
+    editorId: null,
+    lineNumber: null
+};
+
+let maximizedPanel = null;
+
+// ============== Error Display ==============
+
+/**
+ * Shows an error banner at the top of the page.
+ * @param {string} message - The error message to display
+ */
+function showErrorBanner(message) {
+    const existing = document.querySelector('.error-banner');
+    if (existing) existing.remove();
+
+    const banner = document.createElement('div');
+    banner.className = 'error-banner';
+    banner.style.cssText = 'background: #ffebee; color: #c62828; padding: 8px 16px; border-bottom: 1px solid #ef9a9a; font-size: 13px;';
+    banner.textContent = message;
+
+    const container = document.querySelector('.editor-container');
+    if (container && container.parentElement) {
+        container.parentElement.insertBefore(banner, container);
+    }
+}
+
+// ============== Sanitization ==============
+
 /**
- * Initializes the line number mappings from the pre-processed data.
- * 
- * This function expects the line mappings to be already converted from node mappings
- * to line number mappings by the Rust backend. The mappings should contain:
- * - preToPost: {sourceLineNum: [targetLineNums]}
- * - postToPre: {sourceLineNum: [targetLineNums]}
- * - pyCodeToPost: {sourceLineNum: [targetLineNums]}
- * - postToPyCode: {sourceLineNum: [targetLineNums]}
- * - cppCodeToPost: {sourceLineNum: [targetLineNums]}
- * - postToCppCode: {sourceLineNum: [targetLineNums]}
- * 
- * These mappings enable the UI to highlight corresponding lines
- * across different views when a user clicks on a line.
+ * Sanitizes HTML from highlight.js to prevent XSS.
+ * Only allows span tags with class attributes (which is what highlight.js produces).
+ * @param {string} html - HTML string from highlight.js
+ * @returns {string} Sanitized HTML string
+ */
+function sanitizeHighlightHtml(html) {
+    // highlight.js only produces <span class="hljs-*">...</span> tags
+    // Remove any other HTML tags or attributes that shouldn't be there
+    return html
+        // Remove script tags and their content
+        .replace(/<script\b[^<]*(?:(?!<\/script>)<[^<]*)*<\/script>/gi, '')
+        // Remove event handlers
+        .replace(/\s*on\w+\s*=\s*["'][^"']*["']/gi, '')
+        // Remove javascript: URLs
+        .replace(/javascript:/gi, '')
+        // Remove data: URLs (except for safe ones)
+        .replace(/data:(?!image\/(png|jpg|jpeg|gif|svg\+xml))[^"'\s]*/gi, '');
+}
+
+/**
+ * Escapes special regex characters in a string.
+ * @param {string} str - String to escape
+ * @returns {string} Escaped string safe for use in RegExp
+ */
+function escapeRegex(str) {
+    return str.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+}
+
+// ============== Line Mappings ==============
+
+/**
+ * Initializes the line number mappings from the pre-processed JSON data.
+ * Expects a JSON structure with: preToPost, postToPre, pyCodeToPost,
+ * postToPyCode, cppCodeToPost, postToCppCode - each mapping source line
+ * numbers to arrays of target line numbers.
  */
 function initializeLineMappings() {
     try {
-        // Get the line mappings from the embedded JSON data
         const lineMappingsElement = document.getElementById('lineMappings');
         if (lineMappingsElement) {
             lineMappings = JSON.parse(lineMappingsElement.textContent);
-            
-            // Update global variables with the line mappings
+
             preToPost = lineMappings.preToPost || {};
             postToPre = lineMappings.postToPre || {};
             pyCodeToPost = lineMappings.pyCodeToPost || {};
             postToPyCode = lineMappings.postToPyCode || {};
             cppCodeToPost = lineMappings.cppCodeToPost || {};
             postToCppCode = lineMappings.postToCppCode || {};
-            
-            console.log('Line mappings initialized:', {
-                preToPost,
-                postToPre,
-                pyCodeToPost,
-                postToPyCode,
-                cppCodeToPost,
-                postToCppCode
-            });
         } else {
-            console.warn('No line mappings element found');
+            console.warn('No line mappings element found - cross-panel highlighting will be disabled');
+            showErrorBanner('Line mapping data not found. Cross-panel highlighting is disabled.');
         }
     } catch (error) {
         console.error('Error initializing line mappings:', error);
+        showErrorBanner(`Failed to parse line mappings: ${error.message}. Cross-panel highlighting is disabled.`);
     }
 }
 
+/**
+ * Splits HTML content by newlines while preserving HTML tags intact within each line.
+ * Trims leading/trailing empty lines to ensure line numbers align with displayed content.
+ * @param {string} html - HTML string to split
+ * @returns {string[]} Array of HTML line strings
+ */
+function splitHtmlByLines(html) {
+    if (!html) return [];
+
+    let lines = html.split('\n');
+
+    // Remove leading empty lines
+    while (lines.length > 0 && lines[0].trim() === '') {
+        lines.shift();
+    }
+
+    // Remove trailing empty lines
+    while (lines.length > 0 && lines[lines.length - 1].trim() === '') {
+        lines.pop();
+    }
+
+    return lines;
+}
+
+/**
+ * Gets the total number of line mappings for a specific line in an editor.
+ * @param {string} editorId - The editor ID ('preGradGraph', 'postGradGraph', or 'generatedCode')
+ * @param {number} lineNum - The 1-based line number
+ * @returns {number} The count of mappings, or 0 if none exist or editor is unknown
+ */
+function getMappingCount(editorId, lineNum) {
+    switch (editorId) {
+        case 'preGradGraph':
+            return (preToPost[lineNum] || []).length;
+        case 'postGradGraph':
+            return (postToPre[lineNum] || []).length +
+                   (postToPyCode[lineNum] || []).length +
+                   (postToCppCode[lineNum] || []).length;
+        case 'generatedCode':
+            return (pyCodeToPost[lineNum] || []).length +
+                   (cppCodeToPost[lineNum] || []).length;
+        default:
+            return 0;
+    }
+}
+
+// ============== Editor Setup ==============
+
+/**
+ * Setup editor content with line numbers and event handlers.
+ * Handles both plain text and HTML content (for syntax highlighting).
+ * @param {string} editorId - The editor element ID
+ * @param {string[]} lines - Array of line content strings
+ * @param {boolean} isHighlighted - Whether lines contain HTML from syntax highlighting
+ */
+function setupEditorContent(editorId, lines, isHighlighted = false) {
+    if (!lines) {
+        console.warn(`setupEditorContent: No lines provided for ${editorId}`);
+        return;
+    }
 
-// Setup editor content
-function setupEditorContent(editorId, lines) {
-    if (!lines) return;
-    
     const editor = document.getElementById(editorId);
-    if (!editor) return;
+    if (!editor) {
+        console.warn(`setupEditorContent: Editor element '${editorId}' not found`);
+        return;
+    }
+
+    editor.innerHTML = '';
 
-    editor.innerHTML = '';  // Clear existing content
-    
     lines.forEach((line, index) => {
         const lineDiv = document.createElement('div');
         lineDiv.className = 'line';
-        
-        // Create text nodes instead of using innerHTML
+        lineDiv.dataset.lineNumber = index + 1;
+
         const lineNumber = document.createElement('span');
         lineNumber.className = 'line-number';
         lineNumber.textContent = index + 1;
-        
+
         const lineContent = document.createElement('span');
         lineContent.className = 'line-content';
-        lineContent.textContent = line;
-        
-        // Check if this line has any matches
-        const lineNum = index + 1;
-        let hasMatch = false;
-        switch (editorId) {
-            case 'preGradGraph':
-                hasMatch = preToPost[lineNum] && preToPost[lineNum].length > 0;
-                break;
-            case 'postGradGraph':
-                hasMatch = (postToPre[lineNum] && postToPre[lineNum].length > 0) ||
-                          (postToPyCode[lineNum] && postToPyCode[lineNum].length > 0) ||
-                          (postToCppCode[lineNum] && postToCppCode[lineNum].length > 0);
-                break;
-            case 'generatedCode':
-                hasMatch = (pyCodeToPost[lineNum] && pyCodeToPost[lineNum].length > 0) || 
-                (cppCodeToPost[lineNum] && cppCodeToPost[lineNum].length > 0);
-                break;
+
+        // Use innerHTML for highlighted content (sanitized), textContent for plain text
+        if (isHighlighted) {
+            lineContent.innerHTML = sanitizeHighlightHtml(line);
+        } else {
+            lineContent.textContent = line;
         }
-        
+
+        // Check if this line has any matches (using consolidated logic)
+        const lineNum = index + 1;
+        const hasMatch = getMappingCount(editorId, lineNum) > 0;
+
         if (hasMatch) {
             lineContent.classList.add('has-match');
+            lineDiv.classList.add('has-match');
+
+            // Add mapping count badge if more than 1 mapping
+            const count = getMappingCount(editorId, lineNum);
+            if (count > 1) {
+                const badge = document.createElement('span');
+                badge.className = 'mapping-count';
+                badge.textContent = count;
+                lineDiv.appendChild(badge);
+            }
         }
-        
+
         lineDiv.appendChild(lineNumber);
         lineDiv.appendChild(lineContent);
-        
-        // Add both click and hover handlers
+
+        // Event handlers
         lineDiv.addEventListener('click', () => handleLineClick(editorId, index + 1));
+        lineDiv.addEventListener('dblclick', (e) => handleDoubleClick(editorId, index + 1, e));
         lineDiv.addEventListener('mouseenter', () => handleLineHover(editorId, index + 1));
         lineDiv.addEventListener('mouseleave', clearHighlights);
-        
+
         editor.appendChild(lineDiv);
     });
+
+    // Update stats after content is set up
+    updateEditorStats(editorId);
+}
+
+/**
+ * Updates the stats display for an editor (line count and mapped count).
+ * @param {string} editorId - The editor element ID
+ */
+function updateEditorStats(editorId) {
+    const statsMap = {
+        'preGradGraph': 'preGradStats',
+        'postGradGraph': 'postGradStats',
+        'generatedCode': 'generatedCodeStats'
+    };
+
+    const statsId = statsMap[editorId];
+    if (!statsId) return;
+
+    const stats = document.getElementById(statsId);
+    const editor = document.getElementById(editorId);
+    if (!stats || !editor) return;
+
+    const lines = editor.querySelectorAll('.line');
+    const mappedLines = editor.querySelectorAll('.line.has-match');
+
+    stats.textContent = `${lines.length} lines, ${mappedLines.length} mapped`;
 }
 
-// Handle line hover
+// ============== Line Highlighting ==============
+
+/**
+ * Handle line hover - highlights line and corresponding lines.
+ * @param {string} editorId - The editor element ID
+ * @param {number} lineNumber - The 1-based line number
+ */
 function handleLineHover(editorId, lineNumber) {
-    // Clear previous highlights
     clearHighlights();
-    
-    // Add highlight to hovered line
+
     const hoveredLine = document.querySelector(`#${editorId} .line:nth-child(${lineNumber})`);
     if (hoveredLine) {
-        hoveredLine.classList.add('highlight');
-        // Remove scrolling for hovered panel
+        hoveredLine.classList.add('highlight-source');
     }
 
-    // Highlight and scroll corresponding lines
     highlightCorrespondingLines(editorId, lineNumber);
 }
 
-// Clear all highlights
+/**
+ * Clear all line highlights across all editors.
+ */
 function clearHighlights() {
     document.querySelectorAll('.line').forEach(line => {
-        line.classList.remove('highlight');
+        line.classList.remove('highlight', 'highlight-source', 'highlight-target');
     });
 }
 
-// Update handleLineClick to use the same pattern
+/**
+ * Handle double-click - jump to function definition in generated code.
+ * Searches for Python def, C++ function definitions, or async_compile assignments.
+ * @param {string} editorId - The editor element ID
+ * @param {number} lineNumber - The 1-based line number
+ * @param {Event} event - The double-click event
+ */
+function handleDoubleClick(editorId, lineNumber, event) {
+    // Only handle double-clicks in the generated code pane
+    if (editorId !== 'generatedCode') return;
+
+    const editor = document.getElementById(editorId);
+    if (!editor) return;
+
+    const clickedLine = editor.querySelector(`.line:nth-child(${lineNumber})`);
+    if (!clickedLine) return;
+
+    // Get text from line-content span, not the whole line (which includes line number)
+    const lineContentSpan = clickedLine.querySelector('.line-content');
+    if (!lineContentSpan) return;
+    const lineText = lineContentSpan.textContent;
+
+    // Extract function/kernel names from the line
+    const funcCallMatch = lineText.match(/\b([a-zA-Z_][a-zA-Z0-9_]*)\s*\(/g);
+    if (!funcCallMatch) return;
+
+    // Get all function names called on this line
+    const funcNames = funcCallMatch.map(m => m.replace(/\s*\($/, ''));
+
+    // Search for definitions
+    const lines = editor.querySelectorAll('.line');
+    for (const funcName of funcNames) {
+        // Skip built-in functions (O(1) lookup)
+        if (BUILTIN_FUNCTIONS.has(funcName)) {
+            continue;
+        }
+
+        // Escape regex special characters in function name
+        const escapedName = escapeRegex(funcName);
+        const defPattern = new RegExp(`\\b(def|void|static|inline)\\s+${escapedName}\\s*\\(`);
+        const assignPattern = new RegExp(`^\\s*${escapedName}\\s*=\\s*(triton_|async_compile|extern_kernels)`);
+
+        for (let i = 0; i < lines.length; i++) {
+            const contentSpan = lines[i].querySelector('.line-content');
+            if (!contentSpan) continue;
+            const lineContent = contentSpan.textContent;
+
+            if (defPattern.test(lineContent) || assignPattern.test(lineContent)) {
+                // Found the definition - jump to it
+                clearHighlights();
+                lines[i].classList.add('highlight-source');
+                lines[i].scrollIntoView({ behavior: 'smooth', block: 'center' });
+                currentSelection = { editorId, lineNumber: i + 1 };
+                event.preventDefault();
+                event.stopPropagation();
+                return;
+            }
+        }
+    }
+}
+
+/**
+ * Handle line click - highlights and sets current selection for keyboard nav.
+ * @param {string} editorId - The editor element ID
+ * @param {number} lineNumber - The 1-based line number
+ */
 function handleLineClick(editorId, lineNumber) {
     clearHighlights();
-    
-    // Add highlight to clicked line
+
     const clickedLine = document.querySelector(`#${editorId} .line:nth-child(${lineNumber})`);
     if (clickedLine) {
-        clickedLine.classList.add('highlight');
+        clickedLine.classList.add('highlight-source');
         clickedLine.scrollIntoView({
             behavior: 'smooth',
             block: 'center',
@@ -153,40 +365,96 @@ function handleLineClick(editorId, lineNumber) {
         });
     }
 
-    // Highlight corresponding lines
+    currentSelection = { editorId, lineNumber };
     highlightCorrespondingLines(editorId, lineNumber);
 }
 
-// Initialize data from pre-embedded content
+// ============== Syntax Highlighting ==============
+
+/**
+ * Highlight code using highlight.js and return the highlighted HTML.
+ * Falls back to auto-detection if the specified language fails,
+ * and escapes HTML if all highlighting attempts fail.
+ * @param {string} code - Raw code to highlight
+ * @param {string} language - Language identifier (e.g., 'python', 'cpp')
+ * @returns {string} HTML string with syntax highlighting spans
+ */
+function highlightCode(code, language) {
+    if (typeof hljs === 'undefined') {
+        console.warn('highlight.js not loaded, using plain text');
+        return escapeHtml(code);
+    }
+    try {
+        const result = hljs.highlight(code, { language });
+        return result.value;
+    } catch (e) {
+        console.warn(`Syntax highlighting failed for language '${language}': ${e.message}. Trying auto-detection.`);
+        try {
+            const result = hljs.highlightAuto(code);
+            return result.value;
+        } catch (e2) {
+            console.error(`Syntax highlighting completely failed: ${e2.message}. Falling back to plain text.`);
+            return escapeHtml(code);
+        }
+    }
+}
+
+/**
+ * Escape HTML special characters to prevent injection.
+ * @param {string} text - Text to escape
+ * @returns {string} HTML-escaped string
+ */
+function escapeHtml(text) {
+    const div = document.createElement('div');
+    div.textContent = text;
+    return div.innerHTML;
+}
+
+// ============== Initialization ==============
+
+/**
+ * Initialize data from pre-embedded content in the HTML.
+ * Sets up all three editor panels with syntax-highlighted content.
+ */
 function initializeData() {
     try {
-        // Get content from pre tags
-        const preGradGraph = document.querySelector('#preGradGraph pre');
-        const postGradGraph = document.querySelector('#postGradGraph pre');
-        const generatedCode = document.querySelector('#generatedCode pre');
+        const preGradGraph = document.querySelector('#preGradGraph pre code');
+        const postGradGraph = document.querySelector('#postGradGraph pre code');
+        const generatedCode = document.querySelector('#generatedCode pre code');
+
+        // Initialize line mappings first
+        initializeLineMappings();
 
-        if (preGradGraph) preGradGraphData = preGradGraph.textContent.split('\n');
-        if (postGradGraph) postGradGraphData = postGradGraph.textContent.split('\n');
+        // Get raw text content and highlight with highlight.js
+        if (preGradGraph) {
+            const rawCode = preGradGraph.textContent;
+            const highlighted = highlightCode(rawCode, 'python');
+            preGradGraphData = splitHtmlByLines(highlighted);
+        }
+        if (postGradGraph) {
+            const rawCode = postGradGraph.textContent;
+            const highlighted = highlightCode(rawCode, 'python');
+            postGradGraphData = splitHtmlByLines(highlighted);
+        }
         if (generatedCode) {
-            const content = generatedCode.textContent;
-            if (content.includes('AOTInductorModel::run_impl')) {
-                // This is C++ code
-                cppCodeData = content.split('\n');
+            const rawCode = generatedCode.textContent;
+            const isCpp = rawCode.includes('AOTInductorModel::run_impl');
+            const language = isCpp ? 'cpp' : 'python';
+            const highlighted = highlightCode(rawCode, language);
+
+            if (isCpp) {
+                cppCodeData = splitHtmlByLines(highlighted);
                 codeData = null;
             } else {
-                // This is Python code
-                codeData = content.split('\n');
+                codeData = splitHtmlByLines(highlighted);
                 cppCodeData = null;
             }
         }
 
-        // Initialize line mappings from pre-processed data
-        initializeLineMappings();
-
-        // Setup highlighting
-        setupEditorContent('preGradGraph', preGradGraphData);
-        setupEditorContent('postGradGraph', postGradGraphData);
-        setupEditorContent('generatedCode', codeData || cppCodeData);
+        // Setup editors with highlighted content
+        setupEditorContent('preGradGraph', preGradGraphData, true);
+        setupEditorContent('postGradGraph', postGradGraphData, true);
+        setupEditorContent('generatedCode', codeData || cppCodeData, true);
 
         // If it's C++ code, scroll to run_impl
         if (cppCodeData) {
@@ -203,32 +471,40 @@ function initializeData() {
     } catch (error) {
         console.error('Error initializing data:', error);
         console.error(error.stack);
+
+        // Show error to user
+        const container = document.querySelector('.editor-container');
+        if (container) {
+            container.innerHTML = `
+                <div style="padding: 20px; color: #c62828; background: #ffebee; margin: 10px; border-radius: 4px;">
+                    <h3 style="margin-top: 0;">Failed to initialize provenance viewer</h3>
+                    <p>Error: ${escapeHtml(error.message)}</p>
+                    <p>This may be caused by corrupted data or browser issues. Check the console for details.</p>
+                </div>
+            `;
+        }
     }
 }
 
-// Call initialization when the page loads
-window.addEventListener('DOMContentLoaded', initializeData);
-
-// Highlight corresponding lines
+/**
+ * Highlight corresponding lines in other editors based on line mappings.
+ * @param {string} sourceEditorId - The source editor ID
+ * @param {number} lineNumber - The 1-based line number in the source
+ */
 function highlightCorrespondingLines(sourceEditorId, lineNumber) {
     let correspondingLines = findCorrespondingLines(sourceEditorId, lineNumber);
-    
+
     Object.entries(correspondingLines).forEach(([editorId, lines]) => {
-        // Skip scrolling if this is the source editor
         if (lines && editorId !== sourceEditorId) {
-            // Handle both single numbers and arrays of numbers
             const lineNumbers = Array.isArray(lines) ? lines : [lines];
-            
-            // Get the middle line number for scrolling
             const middleIndex = Math.floor(lineNumbers.length / 2);
             let hasScrolled = false;
-            
+
             lineNumbers.forEach((line, index) => {
                 const lineElement = document.querySelector(`#${editorId} .line:nth-child(${line})`);
                 if (lineElement) {
-                    lineElement.classList.add('highlight');
-                    
-                    // Scroll to the middle line of the highlighted range
+                    lineElement.classList.add('highlight-target');
+
                     if (index === middleIndex && !hasScrolled) {
                         lineElement.scrollIntoView({
                             behavior: 'smooth',
@@ -243,10 +519,15 @@ function highlightCorrespondingLines(sourceEditorId, lineNumber) {
     });
 }
 
-// Given a line in sourceEditorId, find the corresponding lines in the other editors that should be highlighted.
+/**
+ * Find corresponding lines for a given source line based on mappings.
+ * @param {string} sourceEditorId - The source editor ID
+ * @param {number} lineNumber - The 1-based line number
+ * @returns {Object} Object with editor IDs as keys and line number arrays as values
+ */
 function findCorrespondingLines(sourceEditorId, lineNumber) {
     let result = {};
-    
+
     switch (sourceEditorId) {
         case 'preGradGraph':
             result.postGradGraph = preToPost[lineNumber] || [];
@@ -265,19 +546,18 @@ function findCorrespondingLines(sourceEditorId, lineNumber) {
                 }
             }
             break;
-            
+
         case 'postGradGraph':
             result.preGradGraph = postToPre[lineNumber] || [];
             if (codeData) {
                 result.generatedCode = postToPyCode[lineNumber] || [];
             } else {
-                result.generatedCode =  postToCppCode[lineNumber] || [];
+                result.generatedCode = postToCppCode[lineNumber] || [];
             }
             break;
-            
+
         case 'generatedCode':
             if (codeData) {
-                // Python code
                 result.postGradGraph = pyCodeToPost[lineNumber] || [];
             } else {
                 result.postGradGraph = cppCodeToPost[lineNumber] || [];
@@ -292,19 +572,190 @@ function findCorrespondingLines(sourceEditorId, lineNumber) {
             }
             break;
     }
-    
+
     return result;
 }
 
-// Resizable Panels Start
+// ============== Keyboard Navigation ==============
+
+/**
+ * Navigate to the next or previous mapped line in the current editor.
+ * @param {number} direction - 1 for next, -1 for previous
+ */
+function navigateToNextMappedLine(direction) {
+    if (!currentSelection.editorId) {
+        // If no selection, start from first mapped line
+        const firstMapped = document.querySelector('.line.has-match');
+        if (firstMapped) {
+            const editorElement = firstMapped.closest('.editor');
+            if (!editorElement || !editorElement.id) {
+                console.warn('navigateToNextMappedLine: Could not find parent editor element');
+                return;
+            }
+            const lineNum = parseInt(firstMapped.dataset.lineNumber, 10);
+            if (isNaN(lineNum)) {
+                console.warn('navigateToNextMappedLine: Invalid line number in dataset');
+                return;
+            }
+            handleLineClick(editorElement.id, lineNum);
+        }
+        return;
+    }
+
+    const editor = document.getElementById(currentSelection.editorId);
+    if (!editor) return;
+
+    const lines = Array.from(editor.querySelectorAll('.line'));
+    const startIndex = currentSelection.lineNumber - 1;
+
+    // Search in the specified direction
+    if (direction > 0) {
+        for (let i = startIndex + 1; i < lines.length; i++) {
+            if (lines[i].classList.contains('has-match')) {
+                handleLineClick(currentSelection.editorId, i + 1);
+                return;
+            }
+        }
+        // Wrap around
+        for (let i = 0; i < startIndex; i++) {
+            if (lines[i].classList.contains('has-match')) {
+                handleLineClick(currentSelection.editorId, i + 1);
+                return;
+            }
+        }
+    } else {
+        for (let i = startIndex - 1; i >= 0; i--) {
+            if (lines[i].classList.contains('has-match')) {
+                handleLineClick(currentSelection.editorId, i + 1);
+                return;
+            }
+        }
+        // Wrap around
+        for (let i = lines.length - 1; i > startIndex; i--) {
+            if (lines[i].classList.contains('has-match')) {
+                handleLineClick(currentSelection.editorId, i + 1);
+                return;
+            }
+        }
+    }
+}
+
+/**
+ * Jump to corresponding line in the next/previous panel.
+ * @param {number} direction - 1 for next panel, -1 for previous
+ */
+function jumpToCorrespondingPanel(direction) {
+    if (!currentSelection.editorId || !currentSelection.lineNumber) return;
+
+    const panels = ['preGradGraph', 'postGradGraph', 'generatedCode'];
+    const currentIndex = panels.indexOf(currentSelection.editorId);
+    const nextIndex = (currentIndex + direction + panels.length) % panels.length;
+
+    const correspondingLines = findCorrespondingLines(currentSelection.editorId, currentSelection.lineNumber);
+    const nextPanel = panels[nextIndex];
+
+    if (correspondingLines[nextPanel] && correspondingLines[nextPanel].length > 0) {
+        const targetLine = correspondingLines[nextPanel][0];
+        handleLineClick(nextPanel, targetLine);
+    } else {
+        // No corresponding line, just focus the panel
+        const editor = document.getElementById(nextPanel);
+        if (editor) {
+            const firstLine = editor.querySelector('.line');
+            if (firstLine) {
+                handleLineClick(nextPanel, 1);
+            }
+        }
+    }
+}
+
+/**
+ * Toggle help overlay visibility.
+ */
+function toggleHelp() {
+    const help = document.getElementById('keyboard-help');
+    if (help) {
+        help.classList.toggle('visible');
+    }
+}
+
+/**
+ * Set up keyboard event handlers for navigation shortcuts.
+ */
+function setupKeyboardNavigation() {
+    document.addEventListener('keydown', (e) => {
+        // Ignore if typing in an input
+        if (e.target.tagName === 'INPUT' || e.target.tagName === 'TEXTAREA') {
+            return;
+        }
+
+        switch (e.key) {
+            case 'n':
+                navigateToNextMappedLine(1);
+                e.preventDefault();
+                break;
+            case 'N':
+                navigateToNextMappedLine(-1);
+                e.preventDefault();
+                break;
+            case 'j':
+                jumpToCorrespondingPanel(1);
+                e.preventDefault();
+                break;
+            case 'J':
+                jumpToCorrespondingPanel(-1);
+                e.preventDefault();
+                break;
+            case '?':
+                toggleHelp();
+                e.preventDefault();
+                break;
+            case 'Escape':
+                clearHighlights();
+                currentSelection = { editorId: null, lineNumber: null };
+                // Also close help if open
+                const help = document.getElementById('keyboard-help');
+                if (help) help.classList.remove('visible');
+                // Restore panels if maximized
+                if (maximizedPanel) {
+                    toggleMaximize(maximizedPanel);
+                }
+                e.preventDefault();
+                break;
+            case '1':
+            case '2':
+            case '3':
+                const panelIndex = parseInt(e.key) - 1;
+                const wrappers = document.querySelectorAll('.editor-wrapper');
+                if (wrappers[panelIndex]) {
+                    toggleMaximize(wrappers[panelIndex]);
+                }
+                e.preventDefault();
+                break;
+        }
+    });
+}
+
+// ============== Resizable Panels ==============
+
+/**
+ * Set up drag-to-resize functionality for panel dividers.
+ */
 function setupResizablePanels() {
     const container = document.querySelector('.editor-container');
-    const pre = document.getElementById('preGradGraph');
-    const post = document.getElementById('postGradGraph');
-    const code = document.getElementById('generatedCode');
+    const wrappers = document.querySelectorAll('.editor-wrapper');
     const divider1 = document.getElementById('divider1');
     const divider2 = document.getElementById('divider2');
 
+    if (wrappers.length < 3 || !divider1 || !divider2) {
+        console.warn('setupResizablePanels: Required elements missing', {
+            wrapperCount: wrappers.length,
+            hasDivider1: !!divider1,
+            hasDivider2: !!divider2
+        });
+        return;
+    }
+
     let isDragging = false;
     let dragDivider = null;
 
@@ -312,19 +763,24 @@ function setupResizablePanels() {
         if (!isDragging || !dragDivider) return;
 
         const containerRect = container.getBoundingClientRect();
-        const totalWidth = containerRect.width;
 
         if (dragDivider === divider1) {
-            const newPreWidth = e.clientX - containerRect.left;
-            const newPostWidth = post.offsetWidth + (pre.offsetWidth - newPreWidth);
-            pre.style.flex = `0 0 ${newPreWidth}px`;
-            post.style.flex = `0 0 ${newPostWidth}px`;
+            const newWidth = e.clientX - containerRect.left;
+            const minWidth = 100;
+            const maxWidth = containerRect.width - 200;
+
+            if (newWidth > minWidth && newWidth < maxWidth) {
+                wrappers[0].style.flex = `0 0 ${newWidth}px`;
+            }
         } else if (dragDivider === divider2) {
-            const preWidth = pre.offsetWidth;
-            const newPostWidth = e.clientX - containerRect.left - preWidth - divider1.offsetWidth;
-            const newCodeWidth = totalWidth - e.clientX + containerRect.left - divider2.offsetWidth;
-            post.style.flex = `0 0 ${newPostWidth}px`;
-            code.style.flex = `0 0 ${newCodeWidth}px`;
+            const firstWidth = wrappers[0].offsetWidth;
+            const newWidth = e.clientX - containerRect.left - firstWidth - divider1.offsetWidth;
+            const minWidth = 100;
+            const maxWidth = containerRect.width - firstWidth - divider1.offsetWidth - 100;
+
+            if (newWidth > minWidth && newWidth < maxWidth) {
+                wrappers[1].style.flex = `0 0 ${newWidth}px`;
+            }
         }
     }
 
@@ -343,10 +799,67 @@ function setupResizablePanels() {
             document.body.style.cursor = 'col-resize';
             document.addEventListener('mousemove', onMouseMove);
             document.addEventListener('mouseup', onMouseUp);
+            e.preventDefault();
+        });
+    });
+}
+
+// ============== Panel Maximize ==============
+
+/**
+ * Set up double-click on headers to maximize/restore panels.
+ */
+function setupPanelMaximize() {
+    const headers = document.querySelectorAll('.editor-header');
+    headers.forEach(header => {
+        header.style.cursor = 'pointer';
+        header.addEventListener('dblclick', () => {
+            const wrapper = header.parentElement;
+            toggleMaximize(wrapper);
         });
     });
 }
 
-window.addEventListener('DOMContentLoaded', setupResizablePanels);
+/**
+ * Toggle maximize state for a panel.
+ * @param {Element} wrapper - The editor-wrapper element to maximize/restore
+ */
+function toggleMaximize(wrapper) {
+    const allWrappers = document.querySelectorAll('.editor-wrapper');
+    const dividers = document.querySelectorAll('.divider');
+
+    if (maximizedPanel === wrapper) {
+        // Restore all panels
+        allWrappers.forEach(w => {
+            w.style.display = '';
+            w.style.flex = '';
+        });
+        dividers.forEach(d => {
+            d.style.display = '';
+        });
+        maximizedPanel = null;
+    } else {
+        // Maximize this panel, hide others
+        allWrappers.forEach(w => {
+            if (w === wrapper) {
+                w.style.display = '';
+                w.style.flex = '1';
+            } else {
+                w.style.display = 'none';
+            }
+        });
+        dividers.forEach(d => {
+            d.style.display = 'none';
+        });
+        maximizedPanel = wrapper;
+    }
+}
+
+// ============== Entry Point ==============
 
-// Resizable Panels End
\ No newline at end of file
+window.addEventListener('DOMContentLoaded', () => {
+    initializeData();
+    setupResizablePanels();
+    setupKeyboardNavigation();
+    setupPanelMaximize();
+});