Convert Const to Allocation in smir

4 files changed, 149 insertions, 3 deletions

diff --git a/compiler/rustc_middle/src/mir/interpret/allocation.rs b/compiler/rustc_middle/src/mir/interpret/allocation.rs
index 2567170f39a..c787481bfbe 100644
--- a/compiler/rustc_middle/src/mir/interpret/allocation.rs
+++ b/compiler/rustc_middle/src/mir/interpret/allocation.rs
@@ -329,6 +329,9 @@ impl<Prov: Provenance, Bytes: AllocBytes> Allocation<Prov, (), Bytes> {
 
     /// Try to create an Allocation of `size` bytes, panics if there is not enough memory
     /// available to the compiler to do so.
+    ///
+    /// Example use case: To obtain an Allocation filled with specific data,
+    /// first call this function and then call write_scalar to fill in the right data.
     pub fn uninit(size: Size, align: Align) -> Self {
         match Self::uninit_inner(size, align, || {
             panic!("Allocation::uninit called with panic_on_fail had allocation failure");
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 39541c845b3..0d9157940f8 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -112,6 +112,10 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::TraitDef(self.create_def_id(did))
     }
 
+    pub fn const_def(&mut self, did: DefId) -> stable_mir::ty::ConstDef {
+        stable_mir::ty::ConstDef(self.create_def_id(did))
+    }
+
     fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
         // FIXME: this becomes inefficient when we have too many ids
         for (i, &d) in self.def_ids.iter().enumerate() {
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index d12de92db8a..c760fd8b69a 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -9,11 +9,11 @@
 
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
-use crate::stable_mir::ty::{FloatTy, IntTy, Movability, RigidTy, TyKind, UintTy};
+use crate::stable_mir::ty::{new_allocation, FloatTy, IntTy, Movability, RigidTy, TyKind, UintTy};
 use crate::stable_mir::{self, Context};
 use rustc_hir as hir;
 use rustc_middle::mir::coverage::CodeRegion;
-use rustc_middle::mir::{self};
+use rustc_middle::mir::{self, ConstantKind};
 use rustc_middle::ty::{self, Ty, TyCtxt, Variance};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc_target::abi::FieldIdx;
@@ -1145,3 +1145,30 @@ impl<'tcx> Stable<'tcx> for ty::TraitDef {
         }
     }
 }
+
+impl<'tcx> Stable<'tcx> for rustc_middle::mir::ConstantKind<'tcx> {
+    type T = stable_mir::ty::ConstantKind;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match self {
+            ConstantKind::Ty(c) => match c.kind() {
+                ty::Value(val) => {
+                    let const_val = tables.tcx.valtree_to_const_val((c.ty(), val));
+                    stable_mir::ty::ConstantKind::Allocated(new_allocation(self, const_val, tables))
+                }
+                _ => todo!(),
+            },
+            ConstantKind::Unevaluated(unev_const, ty) => {
+                stable_mir::ty::ConstantKind::Unevaluated(stable_mir::ty::UnevaluatedConst {
+                    ty: tables.intern_ty(*ty),
+                    def: tables.const_def(unev_const.def),
+                    args: unev_const.args.stable(tables),
+                    promoted: unev_const.promoted.map(|u| u.as_u32()),
+                })
+            }
+            ConstantKind::Val(val, _) => {
+                stable_mir::ty::ConstantKind::Allocated(new_allocation(self, *val, tables))
+            }
+        }
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index c487db5b732..c8bf861be79 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -1,5 +1,10 @@
+use rustc_middle::mir::interpret::{alloc_range, ConstValue, Pointer};
+
 use super::{mir::Mutability, mir::Safety, with, DefId};
-use crate::rustc_internal::Opaque;
+use crate::{
+    rustc_internal::Opaque,
+    rustc_smir::{Stable, Tables},
+};
 
 #[derive(Copy, Clone, Debug)]
 pub struct Ty(pub usize);
@@ -105,6 +110,9 @@ pub struct AliasDef(pub(crate) DefId);
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct TraitDef(pub(crate) DefId);
 
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct ConstDef(pub(crate) DefId);
+
 impl TraitDef {
     pub fn trait_decl(&self) -> TraitDecl {
         with(|cx| cx.trait_decl(self))
@@ -248,6 +256,7 @@ pub type Bytes = Vec<Option<u8>>;
 pub type Size = usize;
 pub type Prov = Opaque;
 pub type Align = u64;
+pub type Promoted = u32;
 pub type InitMaskMaterialized = Vec<u64>;
 
 /// Stores the provenance information of pointers stored in memory.
@@ -266,6 +275,109 @@ pub struct Allocation {
     pub mutability: Mutability,
 }
 
+impl Allocation {
+    /// Creates new empty `Allocation` from given `Align`.
+    fn new_empty_allocation(align: rustc_target::abi::Align) -> Allocation {
+        Allocation {
+            bytes: Vec::new(),
+            provenance: ProvenanceMap { ptrs: Vec::new() },
+            align: align.bytes(),
+            mutability: Mutability::Not,
+        }
+    }
+}
+
+// We need this method instead of a Stable implementation
+// because we need to get `Ty` of the const we are trying to create, to do that
+// we need to have access to `ConstantKind` but we can't access that inside Stable impl.
+pub fn new_allocation<'tcx>(
+    const_kind: &rustc_middle::mir::ConstantKind<'tcx>,
+    const_value: ConstValue<'tcx>,
+    tables: &mut Tables<'tcx>,
+) -> Allocation {
+    match const_value {
+        ConstValue::Scalar(scalar) => {
+            let size = scalar.size();
+            let align = tables
+                .tcx
+                .layout_of(rustc_middle::ty::ParamEnv::reveal_all().and(const_kind.ty()))
+                .unwrap()
+                .align;
+            let mut allocation = rustc_middle::mir::interpret::Allocation::uninit(size, align.abi);
+            allocation
+                .write_scalar(&tables.tcx, alloc_range(rustc_target::abi::Size::ZERO, size), scalar)
+                .unwrap();
+            allocation.stable(tables)
+        }
+        ConstValue::ZeroSized => {
+            let align = tables
+                .tcx
+                .layout_of(rustc_middle::ty::ParamEnv::empty().and(const_kind.ty()))
+                .unwrap()
+                .align;
+            Allocation::new_empty_allocation(align.abi)
+        }
+        ConstValue::Slice { data, start, end } => {
+            let alloc_id = tables.tcx.create_memory_alloc(data);
+            let ptr = Pointer::new(alloc_id, rustc_target::abi::Size::from_bytes(start));
+            let scalar_ptr = rustc_middle::mir::interpret::Scalar::from_pointer(ptr, &tables.tcx);
+            let scalar_len = rustc_middle::mir::interpret::Scalar::from_target_usize(
+                (end - start) as u64,
+                &tables.tcx,
+            );
+            let layout = tables
+                .tcx
+                .layout_of(rustc_middle::ty::ParamEnv::reveal_all().and(const_kind.ty()))
+                .unwrap();
+            let mut allocation =
+                rustc_middle::mir::interpret::Allocation::uninit(layout.size, layout.align.abi);
+            allocation
+                .write_scalar(
+                    &tables.tcx,
+                    alloc_range(rustc_target::abi::Size::ZERO, tables.tcx.data_layout.pointer_size),
+                    scalar_ptr,
+                )
+                .unwrap();
+            allocation
+                .write_scalar(
+                    &tables.tcx,
+                    alloc_range(tables.tcx.data_layout.pointer_size, scalar_len.size()),
+                    scalar_len,
+                )
+                .unwrap();
+            allocation.stable(tables)
+        }
+        ConstValue::ByRef { alloc, offset } => {
+            let ty_size = tables
+                .tcx
+                .layout_of(rustc_middle::ty::ParamEnv::reveal_all().and(const_kind.ty()))
+                .unwrap()
+                .size;
+            let bytes = alloc.0.get_bytes_unchecked(alloc_range(offset, ty_size));
+            let offset_allocation = rustc_middle::mir::interpret::Allocation::from_bytes(
+                bytes,
+                alloc.0.align,
+                alloc.0.mutability,
+            );
+            offset_allocation.stable(tables)
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum ConstantKind {
+    Allocated(Allocation),
+    Unevaluated(UnevaluatedConst),
+}
+
+#[derive(Clone, Debug)]
+pub struct UnevaluatedConst {
+    pub ty: Ty,
+    pub def: ConstDef,
+    pub args: GenericArgs,
+    pub promoted: Option<Promoted>,
+}
+
 pub enum TraitSpecializationKind {
     None,
     Marker,