Extract address generator struct from memory allocator.

3 files changed, 95 insertions, 52 deletions

diff --git a/src/tools/miri/src/alloc_addresses/address_generator.rs b/src/tools/miri/src/alloc_addresses/address_generator.rs
new file mode 100644
index 00000000000..3764f5dcb82
--- /dev/null
+++ b/src/tools/miri/src/alloc_addresses/address_generator.rs
@@ -0,0 +1,78 @@
+use std::ops::Range;
+
+use rand::Rng;
+use rustc_abi::{Align, Size};
+use rustc_const_eval::interpret::{InterpResult, interp_ok};
+use rustc_middle::{err_exhaust, throw_exhaust};
+
+/// Shifts `addr` to make it aligned with `align` by rounding `addr` to the smallest multiple
+/// of `align` that is larger or equal to `addr`
+fn align_addr(addr: u64, align: u64) -> u64 {
+    match addr % align {
+        0 => addr,
+        rem => addr.strict_add(align) - rem,
+    }
+}
+
+/// This provides the logic to generate addresses for memory allocations in a given address range.
+#[derive(Debug)]
+pub struct AddressGenerator {
+    /// This is used as a memory address when a new pointer is casted to an integer. It
+    /// is always larger than any address that was previously made part of a block.
+    next_base_addr: u64,
+    /// This is the last address that can be allocated.
+    end: u64,
+}
+
+impl AddressGenerator {
+    pub fn new(addr_range: Range<u64>) -> Self {
+        Self { next_base_addr: addr_range.start, end: addr_range.end }
+    }
+
+    /// Get the remaining range where this `AddressGenerator` can still allocate addresses.
+    pub fn get_remaining(&self) -> Range<u64> {
+        self.next_base_addr..self.end
+    }
+
+    /// Generate a new address with the specified size and alignment, using the given Rng to add some randomness.
+    /// The returned allocation is guaranteed not to overlap with any address ranges given out by the generator before.
+    /// Returns an error if the allocation request cannot be fulfilled.
+    pub fn generate<'tcx, R: Rng>(
+        &mut self,
+        size: Size,
+        align: Align,
+        rng: &mut R,
+    ) -> InterpResult<'tcx, u64> {
+        // Leave some space to the previous allocation, to give it some chance to be less aligned.
+        // We ensure that `(self.next_base_addr + slack) % 16` is uniformly distributed.
+        let slack = rng.random_range(0..16);
+        // From next_base_addr + slack, round up to adjust for alignment.
+        let base_addr =
+            self.next_base_addr.checked_add(slack).ok_or_else(|| err_exhaust!(AddressSpaceFull))?;
+        let base_addr = align_addr(base_addr, align.bytes());
+
+        // Remember next base address.  If this allocation is zero-sized, leave a gap of at
+        // least 1 to avoid two allocations having the same base address. (The logic in
+        // `alloc_id_from_addr` assumes unique addresses, and different function/vtable pointers
+        // need to be distinguishable!)
+        self.next_base_addr = base_addr
+            .checked_add(size.bytes().max(1))
+            .ok_or_else(|| err_exhaust!(AddressSpaceFull))?;
+        // Even if `Size` didn't overflow, we might still have filled up the address space.
+        if self.next_base_addr > self.end {
+            throw_exhaust!(AddressSpaceFull);
+        }
+        interp_ok(base_addr)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_align_addr() {
+        assert_eq!(align_addr(37, 4), 40);
+        assert_eq!(align_addr(44, 4), 44);
+    }
+}
diff --git a/src/tools/miri/src/alloc_addresses/mod.rs b/src/tools/miri/src/alloc_addresses/mod.rs
index 334503d2994..afd8afbaeab 100644
--- a/src/tools/miri/src/alloc_addresses/mod.rs
+++ b/src/tools/miri/src/alloc_addresses/mod.rs
@@ -1,15 +1,16 @@
 //! This module is responsible for managing the absolute addresses that allocations are located at,
 //! and for casting between pointers and integers based on those addresses.
 
+mod address_generator;
 mod reuse_pool;
 
 use std::cell::RefCell;
-use std::cmp::max;
 
-use rand::Rng;
 use rustc_abi::{Align, Size};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_middle::ty::TyCtxt;
 
+pub use self::address_generator::AddressGenerator;
 use self::reuse_pool::ReusePool;
 use crate::concurrency::VClock;
 use crate::*;
@@ -49,9 +50,8 @@ pub struct GlobalStateInner {
     /// Whether an allocation has been exposed or not. This cannot be put
     /// into `AllocExtra` for the same reason as `base_addr`.
     exposed: FxHashSet<AllocId>,
-    /// This is used as a memory address when a new pointer is casted to an integer. It
-    /// is always larger than any address that was previously made part of a block.
-    next_base_addr: u64,
+    /// The generator for new addresses in a given range.
+    address_generator: AddressGenerator,
     /// The provenance to use for int2ptr casts
     provenance_mode: ProvenanceMode,
 }
