1 files changed, 24 insertions, 8 deletions

diff --git a/zjit/src/distribution.rs b/zjit/src/distribution.rs
index 5927ffa5c9..aa4667b939 100644
--- a/zjit/src/distribution.rs
+++ b/zjit/src/distribution.rs
@@ -1,14 +1,18 @@
+//! Type frequency distribution tracker.
+
+use crate::options::NumProfiles;
+
 /// This implementation was inspired by the type feedback module from Google's S6, which was
 /// written in C++ for use with Python. This is a new implementation in Rust created for use with
 /// Ruby instead of Python.
 #[derive(Debug, Clone)]
 pub struct Distribution<T: Copy + PartialEq + Default, const N: usize> {
     /// buckets and counts have the same length
-    /// buckets[0] is always the most common item
+    /// `buckets[0]` is always the most common item
     buckets: [T; N],
-    counts: [usize; N],
+    counts: [NumProfiles; N],
     /// if there is no more room, increment the fallback
-    other: usize,
+    other: NumProfiles,
     // TODO(max): Add count disparity, which can help determine when to reset the distribution
 }
 
@@ -21,13 +25,13 @@ impl<T: Copy + PartialEq + Default, const N: usize> Distribution<T, N> {
         for (bucket, count) in self.buckets.iter_mut().zip(self.counts.iter_mut()) {
             if *bucket == item || *count == 0 {
                 *bucket = item;
-                *count += 1;
+                *count = count.saturating_add(1);
                 // Keep the most frequent item at the front
                 self.bubble_up();
                 return;
             }
         }
-        self.other += 1;
+        self.other = self.other.saturating_add(1);
     }
 
     /// Keep the highest counted bucket at index 0
@@ -67,7 +71,7 @@ enum DistributionKind {
     SkewedMegamorphic,
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct DistributionSummary<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> {
     kind: DistributionKind,
     buckets: [T; N],
@@ -85,7 +89,7 @@ impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> Distributi
                 assert!(first_count >= count, "First count should be the largest");
             }
         }
-        let num_seen = dist.counts.iter().sum::<usize>() + dist.other;
+        let num_seen = dist.counts.iter().map(|&c| usize::from(c)).sum::<usize>() + usize::from(dist.other);
         let kind = if dist.other == 0 {
             // Seen <= N types total
             if dist.counts[0] == 0 {
@@ -105,17 +109,25 @@ impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> Distributi
                 DistributionKind::Megamorphic
             }
         };
-        Self { kind, buckets: dist.buckets.clone() }
+        Self { kind, buckets: dist.buckets }
     }
 
     pub fn is_monomorphic(&self) -> bool {
         self.kind == DistributionKind::Monomorphic
     }
 
+    pub fn is_polymorphic(&self) -> bool {
+        self.kind == DistributionKind::Polymorphic
+    }
+
     pub fn is_skewed_polymorphic(&self) -> bool {
         self.kind == DistributionKind::SkewedPolymorphic
     }
 
+    pub fn is_megamorphic(&self) -> bool {
+        self.kind == DistributionKind::Megamorphic
+    }
+
     pub fn is_skewed_megamorphic(&self) -> bool {
         self.kind == DistributionKind::SkewedMegamorphic
     }
@@ -124,6 +136,10 @@ impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> Distributi
         assert!(idx < N, "index {idx} out of bounds for buckets[{N}]");
         self.buckets[idx]
     }
+
+    pub fn buckets(&self) -> &[T] {
+        &self.buckets
+    }
 }
 
 #[cfg(test)]