refactor: split `rustc_smir::alloc` into two parts

The previous `rustc_smir::alloc` had many direct calls to rustc queries.

This commit splits it into two parts: `rustc_smir::alloc` and `stable_mir::alloc`.

Following the same pattern as `SmirCtxt` and `SmirInterface`, the `rustc_smir::alloc` handles all direct interactions with rustc queries and performs the actual memory allocations, while the `stable_mir::alloc` is responsible for constructing stable components.

4 files changed, 151 insertions, 99 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/alloc.rs b/compiler/rustc_smir/src/rustc_smir/alloc.rs
index a6d31ac4e13..34692962de7 100644
--- a/compiler/rustc_smir/src/rustc_smir/alloc.rs
+++ b/compiler/rustc_smir/src/rustc_smir/alloc.rs
@@ -1,102 +1,74 @@
-use rustc_abi::{Align, Size};
-use rustc_middle::mir::ConstValue;
-use rustc_middle::mir::interpret::{AllocInit, AllocRange, Pointer, alloc_range};
-use stable_mir::Error;
-use stable_mir::mir::Mutability;
-use stable_mir::ty::{Allocation, ProvenanceMap};
+//! Internal memory allocator implementation for StableMIR.
+//!
+//! This module handles all direct interactions with rustc queries and performs
+//! the actual memory allocations. The stable interface in `stable_mir::alloc`
+//! delegates all query-related operations to this implementation.
 
-use crate::rustc_smir::{Stable, Tables};
-use crate::stable_mir;
+use rustc_abi::{Size, TyAndLayout};
+use rustc_middle::mir::interpret::{
+    AllocId, AllocInit, AllocRange, Allocation, ConstAllocation, Pointer, Scalar, alloc_range,
+};
+use rustc_middle::ty::Ty;
 
-/// Creates new empty `Allocation` from given `Align`.
-fn new_empty_allocation(align: Align) -> Allocation {
-    Allocation {
-        bytes: Vec::new(),
-        provenance: ProvenanceMap { ptrs: Vec::new() },
-        align: align.bytes(),
-        mutability: Mutability::Not,
-    }
+use crate::rustc_smir::{Bridge, SmirError, Tables};
+
+pub fn try_new_scalar<'tcx, B: Bridge>(
+    layout: TyAndLayout<'tcx, Ty<'tcx>>,
+    scalar: Scalar,
+    tables: &mut Tables<'tcx, B>,
+) -> Result<Allocation, B::Error> {
+    let size = scalar.size();
+    let mut allocation = Allocation::new(size, layout.align.abi, AllocInit::Uninit, ());
+    allocation
+        .write_scalar(&tables.tcx, alloc_range(Size::ZERO, size), scalar)
+        .map_err(|e| B::Error::from_internal(e))?;
+
+    Ok(allocation)
 }
 
-// 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.
-#[allow(rustc::usage_of_qualified_ty)]
-pub(crate) fn new_allocation<'tcx>(
-    ty: rustc_middle::ty::Ty<'tcx>,
-    const_value: ConstValue<'tcx>,
-    tables: &mut Tables<'tcx>,
-) -> Allocation {
-    try_new_allocation(ty, const_value, tables)
-        .unwrap_or_else(|_| panic!("Failed to convert: {const_value:?} to {ty:?}"))
+pub fn try_new_slice<'tcx, B: Bridge>(
+    layout: TyAndLayout<'tcx, Ty<'tcx>>,
+    data: ConstAllocation<'tcx>,
+    meta: u64,
+    tables: &mut Tables<'tcx, B>,
+) -> Result<Allocation, B::Error> {
+    let alloc_id = tables.tcx.reserve_and_set_memory_alloc(data);
+    let ptr = Pointer::new(alloc_id.into(), Size::ZERO);
+    let scalar_ptr = Scalar::from_pointer(ptr, &tables.tcx);
+    let scalar_meta: Scalar = Scalar::from_target_usize(meta, &tables.tcx);
+    let mut allocation = Allocation::new(layout.size, layout.align.abi, AllocInit::Uninit, ());
+    allocation
+        .write_scalar(
+            &tables.tcx,
+            alloc_range(Size::ZERO, tables.tcx.data_layout.pointer_size),
+            scalar_ptr,
+        )
+        .map_err(|e| B::Error::from_internal(e))?;
+    allocation
+        .write_scalar(
+            &tables.tcx,
+            alloc_range(tables.tcx.data_layout.pointer_size, scalar_meta.size()),
+            scalar_meta,
+        )
+        .map_err(|e| B::Error::from_internal(e))?;
+
+    Ok(allocation)
 }
 
