From 59d46c2fecf891b5479b91b96ba315042831dc8a Mon Sep 17 00:00:00 2001
From: uint256_t <konndennsa@gmail.com>
Date: Wed, 22 Mar 2023 03:59:54 +0900
Subject: [PATCH] cranelift-entity: improve `EntitySet::cardinality()`
 implementation (#6066)

* Simplify 'EntitySet::cardinality()'

* Fix test
---
 cranelift/entity/src/set.rs | 55 +++++++++++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/cranelift/entity/src/set.rs b/cranelift/entity/src/set.rs
index a5c8b7b295..60578d6e4b 100644
--- a/cranelift/entity/src/set.rs
+++ b/cranelift/entity/src/set.rs
@@ -69,20 +69,14 @@ where
         }
     }
 
-    /// Returns the cardinality of the set.  More precisely, it returns the number of calls to
+    /// Returns the cardinality of the set. More precisely, it returns the number of calls to
     /// `insert` with different key values, that have happened since the the set was most recently
     /// `clear`ed or created with `new`.
     pub fn cardinality(&self) -> usize {
-        let mut n: usize = 0;
-        for idx in 0..self.len / BITS {
-            n += self.elems[idx].count_ones() as usize;
-        }
-        for bit_ix in (self.len / BITS) * BITS..self.len {
-            if (self.elems[bit_ix / BITS] & (1 << (bit_ix % BITS))) != 0 {
-                n += 1;
-            }
-        }
-        n
+        self.elems[..(self.len + (BITS - 1)) / BITS]
+            .iter()
+            .map(|x| x.count_ones() as usize)
+            .sum()
     }
 
     /// Remove all entries from this set.
@@ -254,4 +248,43 @@ mod tests {
 
         assert!(m.is_empty());
     }
+
+    #[test]
+    fn cardinality() {
+        let mut m = EntitySet::new();
+
+        m.insert(E(1));
+        assert!(m.cardinality() == 1);
+
+        m.insert(E(0));
+        assert!(m.cardinality() == 2);
+
+        m.insert(E(1));
+        assert!(m.cardinality() == 2);
+
+        m.insert(E(BITS as u32 - 1));
+        assert!(m.cardinality() == 3);
+
+        m.insert(E(BITS as u32));
+        assert!(m.cardinality() == 4);
+
+        m.insert(E(BITS as u32 - 1));
+        assert!(m.cardinality() == 4);
+
+        assert!(m.pop() == Some(E(BITS as u32)));
+        assert!(m.cardinality() == 3);
+        assert!(m.pop() == Some(E(BITS as u32 - 1)));
+        assert!(m.cardinality() == 2);
+
+        m.insert(E(100));
+        assert!(m.cardinality() == 3);
+
+        assert!(m.pop() == Some(E(100)));
+        assert!(m.cardinality() == 2);
+
+        m.insert(E(100));
+        m.insert(E(101));
+        m.insert(E(102));
+        assert!(m.cardinality() == 5);
+    }
 }