From 0823f0c32b600dafcc707e76dc1898eadc3c2b59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?John=20K=C3=A5re=20Alsaker?= <john.kare.alsaker@gmail.com>
Date: Wed, 16 Aug 2023 10:44:32 +0200
Subject: [PATCH] Remove `count`

---
 compiler/rustc_data_structures/src/sharded.rs | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs
index 39cce5c8ee6fd..52ab5a7fb144e 100644
--- a/compiler/rustc_data_structures/src/sharded.rs
+++ b/compiler/rustc_data_structures/src/sharded.rs
@@ -49,7 +49,7 @@ impl<T> Sharded<T> {
         match self {
             Self::Single(single) => &single,
             #[cfg(parallel_compiler)]
-            Self::Shards(shards) => self.get_shard_by_hash(make_hash(_val)),
+            Self::Shards(..) => self.get_shard_by_hash(make_hash(_val)),
         }
     }
 
@@ -70,21 +70,20 @@ impl<T> Sharded<T> {
         }
     }
 
-    #[inline]
-    fn count(&self) -> usize {
+    pub fn lock_shards(&self) -> Vec<LockGuard<'_, T>> {
         match self {
-            Self::Single(..) => 1,
+            Self::Single(single) => vec![single.lock()],
             #[cfg(parallel_compiler)]
-            Self::Shards(..) => SHARDS,
+            Self::Shards(shards) => shards.iter().map(|shard| shard.0.lock()).collect(),
         }
     }
 
-    pub fn lock_shards(&self) -> Vec<LockGuard<'_, T>> {
-        (0..self.count()).map(|i| self.get_shard_by_index(i).lock()).collect()
-    }
-
     pub fn try_lock_shards(&self) -> Option<Vec<LockGuard<'_, T>>> {
-        (0..self.count()).map(|i| self.get_shard_by_index(i).try_lock()).collect()
+        match self {
+            Self::Single(single) => Some(vec![single.try_lock()?]),
+            #[cfg(parallel_compiler)]
+            Self::Shards(shards) => shards.iter().map(|shard| shard.0.try_lock()).collect(),
+        }
     }
 }