From 1f0e0a8eb4429351412e52b1ba1054f5298cd3b7 Mon Sep 17 00:00:00 2001
From: Ricky Lopez <ricky.lopez@flyzipline.com>
Date: Mon, 29 Jul 2024 15:52:06 -0700
Subject: [PATCH] Add Legend Formatter support

---
 egui_plot/src/legend.rs | 10 ++++++++--
 egui_plot/src/lib.rs    | 43 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/egui_plot/src/legend.rs b/egui_plot/src/legend.rs
index 8e81752..0b44714 100644
--- a/egui_plot/src/legend.rs
+++ b/egui_plot/src/legend.rs
@@ -6,6 +6,7 @@ use egui::{
 };
 
 use super::items::PlotItem;
+use super::LegendFormatterFn;
 
 /// Where to place the plot legend.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -192,6 +193,7 @@ impl LegendWidget {
         config: Legend,
         items: &[Box<dyn PlotItem>],
         hidden_items: &ahash::HashSet<String>, // Existing hidden items in the plot memory.
+        formatter: Option<&LegendFormatterFn>,
     ) -> Option<Self> {
         // If `config.hidden_items` is not `None`, it is used.
         let hidden_items = config.hidden_items.as_ref().unwrap_or(hidden_items);
@@ -203,8 +205,12 @@ impl LegendWidget {
             .iter()
             .filter(|item| !item.name().is_empty())
             .for_each(|item| {
+                let name = formatter.map_or_else(
+                    || item.name().to_owned(),
+                    |formatter| formatter(item.name()).to_owned(),
+                );
                 entries
-                    .entry(item.name().to_owned())
+                    .entry(name.clone())
                     .and_modify(|entry| {
                         if entry.color != item.color() {
                             // Multiple items with different colors
@@ -213,7 +219,7 @@ impl LegendWidget {
                     })
                     .or_insert_with(|| {
                         let color = item.color();
-                        let checked = !hidden_items.contains(item.name());
+                        let checked = !hidden_items.contains(&name);
                         LegendEntry::new(color, checked)
                     });
             });
diff --git a/egui_plot/src/lib.rs b/egui_plot/src/lib.rs
index 50b12cf..00d9d13 100644
--- a/egui_plot/src/lib.rs
+++ b/egui_plot/src/lib.rs
@@ -45,6 +45,9 @@ use legend::LegendWidget;
 type LabelFormatterFn<'a> = dyn Fn(&str, &PlotPoint) -> String + 'a;
 pub type LabelFormatter<'a> = Option<Box<LabelFormatterFn<'a>>>;
 
+type LegendFormatterFn = dyn Fn(&str) -> &str;
+type LegendFormatter = Option<Box<LegendFormatterFn>>;
+
 type GridSpacerFn<'a> = dyn Fn(GridInput) -> Vec<GridMark> + 'a;
 type GridSpacer<'a> = Box<GridSpacerFn<'a>>;
 
@@ -176,6 +179,7 @@ pub struct Plot<'a> {
     show_x: bool,
     show_y: bool,
     label_formatter: LabelFormatter<'a>,
+    legend_formatter: LegendFormatter,
     coordinates_formatter: Option<(Corner, CoordinatesFormatter<'a>)>,
     x_axes: Vec<AxisHints<'a>>, // default x axes
     y_axes: Vec<AxisHints<'a>>, // default y axes
@@ -223,6 +227,7 @@ impl<'a> Plot<'a> {
             show_x: true,
             show_y: true,
             label_formatter: None,
+            legend_formatter: None,
             coordinates_formatter: None,
             x_axes: vec![AxisHints::new(Axis::X)],
             y_axes: vec![AxisHints::new(Axis::Y)],
@@ -416,6 +421,15 @@ impl<'a> Plot<'a> {
         self
     }
 
+    /// Provide a function to customize the legend labels
+    ///
+    /// All items with the same name will be shown with (and have their visibility controlled by) a single label
+    #[inline]
+    pub fn legend_formatter(mut self, legend_formatter: impl Fn(&str) -> &str + 'static) -> Self {
+        self.legend_formatter = Some(Box::new(legend_formatter));
+        self
+    }
+
     /// Show the pointer coordinates in the plot.
     pub fn coordinates_formatter(
         mut self,
@@ -764,6 +778,7 @@ impl<'a> Plot<'a> {
             mut show_x,
             mut show_y,
             label_formatter,
+            legend_formatter,
             coordinates_formatter,
             x_axes,
             y_axes,
@@ -893,20 +908,40 @@ impl<'a> Plot<'a> {
         }
 
         // --- Legend ---
-        let legend = legend_config
-            .and_then(|config| LegendWidget::try_new(plot_rect, config, &items, &mem.hidden_items));
+        let legend = legend_config.and_then(|config| {
+            LegendWidget::try_new(
+                plot_rect,
+                config,
+                &items,
+                &mem.hidden_items,
+                legend_formatter
+                    .as_ref()
+                    .map(|formatter| formatter.as_ref()),
+            )
+        });
         // Don't show hover cursor when hovering over legend.
         if mem.hovered_legend_item.is_some() {
             show_x = false;
             show_y = false;
         }
         // Remove the deselected items.
-        items.retain(|item| !mem.hidden_items.contains(item.name()));
+        items.retain(|item| {
+            !mem.hidden_items.contains(
+                legend_formatter
+                    .as_ref()
+                    .map_or_else(|| item.name(), |formatter| formatter(item.name())),
+            )
+        });
         // Highlight the hovered items.
         if let Some(hovered_name) = &mem.hovered_legend_item {
             items
                 .iter_mut()
-                .filter(|entry| entry.name() == hovered_name)
+                .filter(|entry| {
+                    legend_formatter
+                        .as_ref()
+                        .map_or_else(|| entry.name(), |formatter| formatter(entry.name()))
+                        == hovered_name
+                })
                 .for_each(|entry| entry.highlight());
         }
         // Move highlighted items to front.