From 8f1ea576b769d1d8a518d20eabbaf60343c03d81 Mon Sep 17 00:00:00 2001
From: ouz-a <ouz.agz@gmail.com>
Date: Wed, 9 Aug 2023 21:05:03 +0300
Subject: [PATCH] only allocate bytes within AllocRange

---
 compiler/rustc_smir/src/rustc_smir/mod.rs | 30 +++----------
 compiler/rustc_smir/src/stable_mir/ty.rs  | 51 +++++++++++++++++++----
 2 files changed, 47 insertions(+), 34 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index c760fd8b69a6f..07ef53ca30453 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -9,10 +9,13 @@
 
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
-use crate::stable_mir::ty::{new_allocation, FloatTy, IntTy, Movability, RigidTy, TyKind, UintTy};
+use crate::stable_mir::ty::{
+    allocation_filter, 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::interpret::alloc_range;
 use rustc_middle::mir::{self, ConstantKind};
 use rustc_middle::ty::{self, Ty, TyCtxt, Variance};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
@@ -1080,30 +1083,7 @@ impl<'tcx> Stable<'tcx> for mir::interpret::Allocation {
     type T = stable_mir::ty::Allocation;
 
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
-        let size = self.size();
-        let mut bytes: Vec<Option<u8>> = self
-            .inspect_with_uninit_and_ptr_outside_interpreter(0..size.bytes_usize())
-            .iter()
-            .copied()
-            .map(Some)
-            .collect();
-        for (i, b) in bytes.iter_mut().enumerate() {
-            if !self.init_mask().get(rustc_target::abi::Size::from_bytes(i)) {
-                *b = None;
-            }
-        }
-        stable_mir::ty::Allocation {
-            bytes: bytes,
-            provenance: {
-                let mut ptrs = Vec::new();
-                for (size, prov) in self.provenance().ptrs().iter() {
-                    ptrs.push((size.bytes_usize(), opaque(prov)));
-                }
-                stable_mir::ty::ProvenanceMap { ptrs }
-            },
-            align: self.align.bytes(),
-            mutability: self.mutability.stable(tables),
-        }
+        allocation_filter(self, alloc_range(rustc_target::abi::Size::ZERO, self.size()), tables)
     }
 }
 
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index c8bf861be7938..14b841a467034 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -1,8 +1,8 @@
-use rustc_middle::mir::interpret::{alloc_range, ConstValue, Pointer};
+use rustc_middle::mir::interpret::{alloc_range, AllocRange, ConstValue, Pointer};
 
 use super::{mir::Mutability, mir::Safety, with, DefId};
 use crate::{
-    rustc_internal::Opaque,
+    rustc_internal::{opaque, Opaque},
     rustc_smir::{Stable, Tables},
 };
 
@@ -353,17 +353,50 @@ pub fn new_allocation<'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)
+            allocation_filter(&alloc.0, alloc_range(offset, ty_size), tables)
         }
     }
 }
 
+/// Creates an `Allocation` only from information within the `AllocRange`.
+pub fn allocation_filter<'tcx>(
+    alloc: &rustc_middle::mir::interpret::Allocation,
+    alloc_range: AllocRange,
+    tables: &mut Tables<'tcx>,
+) -> Allocation {
+    let mut bytes: Vec<Option<u8>> = alloc
+        .inspect_with_uninit_and_ptr_outside_interpreter(
+            alloc_range.start.bytes_usize()..alloc_range.end().bytes_usize(),
+        )
+        .iter()
+        .copied()
+        .map(Some)
+        .collect();
+    for (i, b) in bytes.iter_mut().enumerate() {
+        if !alloc
+            .init_mask()
+            .get(rustc_target::abi::Size::from_bytes(i + alloc_range.start.bytes_usize()))
+        {
+            *b = None;
+        }
+    }
+    let mut ptrs = Vec::new();
+    for (offset, prov) in alloc
+        .provenance()
+        .ptrs()
+        .iter()
+        .filter(|a| a.0 >= alloc_range.start && a.0 <= alloc_range.end())
+    {
+        ptrs.push((offset.bytes_usize() - alloc_range.start.bytes_usize(), opaque(prov)));
+    }
+    Allocation {
+        bytes: bytes,
+        provenance: ProvenanceMap { ptrs },
+        align: alloc.align.bytes(),
+        mutability: alloc.mutability.stable(tables),
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum ConstantKind {
     Allocated(Allocation),