From e546275a5dc069489bbe3689090641d1ca1394a4 Mon Sep 17 00:00:00 2001
From: YangKeao <yangkeao@chunibyo.icu>
Date: Mon, 9 May 2022 17:03:27 +0800
Subject: [PATCH] add custom demangler

Signed-off-by: YangKeao <yangkeao@chunibyo.icu>
---
 Cargo.toml    |  2 +-
 src/frames.rs |  7 +++++++
 src/report.rs | 33 +++++++++++++++++++++++++++++++--
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 54f52bc1..c5e4cf3e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,7 +28,7 @@ backtrace = { version = "0.3", optional = true }
 once_cell = "1.9"
 libc = "^0.2.66"
 log = "0.4"
-nix = { version = "0.24", default-features = false, features = ["signal"] }
+nix = { version = "0.24", default-features = false, features = ["signal", "fs"] }
 parking_lot = "0.12"
 tempfile = "3.1"
 thiserror = "1.0"
diff --git a/src/frames.rs b/src/frames.rs
index 83ed61a9..e18c337e 100644
--- a/src/frames.rs
+++ b/src/frames.rs
@@ -104,6 +104,13 @@ impl Symbol {
     }
 
     pub fn name(&self) -> String {
+        self.name_with_demangle(demangle)
+    }
+
+    pub fn name_with_demangle<T>(&self, demangle: T) -> String
+    where
+        T: Fn(&str) -> Cow<str> + 'static,
+    {
         demangle(&String::from_utf8_lossy(self.raw_name())).into_owned()
     }
 
diff --git a/src/report.rs b/src/report.rs
index a19b940e..d41e664f 100644
--- a/src/report.rs
+++ b/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();
@@ -232,7 +261,7 @@ mod protobuf {
                 dedup_str.insert(key.thread_name_or_id());
                 for frame in key.frames.iter() {
                     for symbol in frame {
-                        dedup_str.insert(symbol.name());
+                        dedup_str.insert(symbol.name_with_demangle(&self.demangle));
                         dedup_str.insert(symbol.sys_name().into_owned());
                         dedup_str.insert(symbol.filename().into_owned());
                     }
@@ -260,7 +289,7 @@ mod protobuf {
                 let mut locs = vec![];
                 for frame in key.frames.iter() {
                     for symbol in frame {
-                        let name = symbol.name();
+                        let name = symbol.name_with_demangle(&self.demangle);
                         if let Some(loc_idx) = functions.get(&name) {
                             locs.push(*loc_idx);
                             continue;