Add panic hook, fix cpu legend causing panic on small windows

2 files changed, 46 insertions, 9 deletions

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