add support for custom demangler

CHANGELOG.md

@@ -6,13 +6,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Add `demangle` setter to the `ReportBuilder` to allow using custom demangler (#126)
+
 ### Changed
 - Remove `backtrace-rs` feature, as the default choice when not specified (#130)
 
 ## [0.9.1] - 2022-05-19
 
 ### Fixed
 - Protect the error number in signal handler (#128)
+- Support custom demangler (#126)
 
 ## [0.9.0] - 2022-05-09
 

src/frames.rs

@@ -86,7 +86,10 @@ impl Hash for UnresolvedFrames {
 #[derive(Debug, Clone)]
 pub struct Symbol {
     /// This name is raw name of a symbol (which hasn't been demangled).
-    pub name: Option<Vec<u8>>,
+    pub raw_name: Option<Vec<u8>>,
+
+    /// This name is demangled name of a symbol.
+    pub name: Option<String>,
 
     /// The address of the function. It is not 100% trustworthy.
     pub addr: Option<*mut c_void>,
@@ -100,11 +103,11 @@ pub struct Symbol {
 
 impl Symbol {
     pub fn raw_name(&self) -> &[u8] {
-        self.name.as_deref().unwrap_or(b"Unknown")
+        self.raw_name.as_deref().unwrap_or(b"Unknown")
     }
 
     pub fn name(&self) -> String {
-        demangle(&String::from_utf8_lossy(self.raw_name())).into_owned()
+        self.name.as_deref().unwrap_or("Unknown").to_string()
     }
 
     pub fn sys_name(&self) -> Cow<str> {
@@ -130,8 +133,24 @@ where
     T: crate::backtrace::Symbol,
 {
     fn from(symbol: &T) -> Self {
+        Self::from_with_demangle(symbol, demangle)
+    }
+}
+
+impl Symbol {
+    pub fn from_with_demangle<T, F>(symbol: &T, demangle: F) -> Self
+    where
+        T: crate::backtrace::Symbol,
+        F: Fn(&str) -> Cow<str>,
+    {
+        let raw_name = symbol.name();
+        let name = symbol
+            .name()
+            .map(|name| demangle(&String::from_utf8_lossy(&name)).to_string());
+
         Symbol {
-            name: symbol.name(),
+            raw_name,
+            name,
             addr: symbol.addr(),
             lineno: symbol.lineno(),
             filename: symbol.filename(),
@@ -179,6 +198,15 @@ impl Frames {
 
 impl From<UnresolvedFrames> for Frames {
     fn from(frames: UnresolvedFrames) -> Self {
+        Self::from_with_demangle(frames, demangle)
+    }
+}
+
+impl Frames {
+    pub fn from_with_demangle<F>(frames: UnresolvedFrames, demangle: F) -> Self
+    where
+        F: Fn(&str) -> Cow<str> + Copy,
+    {
         let mut fs = Vec::new();
 
         let mut frame_iter = frames.frames.iter();
@@ -187,7 +215,7 @@ impl From<UnresolvedFrames> for Frames {
             let mut symbols: Vec<Symbol> = Vec::new();
 
             frame.resolve_symbol(|symbol| {
-                let symbol = Symbol::from(symbol);
+                let symbol = Symbol::from_with_demangle(symbol, demangle);
                 symbols.push(symbol);
             });
 
@@ -239,31 +267,17 @@ mod tests {
 
     #[test]
     fn demangle_rust() {
-        let symbol = Symbol {
-            name: Some(b"_ZN3foo3barE".to_vec()),
-            addr: None,
-            lineno: None,
-            filename: None,
-        };
-
-        assert_eq!(&symbol.name(), "foo::bar")
+        let demangled = demangle("_ZN3foo3barE");
+        assert_eq!(demangled, "foo::bar")
     }
 
     #[test]
     fn demangle_cpp() {
         let name =
-            b"_ZNK3MapI10StringName3RefI8GDScriptE10ComparatorIS0_E16DefaultAllocatorE3hasERKS0_"
-                .to_vec();
-
-        let symbol = Symbol {
-            name: Some(name),
-            addr: None,
-            lineno: None,
-            filename: None,
-        };
+            "_ZNK3MapI10StringName3RefI8GDScriptE10ComparatorIS0_E16DefaultAllocatorE3hasERKS0_";
 
         assert_eq!(
-            &symbol.name(),
+            demangle(name),
             "Map<StringName, Ref<GDScript>, Comparator<StringName>, DefaultAllocator>::has(StringName const&) const"
         )
     }

src/report.rs

@@ -1,9 +1,11 @@
 // Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
 
+use std::borrow::Cow;
 use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 
 use parking_lot::RwLock;
+use symbolic_demangle::demangle;
 
 use crate::frames::{Frames, UnresolvedFrames};
 use crate::profiler::Profiler;
@@ -32,6 +34,8 @@ pub struct UnresolvedReport {
 /// A builder of `Report` and `UnresolvedReport`. It builds report from a running `Profiler`.
 pub struct ReportBuilder<'a> {
     frames_post_processor: Option<Box<dyn Fn(&mut Frames)>>,
+    demangle: Box<dyn Fn(&str) -> Cow<str>>,
+
     profiler: &'a RwLock<Result<Profiler>>,
     timing: ReportTiming,
 }
@@ -40,6 +44,8 @@ impl<'a> ReportBuilder<'a> {
     pub(crate) fn new(profiler: &'a RwLock<Result<Profiler>>, timing: ReportTiming) -> Self {
         Self {
             frames_post_processor: None,
+            demangle: Box::new(demangle),
+
             profiler,
             timing,
         }
@@ -57,6 +63,29 @@ impl<'a> ReportBuilder<'a> {
         self
     }
 
+    /// Set `demangle` of a `ReportBuilder`. Before finally building a report,
+    /// `demangle` will be applied to every symbol.
+    /// # Examples
+    ///
+    /// ```
+    /// # use std::borrow::Cow;
+    /// fn demangle(symbol:&str) -> Cow<'_, str> {
+    ///     println!("demangling {}", symbol);
+    ///     Cow::from(symbol)
+    /// };
+    ///
+    /// let guard = pprof::ProfilerGuard::new(100).unwrap();
+    /// guard.report().demangle(demangle).build().unwrap();
+    /// ```
+    pub fn demangle<T>(&mut self, demangle: T) -> &mut Self
+    where
+        T: Fn(&str) -> Cow<str> + 'static,
+    {
+        self.demangle = Box::new(demangle);
+
+        self
+    }
+
     /// Build an `UnresolvedReport`
     pub fn build_unresolved(&self) -> Result<UnresolvedReport> {
         let mut hash_map = HashMap::new();
@@ -108,7 +137,8 @@ impl<'a> ReportBuilder<'a> {
                 profiler.data.try_iter()?.for_each(|entry| {
                     let count = entry.count;
                     if count > 0 {
-                        let mut key = Frames::from(entry.item.clone());
+                        let mut key =
+                            Frames::from_with_demangle(entry.item.clone(), &self.demangle);
                         if let Some(processor) = &self.frames_post_processor {
                             processor(&mut key);
                         }