Add Rust language coloring

src/bin/flamegraph.rs

@@ -82,7 +82,7 @@ struct Opt {
         short = "c",
         long = "colors",
         default_value = defaults::COLORS,
-        possible_values = &["aqua","blue","green","hot","io","java","js","mem","orange","perl","purple","red","wakeup","yellow"],
+        possible_values = &["aqua","blue","green","hot","io","java","js","mem","orange","perl","purple","red","rust","wakeup","yellow"],
         value_name = "STRING"
     )]
     colors: Palette,

src/flamegraph/color/mod.rs

@@ -124,6 +124,8 @@ pub enum MultiPalette {
     Js,
     /// Use Perl semantics to color frames.
     Perl,
+    /// Use Rust semantics to color frames.
+    Rust,
     /// Equivalent to [`BasicPalette::Aqua`] with [`BackgroundColor::Blue`].
     Wakeup,
 }
@@ -196,6 +198,7 @@ impl FromStr for Palette {
             "java" => Ok(Palette::Multi(MultiPalette::Java)),
             "js" => Ok(Palette::Multi(MultiPalette::Js)),
             "perl" => Ok(Palette::Multi(MultiPalette::Perl)),
+            "rust" => Ok(Palette::Multi(MultiPalette::Rust)),
             "red" => Ok(Palette::Basic(BasicPalette::Red)),
             "green" => Ok(Palette::Basic(BasicPalette::Green)),
             "blue" => Ok(Palette::Basic(BasicPalette::Blue)),
@@ -308,6 +311,7 @@ fn rgb_components_for_palette(palette: Palette, name: &str, v1: f32, v2: f32, v3
         Palette::Multi(MultiPalette::Perl) => palettes::perl::resolve(name),
         Palette::Multi(MultiPalette::Js) => palettes::js::resolve(name),
         Palette::Multi(MultiPalette::Wakeup) => palettes::wakeup::resolve(name),
+        Palette::Multi(MultiPalette::Rust) => palettes::rust::resolve(name),
     };
 
     match basic_palette {

src/flamegraph/color/palettes.rs

@@ -99,6 +99,32 @@ pub(super) mod wakeup {
     }
 }
 
+pub(super) mod rust {
+    use crate::flamegraph::color::BasicPalette;
+
+    pub fn resolve(name: &str) -> BasicPalette {
+        if name.starts_with("core::")
+            || name.starts_with("std::")
+            || name.starts_with("alloc::")
+            || name.starts_with("<core::")
+            || name.starts_with("<std::")
+            || name.starts_with("<alloc::")
+        {
+            // Rust system functions
+            BasicPalette::Orange
+        } else if name.contains("::") {
+            // Rust user functions.
+            // Although this will generate false positives for e.g. C++ code
+            // used with Rust, the intention is to color code from user
+            // crates and dependencies differenly than Rust system code.
+            BasicPalette::Aqua
+        } else {
+            // Non-Rust functions
+            BasicPalette::Yellow
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::flamegraph::color::BasicPalette;
@@ -452,4 +478,48 @@ mod tests {
             assert_eq!(result, elem.output);
         }
     }
+
+    #[test]
+    fn rust_returns_correct() {
+        use super::rust;
+
+        let test_names = [
+            TestData {
+                input: String::from("some::not_rust_system::mod"),
+                output: BasicPalette::Aqua,
+            },
+            TestData {
+                input: String::from("core::mod"),
+                output: BasicPalette::Orange,
+            },
+            TestData {
+                input: String::from("std::mod"),
+                output: BasicPalette::Orange,
+            },
+            TestData {
+                input: String::from("alloc::mod"),
+                output: BasicPalette::Orange,
+            },
+            TestData {
+                input: String::from("something_else"),
+                output: BasicPalette::Yellow,
+            },
+            TestData {
+                input: String::from("<alloc::boxed::Box<F,A> as something::else"),
+                output: BasicPalette::Orange,
+            },
+            TestData {
+                input: String::from("<core::something as something::else"),
+                output: BasicPalette::Orange,
+            },
+            TestData {
+                input: String::from("<std::something something::else"),
+                output: BasicPalette::Orange,
+            },
+        ];
+        for elem in test_names.iter() {
+            let result = rust::resolve(&elem.input);
+            assert_eq!(result, elem.output);
+        }
+    }
 }