From 0291e6b0042aef2d386f1b05846cd4b41eac955e Mon Sep 17 00:00:00 2001
From: Dylan Marriner <marriner@amazon.com>
Date: Tue, 25 May 2021 14:06:46 -0700
Subject: [PATCH] Color user-defined async functions correctly for Rust
 coloring

---
 src/flamegraph/color/palettes.rs | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/flamegraph/color/palettes.rs b/src/flamegraph/color/palettes.rs
index 0485513d..20f9f5ef 100644
--- a/src/flamegraph/color/palettes.rs
+++ b/src/flamegraph/color/palettes.rs
@@ -106,7 +106,11 @@ pub(super) mod rust {
         if name.starts_with("core::")
             || name.starts_with("std::")
             || name.starts_with("alloc::")
-            || name.starts_with("<core::")
+            || (name.starts_with("<core::") 
+                // Rust user-defined async functions are desugared into 
+                // GenFutures so we don't want to include those as Rust 
+                // system functions
+                && !name.starts_with("<core::future::from_generator::GenFuture<T>"))
             || name.starts_with("<std::")
             || name.starts_with("<alloc::")
         {
@@ -512,6 +516,12 @@ mod tests {
                 input: String::from("<core::something as something::else"),
                 output: BasicPalette::Orange,
             },
+            TestData {
+                input: String::from(
+                    "<core::future::from_generator::GenFuture<T> as something::else",
+                ),
+                output: BasicPalette::Aqua,
+            },
             TestData {
                 input: String::from("<std::something something::else"),
                 output: BasicPalette::Orange,