From 183ecb9b78a1c22d832e979db5054dcac36d8b7a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kat=20March=C3=A1n?= <kzm@zkat.tech>
Date: Tue, 21 Sep 2021 12:27:43 -0700
Subject: [PATCH] feat(panic): add backtrace support to panic handler and move
 set_panic_hook into fancy features

---
 Cargo.toml   |  4 +++-
 src/panic.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index d34ab7a2..d5b84bbf 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +26,7 @@ supports-hyperlinks = { version = "1.1.0", optional = true }
 supports-color = { version = "1.0.2", optional = true }
 supports-unicode = { version = "1.0.0", optional = true }
 itertools = { version = "0.10.1", optional = true }
+backtrace = { version = "0.3.61", optional = true }
 
 [dev-dependencies]
 semver = "1.0.4"
@@ -49,7 +50,8 @@ fancy = [
     "supports-hyperlinks",
     "supports-color",
     "supports-unicode",
-    "itertools"
+    "itertools",
+    "backtrace"
 ]
 
 [workspace]
diff --git a/src/panic.rs b/src/panic.rs
index a66b81f8..baf4e0c9 100644
--- a/src/panic.rs
+++ b/src/panic.rs
@@ -1,10 +1,14 @@
+#![cfg(feature = "fancy")]
+use std::fmt::Write;
+
+use backtrace::Backtrace;
 use thiserror::Error;
 
 use crate::{self as miette, Context, Diagnostic, Result};
 
 /// Tells miette to render panics using its rendering engine.
 pub fn set_panic_hook() {
-    std::panic::set_hook(Box::new(|info| {
+    std::panic::set_hook(Box::new(move |info| {
         let mut message = "Something went wrong".to_string();
         let payload = info.payload();
         if let Some(msg) = payload.downcast_ref::<&str>() {
@@ -25,6 +29,59 @@ pub fn set_panic_hook() {
 }
 
 #[derive(Debug, Error, Diagnostic)]
-#[error("{0}")]
+#[error("{0}{}", self.maybe_collect_backtrace())]
 #[diagnostic(help("set the `RUST_BACKTRACE=1` environment variable to display a backtrace."))]
 struct Panic(String);
+
+impl Panic {
+    fn maybe_collect_backtrace(&self) -> String {
+        if let Ok(var) = std::env::var("RUST_BACKTRACE") {
+            if !var.is_empty() && var != "0" {
+                // This is all taken from human-panic: https://github.com/rust-cli/human-panic/blob/master/src/report.rs#L55-L107
+                const HEX_WIDTH: usize = std::mem::size_of::<usize>() + 2;
+                //Padding for next lines after frame's address
+                const NEXT_SYMBOL_PADDING: usize = HEX_WIDTH + 6;
+                let mut backtrace = String::new();
+                for (idx, frame) in Backtrace::new().frames().iter().skip(26).enumerate() {
+                    let ip = frame.ip();
+                    let _ = write!(backtrace, "\n{:4}: {:2$?}", idx, ip, HEX_WIDTH);
+
+                    let symbols = frame.symbols();
+                    if symbols.is_empty() {
+                        let _ = write!(backtrace, " - <unresolved>");
+                        continue;
+                    }
+
+                    for (idx, symbol) in symbols.iter().enumerate() {
+                        //Print symbols from this address,
+                        //if there are several addresses
+                        //we need to put it on next line
+                        if idx != 0 {
+                            let _ = write!(backtrace, "\n{:1$}", "", NEXT_SYMBOL_PADDING);
+                        }
+
+                        if let Some(name) = symbol.name() {
+                            let _ = write!(backtrace, " - {}", name);
+                        } else {
+                            let _ = write!(backtrace, " - <unknown>");
+                        }
+
+                        //See if there is debug information with file name and line
+                        if let (Some(file), Some(line)) = (symbol.filename(), symbol.lineno()) {
+                            let _ = write!(
+                                backtrace,
+                                "\n{:3$}at {}:{}",
+                                "",
+                                file.display(),
+                                line,
+                                NEXT_SYMBOL_PADDING
+                            );
+                        }
+                    }
+                }
+                return backtrace;
+            }
+        }
+        "".into()
+    }
+}