-#[allow(rustc::usage_of_qualified_ty)]
-pub(crate) fn try_new_allocation<'tcx>(
-    ty: rustc_middle::ty::Ty<'tcx>,
-    const_value: ConstValue<'tcx>,
-    tables: &mut Tables<'tcx>,
-) -> Result<Allocation, Error> {
-    let layout = tables
-        .tcx
-        .layout_of(rustc_middle::ty::TypingEnv::fully_monomorphized().as_query_input(ty))
-        .map_err(|e| e.stable(tables))?;
-    Ok(match const_value {
-        ConstValue::Scalar(scalar) => {
-            let size = scalar.size();
-            let mut allocation = rustc_middle::mir::interpret::Allocation::new(
-                size,
-                layout.align.abi,
-                AllocInit::Uninit,
-                (),
-            );
-            allocation
-                .write_scalar(&tables.tcx, alloc_range(Size::ZERO, size), scalar)
-                .map_err(|e| e.stable(tables))?;
-            allocation.stable(tables)
-        }
-        ConstValue::ZeroSized => new_empty_allocation(layout.align.abi),
-        ConstValue::Slice { data, meta } => {
-            let alloc_id = tables.tcx.reserve_and_set_memory_alloc(data);
-            let ptr = Pointer::new(alloc_id.into(), Size::ZERO);
-            let scalar_ptr = rustc_middle::mir::interpret::Scalar::from_pointer(ptr, &tables.tcx);
-            let scalar_meta =
-                rustc_middle::mir::interpret::Scalar::from_target_usize(meta, &tables.tcx);
-            let mut allocation = rustc_middle::mir::interpret::Allocation::new(
-                layout.size,
-                layout.align.abi,
-                AllocInit::Uninit,
-                (),
-            );
-            allocation
-                .write_scalar(
-                    &tables.tcx,
-                    alloc_range(Size::ZERO, tables.tcx.data_layout.pointer_size),
-                    scalar_ptr,
-                )
-                .map_err(|e| e.stable(tables))?;
-            allocation
-                .write_scalar(
-                    &tables.tcx,
-                    alloc_range(tables.tcx.data_layout.pointer_size, scalar_meta.size()),
-                    scalar_meta,
-                )
-                .map_err(|e| e.stable(tables))?;
-            allocation.stable(tables)
-        }
-        ConstValue::Indirect { alloc_id, offset } => {
-            let alloc = tables.tcx.global_alloc(alloc_id).unwrap_memory();
-            allocation_filter(&alloc.0, alloc_range(offset, layout.size), tables)
-        }
-    })
+pub fn try_new_indirect<'tcx, B: Bridge>(
+    alloc_id: AllocId,
+    tables: &mut Tables<'tcx, B>,
+) -> ConstAllocation<'tcx> {
+    let alloc = tables.tcx.global_alloc(alloc_id).unwrap_memory();
+
+    alloc
 }
 
 /// Creates an `Allocation` only from information within the `AllocRange`.
