From e548d07c1f016165c206256f2c8bab1bfa4c38c5 Mon Sep 17 00:00:00 2001
From: ClementTsang <cjhtsang@uwaterloo.ca>
Date: Sun, 2 Feb 2020 17:09:42 -0500
Subject: [PATCH] Add panic hook, fix cpu legend causing panic on small windows

---
 Cargo.toml    |  5 +++--
 src/canvas.rs |  2 +-
 src/main.rs   | 53 +++++++++++++++++++++++++++++++++++++++++++--------
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 979e1183..cc7a236b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -29,12 +29,13 @@ futures-timer = "2.0.2"
 futures = "0.3.1"
 heim = "0.0.9"
 log = "0.4.8"
-regex = "1.3.3"
+regex = "1.3.4"
 sysinfo = "0.9.6" #0.9 seems to be the last working version for my Ryzen PC...
-tokio = "0.2.9"
+tokio = "0.2.11"
 winapi = "0.3.8"
 tui = {version = "0.8", features = ["crossterm"], default-features = false }
 lazy_static = "1.4.0"
+backtrace = "0.3"
 
 [dev-dependencies]
 assert_cmd = "0.12"
diff --git a/src/canvas.rs b/src/canvas.rs
index 0e28255e..58d8a09e 100644
--- a/src/canvas.rs
+++ b/src/canvas.rs
@@ -461,7 +461,7 @@ fn draw_cpu_legend<B: backend::Backend>(
 ) {
 	let cpu_data: &[ConvertedCpuData] = &(app_state.canvas_data.cpu_data);
 
-	let num_rows = u64::from(draw_loc.height) - 5;
+	let num_rows = max(0, i64::from(draw_loc.height) - 5) as u64;
 	let start_position = get_start_position(
 		num_rows,
 		&(app_state.scroll_direction),
diff --git a/src/main.rs b/src/main.rs
index 03d12705..c328f6f0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -13,12 +13,14 @@ use crossterm::{
 		KeyModifiers, MouseEvent,
 	},
 	execute,
+	style::Print,
 	terminal::LeaveAlternateScreen,
 	terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen},
 };
 
 use std::{
 	io::{stdout, Write},
+	panic::{self, PanicInfo},
 	sync::mpsc,
 	thread,
 	time::{Duration, Instant},
@@ -135,15 +137,18 @@ fn main() -> error::Result<()> {
 	}
 
 	// Set up up tui and crossterm
-	let mut stdout = stdout();
+	let mut stdout_val = stdout();
 	enable_raw_mode()?;
-	execute!(stdout, EnterAlternateScreen)?;
-	execute!(stdout, EnableMouseCapture)?;
+	execute!(stdout_val, EnterAlternateScreen)?;
+	execute!(stdout_val, EnableMouseCapture)?;
 
-	let mut terminal = Terminal::new(CrosstermBackend::new(stdout))?;
+	let mut terminal = Terminal::new(CrosstermBackend::new(stdout_val))?;
 	terminal.hide_cursor()?;
 	terminal.clear()?;
 
+	// Set panic hook
+	panic::set_hook(Box::new(|info| panic_hook(info)));
+
 	// Set up input handling
 	let (tx, rx) = mpsc::channel();
 	{
@@ -332,27 +337,59 @@ fn main() -> error::Result<()> {
 
 		// Draw!
 		if let Err(err) = canvas::draw_data(&mut terminal, &mut app) {
-			cleanup(&mut terminal)?;
+			cleanup_terminal(&mut terminal)?;
 			error!("{}", err);
 			return Err(err);
 		}
 	}
 
-	cleanup(&mut terminal)?;
+	cleanup_terminal(&mut terminal)?;
 	Ok(())
 }
 
-fn cleanup(
+fn cleanup_terminal(
 	terminal: &mut tui::terminal::Terminal<tui::backend::CrosstermBackend<std::io::Stdout>>,
 ) -> error::Result<()> {
 	disable_raw_mode()?;
-	execute!(terminal.backend_mut(), DisableMouseCapture)?;
 	execute!(terminal.backend_mut(), LeaveAlternateScreen)?;
+	execute!(terminal.backend_mut(), DisableMouseCapture)?;
 	terminal.show_cursor()?;
 
 	Ok(())
 }
 
+/// Based on https://github.com/Rigellute/spotify-tui/blob/master/src/main.rs
+fn panic_hook(panic_info: &PanicInfo<'_>) {
+	let mut stdout = stdout();
+
+	if cfg!(debug_assertions) {
+		let msg = match panic_info.payload().downcast_ref::<&'static str>() {
+			Some(s) => *s,
+			None => match panic_info.payload().downcast_ref::<String>() {
+				Some(s) => &s[..],
+				None => "Box<Any>",
+			},
+		};
+
+		let stacktrace: String = format!("{:?}", backtrace::Backtrace::new());
+
+		execute!(
+			stdout,
+			Print(format!(
+				"thread '<unnamed>' panicked at '{}', {}\n\r{}",
+				msg,
+				panic_info.location().unwrap(),
+				stacktrace
+			)),
+		)
+		.unwrap();
+	}
+
+	disable_raw_mode().unwrap();
+	execute!(stdout, LeaveAlternateScreen).unwrap();
+	execute!(stdout, DisableMouseCapture).unwrap();
+}
+
 fn update_final_process_list(app: &mut app::App) {
 	let mut filtered_process_data: Vec<ConvertedProcessData> = if app.is_grouped() {
 		app.canvas_data