stack2reg: Switch to hashbrown::HashSet

3 files changed, 22 insertions, 8 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 77d1fd0b405..b072d9bd8c2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,12 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
 [[package]]
+name = "ahash"
+version = "0.3.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8fd72866655d1904d6b0997d0b07ba561047d070fbe29de039031c641b61217"
+
+[[package]]
 name = "anyhow"
 version = "1.0.32"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -197,6 +203,7 @@ version = "0.8.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "34f595585f103464d8d2f6e9864682d74c1601fed5e07d62b1c9058dba8246fb"
 dependencies = [
+ "ahash",
  "autocfg",
 ]
 
@@ -325,6 +332,7 @@ dependencies = [
  "cranelift-object",
  "cranelift-simplejit",
  "gimli",
+ "hashbrown",
  "indexmap",
  "libloading",
  "object",
diff --git a/Cargo.toml b/Cargo.toml
index e178e29925b..858a0f0e33e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,6 +23,7 @@ byteorder = "1.2.7"
 indexmap = "1.0.2"
 cfg-if = "0.1.10"
 libloading = { version = "0.6.0", optional = true }
+hashbrown = "0.8.1"
 
 # Uncomment to use local checkout of cranelift
 #[patch."https://github.com/bytecodealliance/wasmtime/"]
diff --git a/src/optimize/stack2reg.rs b/src/optimize/stack2reg.rs
index d1826a9af72..aceced41f27 100644
--- a/src/optimize/stack2reg.rs
+++ b/src/optimize/stack2reg.rs
@@ -13,12 +13,15 @@ use std::collections::BTreeMap;
 use std::fmt;
 use std::ops::Not;
 
-use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::fx::{FxHashSet, FxHasher};
 
 use cranelift_codegen::cursor::{Cursor, FuncCursor};
 use cranelift_codegen::ir::{InstructionData, Opcode, ValueDef};
 use cranelift_codegen::ir::immediates::Offset32;
 
+use hashbrown::HashSet;
+use std::hash::BuildHasherDefault;
+
 use crate::prelude::*;
 
 /// Workaround for `StackSlot` not implementing `Ord`.
@@ -45,9 +48,9 @@ impl Ord for OrdStackSlot {
 
 #[derive(Debug, Default)]
 struct StackSlotUsage {
-    stack_addr: FxHashSet<Inst>,
-    stack_load: FxHashSet<Inst>,
-    stack_store: FxHashSet<Inst>,
+    stack_addr: HashSet<Inst, BuildHasherDefault<FxHasher>>,
+    stack_load: HashSet<Inst, BuildHasherDefault<FxHasher>>,
+    stack_store: HashSet<Inst, BuildHasherDefault<FxHasher>>,
 }
 
 impl StackSlotUsage {
@@ -313,17 +316,19 @@ fn remove_unused_stack_addr_and_stack_load(opt_ctx: &mut OptimizeContext<'_>) {
 
     // Replace all unused stack_addr and stack_load instructions with nop.
     let mut func = &mut opt_ctx.ctx.func;
+
+    // drain_filter() on hashbrown::HashSet drains the items that do *not* match the
+    // predicate. Once hashbrown gets updated to match the behaviour of std::drain_filter
+    // (0.8.2), the predicate will have to be reversed
     for stack_slot_users in opt_ctx.stack_slot_usage_map.values_mut() {
         stack_slot_users
             .stack_addr
-            .drain()
-            .filter(|inst| stack_addr_load_insts_users.get(inst).map(|users| users.is_empty()).unwrap_or(true))
+            .drain_filter(|inst| !(stack_addr_load_insts_users.get(inst).map(|users| users.is_empty()).unwrap_or(true)))
             .for_each(|inst| StackSlotUsage::remove_unused_stack_addr(&mut func, inst));
 
         stack_slot_users
             .stack_load
-            .drain()
-            .filter(|inst| stack_addr_load_insts_users.get(inst).map(|users| users.is_empty()).unwrap_or(true))
+            .drain_filter(|inst| !(stack_addr_load_insts_users.get(inst).map(|users| users.is_empty()).unwrap_or(true)))
             .for_each(|inst| StackSlotUsage::remove_unused_load(&mut func, inst));
     }
 }