Replicate vulkano-rs#2098 on master (vulkano-rs#2110)

hakolao · Feb 20, 2024 · dd54baa · dd54baa
1 parent 4b918a6
commit dd54baa
Showing 1 changed file with 44 additions and 0 deletions.
diff --git a/vulkano/src/memory/allocator/suballocator.rs b/vulkano/src/memory/allocator/suballocator.rs
@@ -497,6 +497,12 @@ impl MemoryAlloc {
     /// [`shift`]: Self::shift
     #[inline]
     pub unsafe fn set_offset(&mut self, new_offset: DeviceSize) {
+        if let Some(ptr) = self.mapped_ptr.as_mut() {
+            *ptr = NonNull::new_unchecked(
+                ptr.as_ptr()
+                    .offset(new_offset as isize - self.offset as isize),
+            );
+        }
         self.offset = new_offset;
     }
 
@@ -2591,6 +2597,44 @@ mod tests {
         ..DUMMY_INFO
     };
 
+    #[test]
+    fn memory_alloc_set_offset() {
+        let (device, _) = gfx_dev_and_queue!();
+        let memory_type_index = device
+            .physical_device()
+            .memory_properties()
+            .memory_types
+            .iter()
+            .position(|memory_type| {
+                memory_type
+                    .property_flags
+                    .contains(MemoryPropertyFlags::HOST_VISIBLE)
+            })
+            .unwrap() as u32;
+        let mut alloc = MemoryAlloc::new(
+            DeviceMemory::allocate(
+                device,
+                MemoryAllocateInfo {
+                    memory_type_index,
+                    allocation_size: 1024,
+                    ..Default::default()
+                },
+            )
+            .unwrap(),
+        )
+        .unwrap();
+        let ptr = alloc.mapped_ptr().unwrap().as_ptr();
+
+        unsafe {
+            alloc.set_offset(16);
+            assert_eq!(alloc.mapped_ptr().unwrap().as_ptr(), ptr.offset(16));
+            alloc.set_offset(0);
+            assert_eq!(alloc.mapped_ptr().unwrap().as_ptr(), ptr.offset(0));
+            alloc.set_offset(32);
+            assert_eq!(alloc.mapped_ptr().unwrap().as_ptr(), ptr.offset(32));
+        }
+    }
+
     #[test]
     fn free_list_allocator_capacity() {
         const THREADS: DeviceSize = 12;