Replace bimap dependency with a more efficient pair of maps.

Cargo.lock

@@ -110,12 +110,6 @@ version = "0.13.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd"
 
-[[package]]
-name = "bimap"
-version = "0.6.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "50ae17cabbc8a38a1e3e4c1a6a664e9a09672dc14d0896fa8d865d3a5a446b07"
-
 [[package]]
 name = "bindgen"
 version = "0.56.0"
@@ -1998,7 +1992,6 @@ name = "rustc_codegen_spirv"
 version = "0.4.0-alpha.17"
 dependencies = [
  "ar",
- "bimap",
  "hashbrown",
  "indexmap",
  "libc",

crates/rustc_codegen_spirv/Cargo.toml

@@ -47,7 +47,6 @@ regex = { version = "1", features = ["perf"] }
 
 # Normal dependencies.
 ar = "0.9.0"
-bimap = "0.6"
 indexmap = "1.6.0"
 rspirv = "0.11"
 rustc-demangle = "0.1.21"

crates/rustc_codegen_spirv/src/builder_spirv.rs

@@ -232,6 +232,7 @@ pub enum SpirvConst {
     // different functions, but of the same type, don't overlap their zombies.
     ZombieUndefForFnAddr,
 
+    // FIXME(eddyb) use `tcx.arena.dropless` to get `&'tcx [_]`, instead of `Rc`.
     Composite(Rc<[Word]>),
 
     /// Pointer to constant data, i.e. `&pointee`, represented as an `OpVariable`

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -144,9 +144,11 @@ impl<'tcx> CodegenCx<'tcx> {
         self.builder.builder(cursor)
     }
 
+    // FIXME(eddyb) this should not clone the `SpirvType` back out, but we'd have
+    // to fix ~100 callsites (to add a `*` deref), if we change this signature.
     #[track_caller]
     pub fn lookup_type(&self, ty: Word) -> SpirvType {
-        self.type_cache.lookup(ty)
+        (*self.type_cache.lookup(ty)).clone()
     }
 
     pub fn debug_type<'cx>(&'cx self, ty: Word) -> SpirvTypePrinter<'cx, 'tcx> {

crates/rustc_codegen_spirv/src/spirv_type.rs

@@ -1,17 +1,18 @@
 use crate::abi::{RecursivePointeeCache, TyLayoutNameKey};
 use crate::builder_spirv::SpirvValue;
 use crate::codegen_cx::CodegenCx;
-use bimap::BiHashMap;
 use indexmap::IndexSet;
 use rspirv::dr::Operand;
 use rspirv::spirv::{Capability, Decoration, Dim, ImageFormat, StorageClass, Word};
 use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::span_bug;
 use rustc_span::def_id::DefId;
 use rustc_span::Span;
 use rustc_target::abi::{Align, Size};
 use std::cell::RefCell;
 use std::fmt;
 use std::iter;
+use std::rc::Rc;
 use std::sync::{LazyLock, Mutex};
 
 /// Spir-v types are represented as simple Words, which are the `result_id` of instructions like
@@ -270,7 +271,7 @@ impl SpirvType {
                 result
             }
         };
-        cx.type_cache.def(result, self);
+        cx.type_cache_def(result, self, def_span);
         result
     }
 
@@ -304,7 +305,7 @@ impl SpirvType {
                 .sess
                 .fatal(&format!("def_with_id invalid for type {:?}", other)),
         };
-        cx.type_cache.def(result, self);
+        cx.type_cache_def(result, self, def_span);
         result
     }
 
@@ -328,6 +329,7 @@ impl SpirvType {
 
     /// Use this if you want a pretty type printing that recursively prints the types within (e.g. struct fields)
     pub fn debug<'cx, 'tcx>(
+        // FIXME(eddyb) don't require a clone of the `SpirvType` here.
         self,
         id: Word,
         cx: &'cx CodegenCx<'tcx>,
@@ -394,6 +396,7 @@ impl SpirvType {
 
 pub struct SpirvTypePrinter<'cx, 'tcx> {
     id: Word,
+    // FIXME(eddyb) don't require a clone of the `SpirvType` here.
     ty: SpirvType,
     cx: &'cx CodegenCx<'tcx>,
 }
@@ -689,8 +692,10 @@ impl SpirvTypePrinter<'_, '_> {
 
 #[derive(Default)]
 pub struct TypeCache<'tcx> {
-    /// Map between ID and structure
-    pub type_defs: RefCell<BiHashMap<Word, SpirvType>>,
+    // FIXME(eddyb) use `tcx.arena.dropless` to get `&'tcx _`, instead of `Rc`.
+    pub id_to_spirv_type: RefCell<FxHashMap<Word, Rc<SpirvType>>>,
+    pub spirv_type_to_id: RefCell<FxHashMap<Rc<SpirvType>, Word>>,
+
     /// Recursive pointer breaking
     pub recursive_pointee_cache: RecursivePointeeCache<'tcx>,
     /// Set of names for a type (only `SpirvType::Adt` currently).
@@ -701,22 +706,53 @@ pub struct TypeCache<'tcx> {
 
 impl TypeCache<'_> {
     fn get(&self, ty: &SpirvType) -> Option<Word> {
-        self.type_defs.borrow().get_by_right(ty).copied()
+        self.spirv_type_to_id.borrow().get(ty).copied()
     }
 
     #[track_caller]
-    pub fn lookup(&self, word: Word) -> SpirvType {
-        self.type_defs
+    pub fn lookup(&self, id: Word) -> Rc<SpirvType> {
+        self.id_to_spirv_type
             .borrow()
-            .get_by_left(&word)
-            .expect("Tried to lookup value that wasn't a type, or has no definition")
+            .get(&id)
+            .expect("tried to lookup ID that wasn't a type, or has no definition")
             .clone()
     }
+}
 
-    fn def(&self, word: Word, ty: SpirvType) {
-        self.type_defs
+impl CodegenCx<'_> {
+    fn type_cache_def(&self, id: Word, ty: SpirvType, def_span: Span) {
+        let ty = Rc::new(ty);
+
+        if let Some(old_ty) = self
+            .type_cache
+            .id_to_spirv_type
             .borrow_mut()
-            .insert_no_overwrite(word, ty)
-            .unwrap();
+            .insert(id, ty.clone())
+        {
+            span_bug!(
+                def_span,
+                "SPIR-V type with ID %{id} is being redefined\n\
+                old type: {old_ty}\n\
+                new type: {ty}",
+                // FIXME(eddyb) don't clone `SpirvType`s here.
+                old_ty = (*old_ty).clone().debug(id, self),
+                ty = (*ty).clone().debug(id, self)
+            );
+        }
+
+        if let Some(old_id) = self
+            .type_cache
+            .spirv_type_to_id
+            .borrow_mut()
+            .insert(ty.clone(), id)
+        {
+            span_bug!(
+                def_span,
+                "SPIR-V type is changing IDs (%{old_id} -> %{id}):\n\
+                {ty}",
+                // FIXME(eddyb) don't clone a `SpirvType` here.
+                ty = (*ty).clone().debug(id, self)
+            );
+        }
     }
 }