cubecl_common/
backtrace.rs

1#[cfg(feature = "std")]
2type BacktraceState = backtrace_std::BacktraceState;
3#[cfg(not(feature = "std"))]
4type BacktraceState = alloc::string::String;
5
6/// Contains the backtrace information if available.
7///
8/// # Notes
9///
10/// We chose BackTrace for the name since Backtrace is often confused with the nighly-only backtrace
11/// feature by `thiserror`.
12#[derive(Clone, Default)]
13pub struct BackTrace {
14    state: Option<BacktraceState>,
15}
16
17impl core::fmt::Debug for BackTrace {
18    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
19        f.write_fmt(format_args!("{self}"))
20    }
21}
22
23impl core::fmt::Display for BackTrace {
24    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
25        match &self.state {
26            Some(state) => f.write_fmt(format_args!("{state}")),
27            None => f.write_str("No backtrace available"),
28        }
29    }
30}
31
32impl BackTrace {
33    /// Creates a new backtrace from the current thread.
34    ///
35    /// # Notes
36    ///
37    /// It is quite cheap to create a backtrace, but quite expensive to display.
38    pub fn capture() -> Self {
39        Self {
40            #[cfg(feature = "std")]
41            state: {
42                // We only resolve the backtrace when displaying the result.
43                //
44                // Making it cheaper to create.
45                //
46                // We capture the backtrace here to reduce the number of frames to ignore.
47                let backtrace = backtrace::Backtrace::new_unresolved();
48                Some(BacktraceState::new(backtrace))
49            },
50            #[cfg(not(feature = "std"))]
51            state: None,
52        }
53    }
54}
55
56#[cfg(feature = "std")]
57mod backtrace_std {
58    use backtrace::BytesOrWideString;
59    use core::fmt::Display;
60    use std::sync::{Arc, Mutex};
61
62    /// A modified version of [backtrace::BacktraceFmt] to skip the first capture.
63    #[derive(Clone)]
64    pub struct BacktraceState {
65        backtrace: Arc<Mutex<backtrace::Backtrace>>,
66    }
67
68    impl BacktraceState {
69        pub fn new(backtrace: backtrace::Backtrace) -> Self {
70            Self {
71                backtrace: Arc::new(Mutex::new(backtrace)),
72            }
73        }
74    }
75
76    impl Display for BacktraceState {
77        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
78            let mut backtrace = self.backtrace.lock().unwrap();
79            backtrace.resolve();
80
81            let cwd = std::env::current_dir();
82
83            let mut print_path =
84                move |fmt: &mut core::fmt::Formatter<'_>, path: BytesOrWideString<'_>| {
85                    let path = path.into_path_buf();
86                    if let Ok(cwd) = &cwd
87                        && let Ok(suffix) = path.strip_prefix(cwd)
88                    {
89                        return core::fmt::Display::fmt(&suffix.display(), fmt);
90                    }
91
92                    core::fmt::Display::fmt(&path.display(), fmt)
93                };
94
95            let mut fmt =
96                backtrace::BacktraceFmt::new(f, backtrace::PrintFmt::Short, &mut print_path);
97            fmt.add_context()?;
98            for frame in backtrace.frames().iter().skip(1) {
99                fmt.frame().backtrace_frame(frame)?;
100            }
101            fmt.finish()?;
102            Ok(())
103        }
104    }
105}
cubecl_common/backtrace.rs

cubecl_common/
backtrace.rs
