diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs index 25c7bd748eebd..e3c84f06543ed 100644 --- a/compiler/rustc_smir/src/rustc_internal/mod.rs +++ b/compiler/rustc_smir/src/rustc_internal/mod.rs @@ -3,8 +3,6 @@ //! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs //! until stable MIR is complete. -use std::ops::{ControlFlow, Index}; - use crate::rustc_internal; use crate::rustc_smir::Tables; use rustc_data_structures::fx; @@ -14,14 +12,18 @@ use rustc_middle::mir::interpret::AllocId; use rustc_middle::ty::TyCtxt; use rustc_span::def_id::{CrateNum, DefId}; use rustc_span::Span; +use stable_mir::ty::IndexedVal; use stable_mir::CompilerError; +use std::fmt::Debug; +use std::hash::Hash; +use std::ops::{ControlFlow, Index}; impl<'tcx> Index for Tables<'tcx> { type Output = DefId; #[inline(always)] fn index(&self, index: stable_mir::DefId) -> &Self::Output { - &self.def_ids.get_index(index.0).unwrap().0 + &self.def_ids[index] } } @@ -30,7 +32,7 @@ impl<'tcx> Index for Tables<'tcx> { #[inline(always)] fn index(&self, index: stable_mir::ty::Span) -> &Self::Output { - &self.spans.get_index(index.0).unwrap().0 + &self.spans[index] } } @@ -96,33 +98,15 @@ impl<'tcx> Tables<'tcx> { } fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId { - if let Some(i) = self.def_ids.get(&did) { - return *i; - } else { - let id = self.def_ids.len(); - self.def_ids.insert(did, stable_mir::DefId(id)); - stable_mir::DefId(id) - } + self.def_ids.create_or_fetch(did) } fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId { - if let Some(i) = self.alloc_ids.get(&aid) { - return *i; - } else { - let id = self.def_ids.len(); - self.alloc_ids.insert(aid, stable_mir::AllocId(id)); - stable_mir::AllocId(id) - } + self.alloc_ids.create_or_fetch(aid) } pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span { - if let Some(i) = self.spans.get(&span) { - return *i; - } else { - let id = self.spans.len(); - self.spans.insert(span, stable_mir::ty::Span(id)); - stable_mir::ty::Span(id) - } + self.spans.create_or_fetch(span) } } @@ -134,9 +118,9 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) { stable_mir::run( Tables { tcx, - def_ids: fx::FxIndexMap::default(), - alloc_ids: fx::FxIndexMap::default(), - spans: fx::FxIndexMap::default(), + def_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() }, + alloc_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() }, + spans: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() }, types: vec![], }, f, @@ -201,3 +185,29 @@ where }) } } + +/// Simmilar to rustc's `FxIndexMap`, `IndexMap` with extra +/// safety features added. +pub struct IndexMap { + index_map: fx::FxIndexMap, +} + +impl IndexMap { + pub fn create_or_fetch(&mut self, key: K) -> V { + let len = self.index_map.len(); + let v = self.index_map.entry(key).or_insert(V::to_val(len)); + *v + } +} + +impl Index + for IndexMap +{ + type Output = K; + + fn index(&self, index: V) -> &Self::Output { + let (k, v) = self.index_map.get_index(index.to_index()).unwrap(); + assert_eq!(*v, index, "Provided value doesn't match with indexed value"); + k + } +} diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs index 757cfa82d3c11..09dc7475ec86b 100644 --- a/compiler/rustc_smir/src/rustc_smir/mod.rs +++ b/compiler/rustc_smir/src/rustc_smir/mod.rs @@ -7,9 +7,9 @@ //! //! For now, we are developing everything inside `rustc`, thus, we keep this module private. +use crate::rustc_internal::IndexMap; use crate::rustc_smir::hir::def::DefKind; use crate::rustc_smir::stable_mir::ty::{BoundRegion, EarlyBoundRegion, Region}; -use rustc_data_structures::fx::FxIndexMap; use rustc_hir as hir; use rustc_middle::mir; use rustc_middle::mir::interpret::{alloc_range, AllocId}; @@ -195,9 +195,9 @@ impl PartialEq for MaybeStable { pub struct Tables<'tcx> { pub tcx: TyCtxt<'tcx>, - pub def_ids: FxIndexMap, - pub alloc_ids: FxIndexMap, - pub spans: FxIndexMap, + pub def_ids: IndexMap, + pub alloc_ids: IndexMap, + pub spans: IndexMap, pub types: Vec>>, } diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs index 104985493efc1..0a5c62eee04a9 100644 --- a/compiler/stable_mir/src/lib.rs +++ b/compiler/stable_mir/src/lib.rs @@ -22,7 +22,8 @@ use std::fmt; use std::fmt::Debug; use self::ty::{ - GenericPredicates, Generics, ImplDef, ImplTrait, Span, TraitDecl, TraitDef, Ty, TyKind, + GenericPredicates, Generics, ImplDef, ImplTrait, IndexedVal, Span, TraitDecl, TraitDef, Ty, + TyKind, }; #[macro_use] @@ -41,7 +42,7 @@ pub type CrateNum = usize; /// A unique identification number for each item accessible for the current compilation unit. #[derive(Clone, Copy, PartialEq, Eq)] -pub struct DefId(pub usize); +pub struct DefId(usize); impl Debug for DefId { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -52,9 +53,28 @@ impl Debug for DefId { } } +impl IndexedVal for DefId { + fn to_val(index: usize) -> Self { + DefId(index) + } + + fn to_index(&self) -> usize { + self.0 + } +} + /// A unique identification number for each provenance #[derive(Clone, Copy, PartialEq, Eq, Debug)] -pub struct AllocId(pub usize); +pub struct AllocId(usize); + +impl IndexedVal for AllocId { + fn to_val(index: usize) -> Self { + AllocId(index) + } + fn to_index(&self) -> usize { + self.0 + } +} /// A list of crate items. pub type CrateItems = Vec; diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs index 6029e3c11adef..691af15da8ca8 100644 --- a/compiler/stable_mir/src/ty.rs +++ b/compiler/stable_mir/src/ty.rs @@ -75,7 +75,7 @@ pub struct Placeholder { } #[derive(Clone, Copy, PartialEq, Eq)] -pub struct Span(pub usize); +pub struct Span(usize); impl Debug for Span { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { @@ -86,6 +86,15 @@ impl Debug for Span { } } +impl IndexedVal for Span { + fn to_val(index: usize) -> Self { + Span(index) + } + fn to_index(&self) -> usize { + self.0 + } +} + #[derive(Clone, Debug)] pub enum TyKind { RigidTy(RigidTy), @@ -565,3 +574,9 @@ pub enum ImplPolarity { Negative, Reservation, } + +pub trait IndexedVal { + fn to_val(index: usize) -> Self; + + fn to_index(&self) -> usize; +}