-pub(super) fn allocation_filter<'tcx>(
+pub fn allocation_filter(
     alloc: &rustc_middle::mir::interpret::Allocation,
     alloc_range: AllocRange,
-    tables: &mut Tables<'tcx>,
-) -> Allocation {
+) -> (Vec<Option<u8>>, Vec<(usize, AllocId)>) {
     let mut bytes: Vec<Option<u8>> = alloc
         .inspect_with_uninit_and_ptr_outside_interpreter(
             alloc_range.start.bytes_usize()..alloc_range.end().bytes_usize(),
@@ -117,15 +89,8 @@ pub(super) fn allocation_filter<'tcx>(
         .iter()
         .filter(|a| a.0 >= alloc_range.start && a.0 <= alloc_range.end())
     {
-        ptrs.push((
-            offset.bytes_usize() - alloc_range.start.bytes_usize(),
-            tables.prov(prov.alloc_id()),
-        ));
-    }
-    Allocation {
-        bytes,
-        provenance: ProvenanceMap { ptrs },
-        align: alloc.align.bytes(),
-        mutability: alloc.mutability.stable(tables),
+        ptrs.push((offset.bytes_usize() - alloc_range.start.bytes_usize(), prov.alloc_id()));
     }
+
+    (bytes, ptrs)
 }
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index edaec47fe7d..4f6919d2724 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -20,7 +20,7 @@ use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 
 use crate::rustc_internal::IndexMap;
 
-mod alloc;
+pub mod alloc;
 mod builder;
 pub mod context;
 mod convert;
diff --git a/compiler/rustc_smir/src/stable_mir/alloc.rs b/compiler/rustc_smir/src/stable_mir/alloc.rs
new file mode 100644
index 00000000000..4e6facac705
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/alloc.rs
@@ -0,0 +1,86 @@
+//! Memory allocation implementation for StableMIR.
+//!
+//! This module is responsible for constructing stable components.
+//! All operations requiring rustc queries must be delegated
+//! to `rustc_smir::alloc` to maintain stability guarantees.
+
+use rustc_abi::Align;
+use rustc_middle::mir::ConstValue;
+use rustc_middle::mir::interpret::AllocRange;
+use rustc_smir::context::SmirCtxt;
+use rustc_smir::{SmirError, Tables, alloc};
+
+use super::Error;
+use super::compiler_interface::BridgeTys;
+use super::convert::Stable;
+use super::mir::Mutability;
+use super::ty::{Allocation, ProvenanceMap};
+use crate::rustc_smir;
+
+/// Creates new empty `Allocation` from given `Align`.
+fn new_empty_allocation(align: 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.
+#[allow(rustc::usage_of_qualified_ty)]
+pub(crate) fn new_allocation<'tcx>(
+    ty: rustc_middle::ty::Ty<'tcx>,
+    const_value: ConstValue<'tcx>,
+    tables: &mut Tables<'tcx, BridgeTys>,
+    cx: &SmirCtxt<'tcx, BridgeTys>,
+) -> Allocation {
+    try_new_allocation(ty, const_value, tables, cx)
+        .unwrap_or_else(|_| panic!("Failed to convert: {const_value:?} to {ty:?}"))
+}
+
+#[allow(rustc::usage_of_qualified_ty)]
+pub(crate) fn try_new_allocation<'tcx>(
+    ty: rustc_middle::ty::Ty<'tcx>,
+    const_value: ConstValue<'tcx>,
+    tables: &mut Tables<'tcx, BridgeTys>,
+    cx: &SmirCtxt<'tcx, BridgeTys>,
+) -> Result<Allocation, Error> {
+    use rustc_smir::context::SmirTypingEnv;
+    let layout = cx
+        .layout_of(cx.fully_monomorphized().as_query_input(ty))
+        .map_err(|e| Error::from_internal(e))?;
+    match const_value {
+        ConstValue::Scalar(scalar) => {
+            alloc::try_new_scalar(layout, scalar, tables).map(|alloc| alloc.stable(tables, cx))
+        }
+        ConstValue::ZeroSized => Ok(new_empty_allocation(layout.align.abi)),
+        ConstValue::Slice { data, meta } => {
+            alloc::try_new_slice(layout, data, meta, tables).map(|alloc| alloc.stable(tables, cx))
+        }
+        ConstValue::Indirect { alloc_id, offset } => {
+            let alloc = alloc::try_new_indirect(alloc_id, tables);
+            use rustc_smir::context::SmirAllocRange;
+            Ok(allocation_filter(&alloc.0, cx.alloc_range(offset, layout.size), tables, cx))
+        }
+    }
+}
+
+/// Creates an `Allocation` only from information within the `AllocRange`.
+pub(super) fn allocation_filter<'tcx>(
+    alloc: &rustc_middle::mir::interpret::Allocation,
+    alloc_range: AllocRange,
+    tables: &mut Tables<'tcx, BridgeTys>,
+    cx: &SmirCtxt<'tcx, BridgeTys>,
+) -> Allocation {
+    let (bytes, ptrs) = alloc::allocation_filter(alloc, alloc_range);
+    let ptrs = ptrs.iter().map(|(i, aid)| (*i, tables.prov(*aid))).collect();
+    Allocation {
+        bytes,
+        provenance: ProvenanceMap { ptrs },
+        align: alloc.align.bytes(),
+        mutability: alloc.mutability.stable(tables, cx),
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 1c992ec5dbb..d3ef5f4001d 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -40,6 +40,7 @@ use stable_mir::ty::{AssocItem, FnDef, ForeignModuleDef, ImplDef, Span, TraitDef
 use crate::{rustc_smir, stable_mir};
 
 pub mod abi;
+mod alloc;
 #[macro_use]
 pub mod crate_def;
 pub mod compiler_interface;