Add a couple of extra features

src/lib.rs

@@ -7,63 +7,72 @@ use std::{
 use regex::Regex;
 use termcolor::{Color, ColorSpec, StandardStream, WriteColor};
 
-#[derive(Debug)]
 pub struct Backtrace {
     frames: Vec<Frame>,
+    panic_info: Option<PanicInfo>,
+}
+
+struct PanicInfo {
+    thread: String,
+    source: SourceInfo,
+    message: Vec<String>,
 }
 
-#[derive(Debug)]
 pub struct Frame {
     function: String,
     frameno: u32,
     source_info: Option<SourceInfo>,
 }
 
-#[derive(Debug)]
 pub struct SourceInfo {
     file: String,
     lineno: usize,
     colno: usize,
 }
 
 impl Backtrace {
-    pub fn render(&self, stdout: &mut StandardStream) -> io::Result<()> {
+    pub fn render(&self, out: &mut StandardStream) -> io::Result<()> {
         if self.frames.is_empty() {
             return Ok(());
         }
         let frameno_width = self.frames.len().ilog10() as usize + 1;
 
         for frame in self.frames.iter().rev() {
-            frame.render(stdout, frameno_width)?;
+            frame.render(out, frameno_width)?;
         }
-        Ok(())
+
+        if let Some(panic_info) = &self.panic_info {
+            panic_info.render(out)?;
+        }
+
+        writeln!(out)
     }
 }
 
 impl Frame {
-    fn render(&self, stdout: &mut StandardStream, width: usize) -> io::Result<()> {
-        write!(stdout, "{:>width$}: ", self.frameno)?;
+    fn render(&self, out: &mut StandardStream, width: usize) -> io::Result<()> {
+        write!(out, "{:>width$}: ", self.frameno)?;
 
-        stdout.set_color(ColorSpec::new().set_fg(Some(Color::Green)))?;
-        writeln!(stdout, "{}", self.function)?;
-        stdout.set_color(&ColorSpec::new())?;
+        out.set_color(ColorSpec::new().set_fg(Some(Color::Green)))?;
+        writeln!(out, "{}", self.function)?;
+        out.set_color(&ColorSpec::new())?;
 
         if let Some(source_info) = &self.source_info {
-            source_info.render(stdout, width)?;
+            source_info.render(out, width)?;
         }
         Ok(())
     }
 }
 
 impl SourceInfo {
-    fn render(&self, stdout: &mut StandardStream, width: usize) -> io::Result<()> {
-        write!(stdout, "{:width$}  at ", "")?;
-        writeln!(stdout, "{}:{}:{}", self.file, self.lineno, self.colno)?;
-        self.render_code(stdout, width)?;
+    fn render(&self, out: &mut StandardStream, width: usize) -> io::Result<()> {
+        write!(out, "{:width$}  at ", "")?;
+        writeln!(out, "{self}")?;
+        self.render_code(out, width)?;
         Ok(())
     }
 
-    fn render_code(&self, stdout: &mut StandardStream, width: usize) -> io::Result<()> {
+    fn render_code(&self, out: &mut StandardStream, width: usize) -> io::Result<()> {
         let path = Path::new(&self.file);
         if path.exists() {
             let lineno = self.lineno - 1;
@@ -76,86 +85,179 @@ impl SourceInfo {
                 .take(5);
             for (i, line) in viewport {
                 if i == lineno {
-                    stdout.set_color(ColorSpec::new().set_bold(true))?;
+                    out.set_color(ColorSpec::new().set_bold(true))?;
                 }
-                write!(stdout, "{:width$}    {} | ", "", i + 1)?;
-                writeln!(stdout, "{}", line?)?;
+                write!(out, "{:width$}    {} | ", "", i + 1)?;
+                writeln!(out, "{}", line?)?;
                 if i == lineno {
-                    stdout.set_color(ColorSpec::new().set_bold(false))?;
+                    out.set_color(ColorSpec::new().set_bold(false))?;
                 }
             }
         }
         Ok(())
     }
 }
 
+impl std::fmt::Display for SourceInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}:{}", self.file, self.lineno, self.colno)
+    }
+}
+
+impl PanicInfo {
+    fn render(&self, out: &mut StandardStream) -> io::Result<()> {
+        out.set_color(ColorSpec::new().set_fg(Some(Color::Red)))?;
+        writeln!(out, "thread '{}' panicked at {}", self.thread, self.source)?;
+        out.set_color(&ColorSpec::new())?;
+
+        for line in &self.message {
+            writeln!(out, ">> {}", line)?;
+        }
+        Ok(())
+    }
+}
+
 pub struct Parser {
-    re1: Regex,
-    re2: Regex,
-    lines: Vec<BacktraceLine>,
+    panic_regex: Regex,
+    function_regex: Regex,
+    source_regex: Regex,
+    lines: Vec<ParsedLine>,
 }
 
-enum BacktraceLine {
-    Function { function: String, frameno: u32 },
-    Source(SourceInfo),
+enum ParsedLine {
+    /// A line reporting a panic, e.g.,
+    /// ```ignore
+    /// thread 'rustc' panicked at /rustc/b3aa8e7168a3d940122db3561289ffbf3f587262/compiler/rustc_errors/src/lib.rs:1651:9:
+    /// ```
+    ThreadPanic { thread: String, source: SourceInfo },
+    /// The begining of a trace starts with `stack backtrace:`
+    BacktraceStart,
+    /// The "header" of a frame containing the frame number and the function's name, e.g.,
+    /// ```ignore
+    ///   28: rustc_middle::ty::context::tls::enter_context`
+    /// ```
+    BacktraceHeader { function: String, frameno: u32 },
+    /// Line containing source information about a frame, e.g.,
+    /// ```ignore
+    ///              at /rustc/b3aa8e7168a3d940122db3561289ffbf3f587262/compiler/rustc_middle/src/ty/context/tls.rs:79:9
+    /// ```
+    BacktraceSource(SourceInfo),
+    /// A line that doesn't match any patter
+    Other(String),
 }
 
 impl Parser {
     pub fn new() -> Parser {
-        let re1 = Regex::new(r"^\s+(?P<frameno>\d+):\s+((\w+)\s+-\s+)?(?P<function>.+)").unwrap();
-        let re2 = Regex::new(r"^\s+at\s+(?P<file>[^:]+):(?P<lineno>\d+):(?P<colno>\d+)").unwrap();
+        let panic_regex = Regex::new(r"^thread\s+'(?P<thread>[^']+)'\spanicked\s+at\s+(?P<file>[^:]+):(?P<lineno>\d+):(?P<colno>\d+)").unwrap();
+        let function_regex =
+            Regex::new(r"^\s+(?P<frameno>\d+):\s+((\w+)\s+-\s+)?(?P<function>.+)").unwrap();
+        let source_regex =
+            Regex::new(r"^\s+at\s+(?P<file>[^:]+):(?P<lineno>\d+):(?P<colno>\d+)").unwrap();
         Parser {
-            re1,
-            re2,
+            panic_regex,
+            function_regex,
+            source_regex,
             lines: vec![],
         }
     }
 
-    pub fn parse_line(&mut self, line: &str) -> bool {
-        if let Some(captures) = self.re1.captures(&line) {
+    pub fn parse_line(&mut self, line: String) {
+        let parsed = if line.eq_ignore_ascii_case("stack backtrace:") {
+            ParsedLine::BacktraceStart
+        } else if let Some(captures) = self.panic_regex.captures(&line) {
+            let thread = captures.name("thread").unwrap().as_str().to_string();
+            let file = captures.name("file").unwrap().as_str().to_string();
+            let lineno = captures.name("lineno").unwrap().as_str();
+            let colno = captures.name("colno").unwrap().as_str();
+            ParsedLine::ThreadPanic {
+                thread,
+                source: SourceInfo {
+                    file,
+                    lineno: lineno.parse().unwrap(),
+                    colno: colno.parse().unwrap(),
+                },
+            }
+        } else if let Some(captures) = self.function_regex.captures(&line) {
             let frameno = captures.name("frameno").unwrap().as_str().to_string();
             let function = captures.name("function").unwrap().as_str().to_string();
-            self.lines.push(BacktraceLine::Function {
+            ParsedLine::BacktraceHeader {
                 function,
                 frameno: frameno.parse().unwrap(),
-            });
-            true
-        } else if let Some(captures) = self.re2.captures(&line) {
+            }
+        } else if let Some(captures) = self.source_regex.captures(&line) {
             let file = captures.name("file").unwrap().as_str().to_string();
             let lineno = captures.name("lineno").unwrap().as_str();
             let colno = captures.name("colno").unwrap().as_str();
-            self.lines.push(BacktraceLine::Source(SourceInfo {
+            ParsedLine::BacktraceSource(SourceInfo {
                 file,
                 lineno: lineno.parse().unwrap(),
                 colno: colno.parse().unwrap(),
-            }));
-            true
+            })
         } else {
-            false
-        }
+            ParsedLine::Other(line)
+        };
+        self.lines.push(parsed)
     }
 
-    pub fn into_backtrace(self) -> Backtrace {
+    pub fn into_backtraces(self) -> Vec<Backtrace> {
+        let mut backtraces = vec![];
         let mut frames = vec![];
         let mut lines = self.lines.into_iter().peekable();
+        let mut panic_info = None;
+        let mut in_panic_info = false;
         while let Some(line) = lines.next() {
-            if let BacktraceLine::Function { function, frameno } = line {
-                let source_info = lines
-                    .next_if(|line| matches!(line, BacktraceLine::Source(..)))
-                    .and_then(|line| {
-                        if let BacktraceLine::Source(source_info) = line {
-                            Some(source_info)
-                        } else {
-                            None
-                        }
+            match line {
+                ParsedLine::ThreadPanic { thread, source } => {
+                    in_panic_info = true;
+                    panic_info = Some(PanicInfo {
+                        thread,
+                        source,
+                        message: vec![],
                     });
-                frames.push(Frame {
-                    function,
-                    frameno,
-                    source_info,
-                })
+                }
+                ParsedLine::Other(line) => {
+                    if let Some(panic_info) = &mut panic_info {
+                        if in_panic_info {
+                            panic_info.message.push(line);
+                        }
+                    }
+                }
+                ParsedLine::BacktraceStart => {
+                    in_panic_info = false;
+                    if !frames.is_empty() {
+                        backtraces.push(Backtrace {
+                            frames: std::mem::take(&mut frames),
+                            panic_info: std::mem::take(&mut panic_info),
+                        });
+                    }
+                }
+                ParsedLine::BacktraceHeader { function, frameno } => {
+                    in_panic_info = false;
+                    let source_info = lines
+                        .next_if(|line| matches!(line, ParsedLine::BacktraceSource(..)))
+                        .and_then(|line| {
+                            if let ParsedLine::BacktraceSource(source_info) = line {
+                                Some(source_info)
+                            } else {
+                                None
+                            }
+                        });
+                    frames.push(Frame {
+                        function,
+                        frameno,
+                        source_info,
+                    })
+                }
+                ParsedLine::BacktraceSource(..) => {
+                    // This case is in theory never reached because source lines should be consumed
+                    // in the `BacktraceHeader` case.
+                    in_panic_info = false;
+                }
             }
         }
-        Backtrace { frames }
+        if !frames.is_empty() {
+            backtraces.push(Backtrace { frames, panic_info });
+        }
+        backtraces
     }
 }

src/main.rs

@@ -4,45 +4,63 @@ use std::process::{Command, Stdio};
 use clap::Parser;
 use termcolor::{ColorChoice, StandardStream};
 
-#[derive(clap::Parser, Debug)]
-#[command(author, version, about, long_about = None)]
+#[derive(clap::Parser)]
+#[command(version, about, long_about = None)]
 struct Args {
     #[arg(trailing_var_arg(true))]
     cmd: Vec<String>,
 
-    #[arg(long, short, default_value = "1")]
-    backtrace: String,
+    #[arg(long, short, default_value = "short")]
+    style: BacktraceStyle,
+
+    #[arg(long, short)]
+    no_lib_backtrace: bool,
+}
+
+#[derive(clap::ValueEnum, Copy, Clone)]
+enum BacktraceStyle {
+    Short,
+    Full,
 }
 
+impl BacktraceStyle {
+    fn env_var_str(&self) -> &'static str {
+        match self {
+            BacktraceStyle::Short => "1",
+            BacktraceStyle::Full => "full",
+        }
+    }
+}
 fn main() -> io::Result<()> {
     let mut args = Args::parse();
 
     if args.cmd.len() == 0 {
         std::process::exit(1);
     }
 
-    let mut cmd = Command::new(args.cmd.remove(0));
-    for arg in args.cmd {
-        cmd.arg(arg);
+    let mut env_vars = vec![("RUST_BACKTRACE", args.style.env_var_str())];
+    if args.no_lib_backtrace {
+        env_vars.push(("RUST_LIB_BACKTRACE", "0"));
     }
-    let child = cmd
+
+    let child = Command::new(args.cmd.remove(0))
+        .args(args.cmd)
         .stderr(Stdio::piped())
-        .env("RUST_BACKTRACE", args.backtrace)
-        .spawn()
-        .expect("Failed to execute command");
+        .envs(env_vars)
+        .spawn()?;
 
     let mut parser = backtracetk::Parser::new();
-    let stderr = child.stderr.expect("Failed to open stderr");
+    let stderr = child.stderr.expect("failed to open stderr");
     for line in BufReader::new(stderr).lines() {
         let line = line?;
-        if !parser.parse_line(&line) {
-            eprintln!("{line}");
-        }
+        eprintln!("{line}");
+        parser.parse_line(line);
     }
-    let backtrace = parser.into_backtrace();
 
     let mut stderr = StandardStream::stderr(ColorChoice::Auto);
-    backtrace.render(&mut stderr)?;
+    for backtrace in parser.into_backtraces() {
+        backtrace.render(&mut stderr)?;
+    }
 
     Ok(())
 }