From ea10be67cb1c13118a7aaf18619a79f54c0e2b5e Mon Sep 17 00:00:00 2001
From: Jeremy Leibs <jeremy@rerun.io>
Date: Fri, 7 Jun 2024 22:04:01 -0400
Subject: [PATCH] Fix time-histogram removal book-keeping

---
 crates/re_entity_db/src/time_histogram_per_timeline.rs |  9 +++------
 crates/re_entity_db/tests/time_histograms.rs           | 10 +++++++++-
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/crates/re_entity_db/src/time_histogram_per_timeline.rs b/crates/re_entity_db/src/time_histogram_per_timeline.rs
index d7c38d07e0b2..c2452e2bc406 100644
--- a/crates/re_entity_db/src/time_histogram_per_timeline.rs
+++ b/crates/re_entity_db/src/time_histogram_per_timeline.rs
@@ -100,12 +100,9 @@ impl TimeHistogramPerTimeline {
                 });
         } else {
             for (timeline, time_value) in timepoint.iter() {
-                let remaining_count = self
-                    .times
-                    .entry(*timeline)
-                    .or_default()
-                    .decrement(time_value.as_i64(), n);
-                if remaining_count == 0 {
+                let hist = self.times.entry(*timeline).or_default();
+                hist.decrement(time_value.as_i64(), n);
+                if hist.is_empty() {
                     self.times.remove(timeline);
                 }
             }
diff --git a/crates/re_entity_db/tests/time_histograms.rs b/crates/re_entity_db/tests/time_histograms.rs
index 0aacceca630d..b1282f33ba68 100644
--- a/crates/re_entity_db/tests/time_histograms.rs
+++ b/crates/re_entity_db/tests/time_histograms.rs
@@ -668,6 +668,10 @@ fn assert_recursive_histogram<'a>(
         let histo = db.time_histogram(timeline);
 
         if let Some(expected) = expected_times {
+            if expected.is_empty() {
+                assert!(histo.is_none());
+                continue;
+            }
             let histo = histo.unwrap();
             let ranges = histo.range(i64::MIN.., 0).collect::<Vec<_>>();
             let expected: Vec<_> = expected.to_vec();
@@ -692,9 +696,13 @@ fn assert_histogram_for_component<'a>(
             .and_then(|tree| tree.time_histogram_for_component(timeline, component_name));
 
         if let Some(expected) = expected_times {
+            let expected: Vec<_> = expected.to_vec();
+            if expected.is_empty() {
+                assert!(histo.is_none());
+                continue;
+            }
             let histo = histo.unwrap();
             let ranges = histo.range(i64::MIN.., 0).collect::<Vec<_>>();
-            let expected: Vec<_> = expected.to_vec();
             similar_asserts::assert_eq!(expected, ranges);
         } else {
             assert!(histo.is_none());