fix: consistent colors for packages

crates/turborepo-ui/src/color_selector.rs

@@ -1,21 +1,20 @@
 use std::{
     collections::HashMap,
+    hash::{DefaultHasher, Hash, Hasher},
     sync::{Arc, OnceLock, RwLock},
+    u8,
 };
 
 use console::{Style, StyledObject};
 
-static COLORS: OnceLock<[Style; 5]> = OnceLock::new();
+static COLORS: OnceLock<[Style; u8::MAX as usize]> = OnceLock::new();
 
-pub fn get_terminal_package_colors() -> &'static [Style; 5] {
+pub fn get_terminal_package_colors() -> &'static [Style; u8::MAX as usize] {
     COLORS.get_or_init(|| {
-        [
-            Style::new().cyan(),
-            Style::new().magenta(),
-            Style::new().green(),
-            Style::new().yellow(),
-            Style::new().blue(),
-        ]
+        let colors: [Style; u8::MAX as usize] =
+            core::array::from_fn(|index| Style::new().color256(index as u8));
+
+        colors
     })
 }
 
@@ -26,10 +25,20 @@ pub struct ColorSelector {
     inner: Arc<RwLock<ColorSelectorInner>>,
 }
 
-#[derive(Default)]
 struct ColorSelectorInner {
-    idx: usize,
     cache: HashMap<String, &'static Style>,
+    colors_taken_state: [bool; u8::MAX as usize],
+    total_colors_taken: u8,
+}
+
+impl Default for ColorSelectorInner {
+    fn default() -> Self {
+        Self {
+            cache: Default::default(),
+            colors_taken_state: [false; u8::MAX as usize],
+            total_colors_taken: 0,
+        }
+    }
 }
 
 impl ColorSelector {
@@ -65,13 +74,40 @@ impl ColorSelectorInner {
 
     fn insert_color(&mut self, key: String) -> &'static Style {
         let colors = get_terminal_package_colors();
-        let chosen_color = &colors[self.idx % colors.len()];
+        let chosen_color = Self::get_color_id_by_key(self, &key);
         // A color might have been chosen by the time we get to inserting
-        self.cache.entry(key).or_insert_with(|| {
-            // If a color hasn't been chosen, then we increment the index
-            self.idx += 1;
-            chosen_color
-        })
+        self.cache
+            .entry(key)
+            .or_insert_with(|| &colors[chosen_color])
+    }
+
+    pub fn get_color_hash_by_key(key: &str) -> u64 {
+        let mut hasher = DefaultHasher::new();
+        key.hash(&mut hasher);
+        hasher.finish()
+    }
+
+    fn get_color_id_by_key(&mut self, key: &str) -> usize {
+        let colors = get_terminal_package_colors();
+
+        if self.total_colors_taken == u8::MAX {
+            self.colors_taken_state = [false; u8::MAX as usize];
+            self.total_colors_taken = 0;
+        }
+
+        let mut color_id: usize =
+            (Self::get_color_hash_by_key(&key) % colors.len() as u64) as usize;
+
+        let mut state: bool = *(self.colors_taken_state.get(color_id).unwrap());
+
+        while state {
+            color_id = (color_id + 1) % colors.len();
+            state = *(self.colors_taken_state.get(color_id).unwrap());
+        }
+
+        self.total_colors_taken += 1;
+        self.colors_taken_state[color_id] = true;
+        return color_id;
     }
 }
 
@@ -110,15 +146,67 @@ mod tests {
             });
         });
         // We only inserted 2 keys so next index should be 2
-        assert_eq!(selector.inner.read().unwrap().idx, 2);
+        assert_eq!(selector.inner.read().unwrap().total_colors_taken, 2);
     }
 
     #[test]
-    fn test_color_selector_wraps_around() {
+    fn test_rotation_after_all_colors_are_taken() {
         let selector = super::ColorSelector::default();
-        for key in &["1", "2", "3", "4", "5", "6"] {
-            selector.color_for_key(key);
+
+        let colors = super::get_terminal_package_colors();
+        let num_colors = colors.len();
+
+        // Exhaust all colors
+        for i in 0..num_colors {
+            let key = format!("package{}", i);
+            selector.color_for_key(&key);
         }
-        assert_eq!(selector.color_for_key("1"), selector.color_for_key("6"));
+
+        // At this point, all colors should be taken
+        for state in selector
+            .inner
+            .read()
+            .expect("lock poisoned")
+            .colors_taken_state
+            .iter()
+            .take(num_colors)
+        {
+            assert_eq!(*state, true);
+        }
+
+        // The next key should start rotating from the beginning
+        let key_next = format!("package{}", num_colors + 1);
+        let next_color_id = selector.color_for_key(&key_next);
+
+        // It should be the first color in the rotation again
+        let next_key_color_id = (super::ColorSelectorInner::get_color_hash_by_key(&key_next)
+            % colors.len() as u64) as usize;
+        assert_eq!(next_color_id, &colors[next_key_color_id]);
+
+        // At this point, all colors should be not taken, expect the one taken with the
+        // latest package
+        for (index, state) in selector
+            .inner
+            .read()
+            .expect("lock poisoned")
+            .colors_taken_state
+            .iter()
+            .enumerate()
+        {
+            if index == next_key_color_id {
+                assert_eq!(*state, true);
+            } else {
+                assert_eq!(*state, false);
+            }
+        }
+
+        assert_eq!(
+            selector
+                .inner
+                .read()
+                .expect("lock poisoned")
+                .total_colors_taken,
+            1
+        );
     }
 }