1 files changed, 110 insertions, 70 deletions

diff --git a/crates/ra_prof/src/lib.rs b/crates/ra_prof/src/lib.rs
index da541005a..d38ff397e 100644
--- a/crates/ra_prof/src/lib.rs
+++ b/crates/ra_prof/src/lib.rs
@@ -6,10 +6,9 @@ mod google_cpu_profiler;
 
 use std::{
     cell::RefCell,
+    collections::BTreeMap,
     collections::HashSet,
     io::{stderr, Write},
-    iter::repeat,
-    mem,
     sync::{
         atomic::{AtomicBool, Ordering},
         RwLock,
@@ -17,7 +16,6 @@ use std::{
     time::{Duration, Instant},
 };
 
-use itertools::Itertools;
 use once_cell::sync::Lazy;
 
 pub use crate::memory_usage::{Bytes, MemoryUsage};
@@ -51,6 +49,8 @@ pub fn set_filter(f: Filter) {
     *old = filter_data;
 }
 
+pub type Label = &'static str;
+
 /// This function starts a profiling scope in the current execution stack with a given description.
 /// It returns a Profile structure and measure elapsed time between this method invocation and Profile structure drop.
 /// It supports nested profiling scopes in case when this function invoked multiple times at the execution stack. In this case the profiling information will be nested at the output.
@@ -78,10 +78,10 @@ pub fn set_filter(f: Filter) {
 ///  0ms - profile
 ///      0ms - profile2
 /// ```
-pub fn profile(desc: &str) -> Profiler {
-    assert!(!desc.is_empty());
+pub fn profile(label: Label) -> Profiler {
+    assert!(!label.is_empty());
     if !PROFILING_ENABLED.load(Ordering::Relaxed) {
-        return Profiler { desc: None };
+        return Profiler { label: None };
     }
 
     PROFILE_STACK.with(|stack| {
@@ -94,20 +94,35 @@ pub fn profile(desc: &str) -> Profiler {
             };
         }
         if stack.starts.len() > stack.filter_data.depth {
-            return Profiler { desc: None };
+            return Profiler { label: None };
         }
         let allowed = &stack.filter_data.allowed;
-        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(desc) {
-            return Profiler { desc: None };
+        if stack.starts.is_empty() && !allowed.is_empty() && !allowed.contains(label) {
+            return Profiler { label: None };
         }
 
         stack.starts.push(Instant::now());
-        Profiler { desc: Some(desc.to_string()) }
+        Profiler { label: Some(label) }
     })
 }
 
+pub fn print_time(label: Label) -> impl Drop {
+    struct Guard {
+        label: Label,
+        start: Instant,
+    }
+
+    impl Drop for Guard {
+        fn drop(&mut self) {
+            eprintln!("{}: {:?}", self.label, self.start.elapsed())
+        }
+    }
+
+    Guard { label, start: Instant::now() }
+}
+
 pub struct Profiler {
-    desc: Option<String>,
+    label: Option<Label>,
 }
 
 pub struct Filter {
@@ -160,7 +175,7 @@ struct ProfileStack {
 struct Message {
     level: usize,
     duration: Duration,
-    message: String,
+    label: Label,
 }
 
 impl ProfileStack {
@@ -186,14 +201,13 @@ thread_local!(static PROFILE_STACK: RefCell<ProfileStack> = RefCell::new(Profile
 impl Drop for Profiler {
     fn drop(&mut self) {
         match self {
-            Profiler { desc: Some(desc) } => {
+            Profiler { label: Some(label) } => {
                 PROFILE_STACK.with(|stack| {
                     let mut stack = stack.borrow_mut();
                     let start = stack.starts.pop().unwrap();
                     let duration = start.elapsed();
                     let level = stack.starts.len();
-                    let message = mem::replace(desc, String::new());
-                    stack.messages.push(Message { level, duration, message });
+                    stack.messages.push(Message { level, duration, label: label });
                     if level == 0 {
                         let stdout = stderr();
                         let longer_than = stack.filter_data.longer_than;
@@ -201,70 +215,96 @@ impl Drop for Profiler {
                         // (otherwise we could print `0ms` despite user's `>0` filter when
                         // `duration` is just a few nanos).
                         if duration.as_millis() > longer_than.as_millis() {
-                            print(0, &stack.messages, &mut stdout.lock(), longer_than, None);
+                            print(&stack.messages, longer_than, &mut stdout.lock());
                         }
                         stack.messages.clear();
                     }
                 });
             }
-            Profiler { desc: None } => (),
+            Profiler { label: None } => (),
         }
     }
 }
 
-fn print(
-    lvl: usize,
-    msgs: &[Message],
-    out: &mut impl Write,
-    longer_than: Duration,
-    total: Option<Duration>,
-) {
+fn print(msgs: &[Message], longer_than: Duration, out: &mut impl Write) {
     if msgs.is_empty() {
         return;
     }
-    // The index of the first element that will be included in the slice when we recurse.
-    let mut next_start = 0;
-    let indent = repeat("    ").take(lvl).collect::<String>();
-    // We output hierarchy for long calls, but sum up all short calls
-    let mut short = Vec::new();
+    let children_map = idx_to_children(msgs);
+    let root_idx = msgs.len() - 1;
+    print_for_idx(root_idx, &children_map, msgs, longer_than, out);
+}
+
+fn print_for_idx(
+    current_idx: usize,
+    children_map: &[Vec<usize>],
+    msgs: &[Message],
+    longer_than: Duration,
+    out: &mut impl Write,
+) {
+    let current = &msgs[current_idx];
+    let current_indent = "    ".repeat(current.level);
+    writeln!(out, "{}{:5}ms - {}", current_indent, current.duration.as_millis(), current.label)
+        .expect("printing profiling info");
+
+    let longer_than_millis = longer_than.as_millis();
+    let children_indices = &children_map[current_idx];
     let mut accounted_for = Duration::default();
-    for (i, &Message { level, duration, message: ref msg }) in msgs.iter().enumerate() {
-        if level != lvl {
-            continue;
-        }
-        accounted_for += duration;
-        if duration.as_millis() > longer_than.as_millis() {
-            writeln!(out, "{}{:5}ms - {}", indent, duration.as_millis(), msg)
-                .expect("printing profiling info to stdout");
+    let mut short_children = BTreeMap::new(); // Use `BTreeMap` to get deterministic output.
 
-            print(lvl + 1, &msgs[next_start..i], out, longer_than, Some(duration));
+    for child_idx in children_indices.iter() {
+        let child = &msgs[*child_idx];
+        if child.duration.as_millis() > longer_than_millis {
+            print_for_idx(*child_idx, children_map, msgs, longer_than, out);
         } else {
-            short.push((msg, duration))
+            let pair = short_children.entry(child.label).or_insert((Duration::default(), 0));
+            pair.0 += child.duration;
+            pair.1 += 1;
         }
+        accounted_for += child.duration;
+    }
 
-        next_start = i + 1;
+    for (child_msg, (duration, count)) in short_children.iter() {
+        let millis = duration.as_millis();
+        writeln!(out, "    {}{:5}ms - {} ({} calls)", current_indent, millis, child_msg, count)
+            .expect("printing profiling info");
     }
-    short.sort_by_key(|(msg, _time)| *msg);
-    for (msg, entires) in short.iter().group_by(|(msg, _time)| msg).into_iter() {
-        let mut count = 0;
-        let mut total_duration = Duration::default();
-        entires.for_each(|(_msg, time)| {
-            count += 1;
-            total_duration += *time;
-        });
-        writeln!(out, "{}{:5}ms - {} ({} calls)", indent, total_duration.as_millis(), msg, count)
-            .expect("printing profiling info to stdout");
+
+    let unaccounted_millis = (current.duration - accounted_for).as_millis();
+    if !children_indices.is_empty()
+        && unaccounted_millis > 0
+        && unaccounted_millis > longer_than_millis
+    {
+        writeln!(out, "    {}{:5}ms - ???", current_indent, unaccounted_millis)
+            .expect("printing profiling info");
     }
+}
 
-    if let Some(total) = total {
-        if let Some(unaccounted) = total.checked_sub(accounted_for) {
-            let unaccounted_millis = unaccounted.as_millis();
-            if unaccounted_millis > longer_than.as_millis() && unaccounted_millis > 0 {
-                writeln!(out, "{}{:5}ms - ???", indent, unaccounted_millis)
-                    .expect("printing profiling info to stdout");
-            }
+/// Returns a mapping from an index in the `msgs` to the vector with the indices of its children.
+///
+/// This assumes that the entries in `msgs` are in the order of when the calls to `profile` finish.
+/// In other words, a postorder of the call graph. In particular, the root is the last element of
+/// `msgs`.
+fn idx_to_children(msgs: &[Message]) -> Vec<Vec<usize>> {
+    // Initialize with the index of the root; `msgs` and `ancestors` should be never empty.
+    assert!(!msgs.is_empty());
+    let mut ancestors = vec![msgs.len() - 1];
+    let mut result: Vec<Vec<usize>> = vec![vec![]; msgs.len()];
+    for (idx, msg) in msgs[..msgs.len() - 1].iter().enumerate().rev() {
+        // We need to find the parent of the current message, i.e., the last ancestor that has a
+        // level lower than the current message.
+        while msgs[*ancestors.last().unwrap()].level >= msg.level {
+            ancestors.pop();
         }
+        result[*ancestors.last().unwrap()].push(idx);
+        ancestors.push(idx);
+    }
+    // Note that above we visited all children from the last to the first one. Let's reverse vectors
+    // to get the more natural order where the first element is the first child.
+    for vec in result.iter_mut() {
+        vec.reverse();
     }
+    result
 }
 
 /// Prints backtrace to stderr, useful for debugging.
@@ -369,11 +409,11 @@ mod tests {
     fn test_longer_than() {
         let mut result = vec![];
         let msgs = vec![
-            Message { level: 1, duration: Duration::from_nanos(3), message: "bar".to_owned() },
-            Message { level: 1, duration: Duration::from_nanos(2), message: "bar".to_owned() },
-            Message { level: 0, duration: Duration::from_millis(1), message: "foo".to_owned() },
+            Message { level: 1, duration: Duration::from_nanos(3), label: "bar" },
+            Message { level: 1, duration: Duration::from_nanos(2), label: "bar" },
+            Message { level: 0, duration: Duration::from_millis(1), label: "foo" },
         ];
-        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
+        print(&msgs, Duration::from_millis(0), &mut result);
         // The calls to `bar` are so short that they'll be rounded to 0ms and should get collapsed
         // when printing.
         assert_eq!(
@@ -386,10 +426,10 @@ mod tests {
     fn test_unaccounted_for_topmost() {
         let mut result = vec![];
         let msgs = vec![
-            Message { level: 1, duration: Duration::from_millis(2), message: "bar".to_owned() },
-            Message { level: 0, duration: Duration::from_millis(5), message: "foo".to_owned() },
+            Message { level: 1, duration: Duration::from_millis(2), label: "bar" },
+            Message { level: 0, duration: Duration::from_millis(5), label: "foo" },
         ];
-        print(0, &msgs, &mut result, Duration::from_millis(0), Some(Duration::from_millis(1)));
+        print(&msgs, Duration::from_millis(0), &mut result);
         assert_eq!(
             std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
             vec![
@@ -405,13 +445,13 @@ mod tests {
     fn test_unaccounted_for_multiple_levels() {
         let mut result = vec![];
         let msgs = vec![
-            Message { level: 2, duration: Duration::from_millis(3), message: "baz".to_owned() },
-            Message { level: 1, duration: Duration::from_millis(5), message: "bar".to_owned() },
-            Message { level: 2, duration: Duration::from_millis(2), message: "baz".to_owned() },
-            Message { level: 1, duration: Duration::from_millis(4), message: "bar".to_owned() },
-            Message { level: 0, duration: Duration::from_millis(9), message: "foo".to_owned() },
+            Message { level: 2, duration: Duration::from_millis(3), label: "baz" },
+            Message { level: 1, duration: Duration::from_millis(5), label: "bar" },
+            Message { level: 2, duration: Duration::from_millis(2), label: "baz" },
+            Message { level: 1, duration: Duration::from_millis(4), label: "bar" },
+            Message { level: 0, duration: Duration::from_millis(9), label: "foo" },
         ];
-        print(0, &msgs, &mut result, Duration::from_millis(0), None);
+        print(&msgs, Duration::from_millis(0), &mut result);
         assert_eq!(
             std::str::from_utf8(&result).unwrap().lines().collect::<Vec<_>>(),
             vec![