@@ -64,7 +64,7 @@ impl VisitProvenance for GlobalStateInner {
             prepared_alloc_bytes: _,
             reuse: _,
             exposed: _,
-            next_base_addr: _,
+            address_generator: _,
             provenance_mode: _,
         } = self;
         // Though base_addr, int_to_ptr_map, and exposed contain AllocIds, we do not want to visit them.
@@ -77,14 +77,14 @@ impl VisitProvenance for GlobalStateInner {
 }
 
 impl GlobalStateInner {
-    pub fn new(config: &MiriConfig, stack_addr: u64) -> Self {
+    pub fn new<'tcx>(config: &MiriConfig, stack_addr: u64, tcx: TyCtxt<'tcx>) -> Self {
         GlobalStateInner {
             int_to_ptr_map: Vec::default(),
             base_addr: FxHashMap::default(),
             prepared_alloc_bytes: FxHashMap::default(),
             reuse: ReusePool::new(config),
             exposed: FxHashSet::default(),
-            next_base_addr: stack_addr,
+            address_generator: AddressGenerator::new(stack_addr..tcx.target_usize_max()),
             provenance_mode: config.provenance_mode,
         }
     }
@@ -96,15 +96,6 @@ impl GlobalStateInner {
     }
 }
 
-/// Shifts `addr` to make it aligned with `align` by rounding `addr` to the smallest multiple
-/// of `align` that is larger or equal to `addr`
-fn align_addr(addr: u64, align: u64) -> u64 {
-    match addr % align {
-        0 => addr,
-        rem => addr.strict_add(align) - rem,
-    }
-}
-
 impl<'tcx> EvalContextExtPriv<'tcx> for crate::MiriInterpCx<'tcx> {}
 trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn addr_from_alloc_id_uncached(
@@ -194,34 +185,17 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
             interp_ok(reuse_addr)
         } else {
             // We have to pick a fresh address.
-            // Leave some space to the previous allocation, to give it some chance to be less aligned.
-            // We ensure that `(global_state.next_base_addr + slack) % 16` is uniformly distributed.
-            let slack = rng.random_range(0..16);
-            // From next_base_addr + slack, round up to adjust for alignment.
-            let base_addr = global_state
-                .next_base_addr
-                .checked_add(slack)
-                .ok_or_else(|| err_exhaust!(AddressSpaceFull))?;
-            let base_addr = align_addr(base_addr, info.align.bytes());
-
-            // Remember next base address.  If this allocation is zero-sized, leave a gap of at
-            // least 1 to avoid two allocations having the same base address. (The logic in
-            // `alloc_id_from_addr` assumes unique addresses, and different function/vtable pointers
-            // need to be distinguishable!)
-            global_state.next_base_addr = base_addr
-                .checked_add(max(info.size.bytes(), 1))
-                .ok_or_else(|| err_exhaust!(AddressSpaceFull))?;
-            // Even if `Size` didn't overflow, we might still have filled up the address space.
-            if global_state.next_base_addr > this.target_usize_max() {
-                throw_exhaust!(AddressSpaceFull);
-            }
+            let new_addr =
+                global_state.address_generator.generate(info.size, info.align, &mut rng)?;
+
             // If we filled up more than half the address space, start aggressively reusing
             // addresses to avoid running out.
-            if global_state.next_base_addr > u64::try_from(this.target_isize_max()).unwrap() {
+            let remaining_range = global_state.address_generator.get_remaining();
+            if remaining_range.start > remaining_range.end / 2 {
                 global_state.reuse.address_space_shortage();
             }
 
-            interp_ok(base_addr)
+            interp_ok(new_addr)
         }
     }
 }
@@ -519,14 +493,3 @@ impl<'tcx> MiriMachine<'tcx> {
         })
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_align_addr() {
-        assert_eq!(align_addr(37, 4), 40);
-        assert_eq!(align_addr(44, 4), 44);
-    }
-}
diff --git a/src/tools/miri/src/machine.rs b/src/tools/miri/src/machine.rs
index 8f0814a070c..9482769a2fc 100644
--- a/src/tools/miri/src/machine.rs
+++ b/src/tools/miri/src/machine.rs
@@ -674,11 +674,13 @@ impl<'tcx> MiriMachine<'tcx> {
             thread_cpu_affinity
                 .insert(threads.active_thread(), CpuAffinityMask::new(&layout_cx, config.num_cpus));
         }
+        let alloc_addresses =
+            RefCell::new(alloc_addresses::GlobalStateInner::new(config, stack_addr, tcx));
         MiriMachine {
             tcx,
             borrow_tracker,
             data_race,
-            alloc_addresses: RefCell::new(alloc_addresses::GlobalStateInner::new(config, stack_addr)),
+            alloc_addresses,
             // `env_vars` depends on a full interpreter so we cannot properly initialize it yet.
             env_vars: EnvVars::default(),
             main_fn_ret_place: None,