Refactor ChunkState to encode space index

Author: qinsoon

src/policy/immix/immixspace.rs

@@ -524,7 +524,7 @@ impl<VM: VMBinding> ImmixSpace<VM> {
         self.defrag.notify_new_clean_block(copy);
         let block = Block::from_aligned_address(block_address);
         block.init(copy);
-        self.chunk_map.set(block.chunk(), ChunkState::Allocated);
+        self.chunk_map.set(block.chunk(), ChunkState::allocated(self.common().descriptor.get_index()));
         self.lines_consumed
             .fetch_add(Block::LINES, Ordering::SeqCst);
         Some(block)
@@ -899,7 +899,7 @@ struct SweepChunk<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, mmtk: &'static MMTK<VM>) {
-        assert_eq!(self.space.chunk_map.get(self.chunk), ChunkState::Allocated);
+        assert!(self.space.chunk_map.get(self.chunk).is_allocated());
 
         let mut histogram = self.space.defrag.new_histogram();
         let line_mark_state = if super::BLOCK_ONLY {
@@ -950,7 +950,7 @@ impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
         probe!(mmtk, sweep_chunk, allocated_blocks);
         // Set this chunk as free if there is not live blocks.
         if allocated_blocks == 0 {
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set(self.chunk, ChunkState::free())
         }
         self.space.defrag.add_completed_mark_histogram(histogram);
         self.epilogue.finish_one_work_packet();

src/policy/marksweepspace/native_ms/global.rs

@@ -402,7 +402,7 @@ impl<VM: VMBinding> MarkSweepSpace<VM> {
 
     pub fn record_new_block(&self, block: Block) {
         block.init();
-        self.chunk_map.set(block.chunk(), ChunkState::Allocated);
+        self.chunk_map.set(block.chunk(), ChunkState::allocated(self.common.descriptor.get_index()));
     }
 
     pub fn prepare(&mut self, full_heap: bool) {
@@ -567,7 +567,7 @@ struct PrepareChunkMap<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for PrepareChunkMap<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, _mmtk: &'static MMTK<VM>) {
-        debug_assert!(self.space.chunk_map.get(self.chunk) == ChunkState::Allocated);
+        debug_assert!(self.space.chunk_map.get(self.chunk).is_allocated());
         // number of allocated blocks.
         let mut n_occupied_blocks = 0;
         self.chunk
@@ -581,7 +581,7 @@ impl<VM: VMBinding> GCWork<VM> for PrepareChunkMap<VM> {
             });
         if n_occupied_blocks == 0 {
             // Set this chunk as free if there is no live blocks.
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set(self.chunk, ChunkState::free())
         } else {
             // Otherwise this chunk is occupied, and we reset the mark bit if it is on the side.
             if let MetadataSpec::OnSide(side) = *VM::VMObjectModel::LOCAL_MARK_BIT_SPEC {
@@ -617,7 +617,7 @@ struct SweepChunk<VM: VMBinding> {
 
 impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
     fn do_work(&mut self, _worker: &mut GCWorker<VM>, _mmtk: &'static MMTK<VM>) {
-        assert_eq!(self.space.chunk_map.get(self.chunk), ChunkState::Allocated);
+        assert!(self.space.chunk_map.get(self.chunk).is_allocated());
 
         // number of allocated blocks.
         let mut allocated_blocks = 0;
@@ -636,7 +636,7 @@ impl<VM: VMBinding> GCWork<VM> for SweepChunk<VM> {
         probe!(mmtk, sweep_chunk, allocated_blocks);
         // Set this chunk as free if there is not live blocks.
         if allocated_blocks == 0 {
-            self.space.chunk_map.set(self.chunk, ChunkState::Free)
+            self.space.chunk_map.set(self.chunk, ChunkState::free());
         }
         self.epilogue.finish_one_work_packet();
     }

src/util/heap/chunk_map.rs

@@ -45,13 +45,44 @@ impl Chunk {
 }
 
 /// Chunk allocation state
-#[repr(u8)]
-#[derive(Debug, PartialEq, Clone, Copy)]
-pub enum ChunkState {
-    /// The chunk is not allocated.
-    Free = 0,
-    /// The chunk is allocated.
-    Allocated = 1,
+/// Highest bit: 0 = free, 1 = allocated
+/// Lower 4 bits: Space index (0-15)
+#[repr(transparent)]
+#[derive(PartialEq, Clone, Copy)]
+pub struct ChunkState(u8);
+
+impl ChunkState {
+    pub fn allocated(space_index: usize) -> ChunkState {
+        debug_assert!(space_index < crate::util::heap::layout::heap_parameters::MAX_SPACES);
+        let mut encode = space_index as u8;
+        encode |= 0x80;
+        ChunkState(encode)
+    }
+    pub fn free() -> ChunkState {
+        ChunkState(0)
+    }
+    pub fn is_free(&self) -> bool {
+        self.0 & 0x80 == 0
+    }
+    pub fn is_allocated(&self) -> bool {
+        !self.is_free()
+    }
+    pub fn get_space_index(&self) -> usize {
+        debug_assert!(self.is_allocated());
+        let index = (self.0 & 0x0F) as usize;
+        debug_assert!(index < crate::util::heap::layout::heap_parameters::MAX_SPACES);
+        index
+    }
+}
+
+impl std::fmt::Debug for ChunkState {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.is_free() {
+            write!(f, "Free")
+        } else {
+            write!(f, "Allocated in space {}", self.get_space_index())
+        }
+    }
 }
 
 /// A byte-map to record all the allocated chunks.
@@ -78,10 +109,17 @@ impl ChunkMap {
         if self.get(chunk) == state {
             return;
         }
+        #[cfg(debug_assertions)]
+        {
+            let old_state = self.get(chunk);
+            if state.is_allocated() {
+                assert!(old_state.is_free() || old_state.get_space_index() == state.get_space_index(), "Chunk {:?}: old state {:?}, new state {:?}. Cannot set to new state.", chunk, old_state, state);
+            }
+        }
         // Update alloc byte
-        unsafe { Self::ALLOC_TABLE.store::<u8>(chunk.start(), state as u8) };
+        unsafe { Self::ALLOC_TABLE.store::<u8>(chunk.start(), state.0) };
         // If this is a newly allcoated chunk, then expand the chunk range.
-        if state == ChunkState::Allocated {
+        if state.is_allocated() {
             debug_assert!(!chunk.start().is_zero());
             let mut range = self.chunk_range.lock();
             if range.start == Chunk::ZERO {
@@ -99,11 +137,7 @@ impl ChunkMap {
     /// Get chunk state
     pub fn get(&self, chunk: Chunk) -> ChunkState {
         let byte = unsafe { Self::ALLOC_TABLE.load::<u8>(chunk.start()) };
-        match byte {
-            0 => ChunkState::Free,
-            1 => ChunkState::Allocated,
-            _ => unreachable!(),
-        }
+        ChunkState(byte)
     }
 
     /// A range of all chunks in the heap.
@@ -120,7 +154,7 @@ impl ChunkMap {
         let mut work_packets: Vec<Box<dyn GCWork<VM>>> = vec![];
         for chunk in self
             .all_chunks()
-            .filter(|c| self.get(*c) == ChunkState::Allocated)
+            .filter(|c| self.get(*c).is_allocated())
         {
             work_packets.push(func(chunk));
         }

src/util/object_enum.rs

@@ -85,7 +85,7 @@ pub(crate) fn enumerate_blocks_from_chunk_map<B>(
     B: BlockMayHaveObjects,
 {
     for chunk in chunk_map.all_chunks() {
-        if chunk_map.get(chunk) == ChunkState::Allocated {
+        if chunk_map.get(chunk).is_allocated() {
             for block in chunk.iter_region::<B>() {
                 if block.may_have_objects() {
                     enumerator.visit_address_range(block.start(), block.end());

src/vm/tests/mock_tests/mock_test_allocate_nonmoving.rs

@@ -0,0 +1,44 @@
+// GITHUB-CI: MMTK_PLAN=all
+
+use lazy_static::lazy_static;
+
+use super::mock_test_prelude::*;
+use crate::plan::AllocationSemantics;
+
+#[test]
+pub fn allocate_alignment() {
+    with_mockvm(
+        || -> MockVM {
+            MockVM {
+                is_collection_enabled: MockMethod::new_fixed(Box::new(|_| false)),
+                ..MockVM::default()
+            }
+        },
+        || {
+            // 1MB heap
+            const MB: usize = 1024 * 1024;
+            let mut fixture = MutatorFixture::create_with_heapsize(MB);
+
+            // Normal alloc
+            let addr = memory_manager::alloc(
+                &mut fixture.mutator,
+                16,
+                8,
+                0,
+                AllocationSemantics::Default,
+            );
+            assert!(!addr.is_zero());
+
+            // Non moving alloc
+            let addr = memory_manager::alloc(
+                &mut fixture.mutator,
+                16,
+                8,
+                0,
+                AllocationSemantics::NonMoving,
+            );
+            assert!(!addr.is_zero());
+        },
+        no_cleanup,
+    )
+}

src/vm/tests/mock_tests/mod.rs

@@ -24,6 +24,7 @@ pub(crate) mod mock_test_prelude {
 }
 
 mod mock_test_allocate_align_offset;
+mod mock_test_allocate_nonmoving;
 mod mock_test_allocate_with_disable_collection;
 mod mock_test_allocate_with_initialize_collection;
 mod mock_test_allocate_with_re_enable_collection;