feat(wayland/shm): impl shm ontop of dmabuf (again), and make the upload async

Signed-off-by: Schmarni <marnistromer@gmail.com>

src/wayland/core/buffer.rs

@@ -62,11 +62,20 @@ impl Buffer {
 	) -> Option<Handle<Image>> {
 		tracing::debug!("Updating texture for buffer {:?}", self.id);
 		match &self.backing {
-			BufferBacking::Shm(backing) => backing.update_tex(images),
+			BufferBacking::Shm(backing) => backing.update_tex(dmatexes, images),
 			BufferBacking::Dmabuf(backing) => backing.update_tex(dmatexes, images),
 		}
 	}
 
+	#[tracing::instrument(level = "debug", skip_all)]
+	pub fn on_commit(&self) {
+		tracing::debug!("running on_commit for buffer {:?}", self.id);
+		match &self.backing {
+			BufferBacking::Shm(backing) => backing.on_commit(),
+			BufferBacking::Dmabuf(_backing) => {}
+		}
+	}
+
 	pub fn is_transparent(&self) -> bool {
 		match &self.backing {
 			BufferBacking::Shm(backing) => backing.is_transparent(),
@@ -81,7 +90,10 @@ impl Buffer {
 		}
 	}
 	pub fn uses_buffer_usage(&self) -> bool {
-		matches!(self.backing, BufferBacking::Dmabuf(_))
+		matches!(
+			self.backing,
+			BufferBacking::Dmabuf(_) | BufferBacking::Shm(_)
+		)
 	}
 }
 

src/wayland/core/shm_buffer_backing.rs

@@ -1,20 +1,66 @@
+use crate::wayland::{RENDER_DEVICE, vulkano_data::VULKANO_CONTEXT};
+
 use super::shm_pool::ShmPool;
 use bevy::{
 	asset::{Assets, Handle, RenderAssetUsages},
 	image::Image,
 	render::render_resource::{Extent3d, TextureDimension, TextureFormat},
 };
+use bevy_dmabuf::{
+	dmatex::{Dmatex, DmatexPlane, Resolution},
+	format_mapping::{drm_fourcc_to_vk_format, vk_format_to_drm_fourcc},
+	import::{DropCallback, ImportedDmatexs, ImportedTexture, import_texture},
+};
+use drm_fourcc::DrmFourcc;
 use mint::Vector2;
-use std::sync::Arc;
+use parking_lot::Mutex;
+use std::{
+	os::fd::OwnedFd,
+	sync::{Arc, OnceLock},
+};
+use tracing::debug_span;
+use vulkano::{
+	buffer::{BufferCreateFlags, BufferUsage},
+	command_buffer::{
+		AutoCommandBufferBuilder, CommandBufferUsage, CopyBufferToImageInfo,
+		PrimaryCommandBufferAbstract,
+	},
+	image::{
+		ImageAspect, ImageCreateFlags, ImageCreateInfo, ImageMemory, ImageTiling, ImageUsage,
+		sys::RawImage,
+	},
+	memory::{
+		DedicatedAllocation, DeviceMemory, ExternalMemoryHandleType, MemoryAllocateInfo,
+		MemoryPropertyFlags, ResourceMemory,
+		allocator::{AllocationCreateInfo, MemoryTypeFilter},
+	},
+	sync::GpuFuture,
+};
 use waynest::server::protocol::core::wayland::wl_shm::Format;
 /// Parameters for a shared memory buffer
-#[derive(Debug)]
 pub struct ShmBufferBacking {
 	pool: Arc<ShmPool>,
 	offset: usize,
 	stride: usize,
 	size: Vector2<usize>,
-	format: Format,
+	wl_format: Format,
+	image: Arc<vulkano::image::Image>,
+	tex: OnceLock<Handle<Image>>,
+	pending_imported_dmatex: Mutex<Option<ImportedTexture>>,
+}
+
+impl std::fmt::Debug for ShmBufferBacking {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("ShmBufferBacking")
+			.field("pool", &self.pool)
+			.field("offset", &self.offset)
+			.field("stride", &self.stride)
+			.field("size", &self.size)
+			.field("wl_format", &self.wl_format)
+			.field("image", &self.image)
+			.field("tex", &self.tex)
+			.finish()
+	}
 }
 
 impl ShmBufferBacking {
@@ -23,70 +69,198 @@ impl ShmBufferBacking {
 		offset: usize,
 		stride: usize,
 		size: Vector2<usize>,
-		format: Format,
+		wl_format: Format,
 	) -> Self {
+		let vk = VULKANO_CONTEXT.wait();
+		let format = match wl_format {
+			Format::Argb8888 | Format::Xrgb8888 => vulkano::format::Format::B8G8R8A8_SRGB,
+			_ => unimplemented!(),
+		};
+		let modifiers = vk
+			.phys_dev
+			.format_properties(format)
+			.unwrap()
+			.drm_format_modifier_properties
+			.into_iter()
+			.filter_map(|v| {
+				(v.drm_format_modifier_plane_count == 1).then_some(v.drm_format_modifier)
+			})
+			.collect();
+		let raw_image = RawImage::new(
+			vk.dev.clone(),
+			ImageCreateInfo {
+				flags: ImageCreateFlags::empty(),
+				image_type: vulkano::image::ImageType::Dim2d,
+				format,
+				extent: [size.x as u32, size.y as u32, 1],
+				tiling: ImageTiling::DrmFormatModifier,
+				usage: ImageUsage::TRANSFER_DST,
+				drm_format_modifiers: modifiers,
+				external_memory_handle_types: ExternalMemoryHandleType::DmaBuf.into(),
+				..Default::default()
+			},
+		)
+		.unwrap();
+		let (modifier, num_planes) = raw_image.drm_format_modifier().unwrap();
+
+		let mem_reqs = raw_image.memory_requirements()[0];
+
+		let index = vk
+			.phys_dev
+			.memory_properties()
+			.memory_types
+			.iter()
+			.enumerate()
+			.map(|(i, _v)| i as u32)
+			.find(|i| mem_reqs.memory_type_bits & (1 << i) != 0)
+			.expect("no valid memory type");
+
+		let mem = ResourceMemory::new_dedicated(
+			DeviceMemory::allocate(
+				vk.dev.clone(),
+				MemoryAllocateInfo {
+					allocation_size: mem_reqs.layout.size(),
+					memory_type_index: index,
+					dedicated_allocation: Some(DedicatedAllocation::Image(&raw_image)),
+					export_handle_types: ExternalMemoryHandleType::DmaBuf.into(),
+					..Default::default()
+				},
+			)
+			.unwrap(),
+		);
+		let Ok(image) = raw_image.bind_memory([mem]) else {
+			panic!("unable to bind memory")
+		};
+		let image = Arc::new(image);
+		let ImageMemory::Normal(mem) = image.memory() else {
+			unreachable!()
+		};
+
+		let [mem] = mem.as_slice() else {
+			unreachable!()
+		};
+
+		let fd = OwnedFd::from(
+			mem.device_memory()
+				.export_fd(ExternalMemoryHandleType::DmaBuf)
+				.unwrap(),
+		);
+
+		let planes = (0..num_planes)
+			.filter_map(|i| {
+				Some(match i {
+					0 => ImageAspect::MemoryPlane0,
+					1 => ImageAspect::MemoryPlane1,
+					2 => ImageAspect::MemoryPlane2,
+					3 => ImageAspect::MemoryPlane3,
+					_ => return None,
+				})
+			})
+			.map(|aspect| {
+				let plane_layout = image.subresource_layout(aspect, 0, 0).unwrap();
+
+				DmatexPlane {
+					dmabuf_fd: fd.try_clone().unwrap().into(),
+					modifier,
+					offset: plane_layout.offset as u32,
+					stride: plane_layout.row_pitch as i32,
+				}
+			})
+			.collect::<Vec<_>>();
+
+		let dmatex = Dmatex {
+			planes,
+			res: Resolution {
+				x: size.x as u32,
+				y: size.y as u32,
+			},
+			format: DrmFourcc::Argb8888 as u32,
+			flip_y: false,
+			srgb: true,
+		};
+		let imported_dmatex =
+			import_texture(RENDER_DEVICE.wait(), dmatex, DropCallback(None)).unwrap();
 		Self {
 			pool,
 			offset,
 			stride,
 			size,
-			format,
+			wl_format,
+			image,
+			pending_imported_dmatex: Mutex::new(Some(imported_dmatex)),
+			tex: OnceLock::new(),
 		}
 	}
+	pub fn on_commit(&self) {
+		let vk = VULKANO_CONTEXT.wait();
+		let data_len = self.size.x * self.size.y * 4;
+		let gpu_buffer = vulkano::buffer::Buffer::new_slice::<u8>(
+			vk.alloc.clone(),
+			vulkano::buffer::BufferCreateInfo {
+				usage: BufferUsage::TRANSFER_SRC,
+				..Default::default()
+			},
+			AllocationCreateInfo {
+				memory_type_filter: MemoryTypeFilter::HOST_SEQUENTIAL_WRITE,
+				..Default::default()
+			},
+			data_len as u64,
+		)
+		.unwrap();
+		{
+			let _span = debug_span!("copy to gpu buffer").entered();
+			let shm_data_lock = self.pool.data_lock();
+			let mut gpu_slice = gpu_buffer.write().unwrap();
+			for (shm_offset, gpu_offset) in
+				(0..self.size.y).map(|v| (self.offset + (v * self.stride), (v * (self.size.x * 4))))
+			{
+				let line_slice = &shm_data_lock[shm_offset..(shm_offset + (self.size.x * 4))];
+				let gpu_subslice = &mut gpu_slice[gpu_offset..(gpu_offset + (self.size.x * 4))];
+				gpu_subslice.copy_from_slice(line_slice);
+			}
+		}
+		let mut command_buffer = AutoCommandBufferBuilder::primary(
+			vk.command_buffer_alloc.clone(),
+			vk.queue.queue_family_index(),
+			CommandBufferUsage::OneTimeSubmit,
+		)
+		.unwrap();
+
+		command_buffer
+			.copy_buffer_to_image(CopyBufferToImageInfo::buffer_image(
+				gpu_buffer.clone(),
+				self.image.clone(),
+			))
+			.unwrap();
+
+		let command_buffer = command_buffer.build().unwrap();
+		command_buffer
+			.execute(vk.queue.clone())
+			.unwrap()
+			.then_signal_fence_and_flush()
+			.unwrap()
+			.wait(None)
+			.unwrap();
+	}
 
 	#[tracing::instrument("debug", skip_all)]
-	pub fn update_tex(&self, images: &mut Assets<Image>) -> Option<Handle<Image>> {
+	pub fn update_tex(
-		let src_data_lock = self.pool.data_lock();
+		&self,
-		let mut src_cursor = self.offset;
+		dmatexes: &ImportedDmatexs,
-
+		images: &mut Assets<Image>,
-		// Calculate maximum cursor position needed - stride is already in bytes
+	) -> Option<Handle<Image>> {
-		let max_cursor = self.offset + (self.size.y * self.stride);
+		self.pending_imported_dmatex
-
+			.lock()
-		// Check if we have enough data
+			.take()
-		if max_cursor > src_data_lock.len() {
+			.map(|tex| dmatexes.insert_imported_dmatex(images, tex))
-			return None;
+			.inspect(|handle| {
-		}
+				_ = self.tex.set(handle.clone());
-		let data_len = self.size.x * self.size.y * 4;
+			});
-		if src_data_lock.len() < data_len {
+		self.tex.get().cloned()
-			return None;
-		}
-		let mut dst_cursor = 0;
-
-		let mut dst_data = vec![0u8; data_len];
-		for _y in 0..self.size.y {
-			for _x in 0..self.size.x {
-				match self.format {
-					Format::Argb8888 | Format::Xrgb8888 => {
-						dst_data[dst_cursor] = src_data_lock[src_cursor + 2]; // Red is byte 2
-						dst_data[dst_cursor + 1] = src_data_lock[src_cursor + 1]; // Green is byte 1
-						dst_data[dst_cursor + 2] = src_data_lock[src_cursor]; // Blue is byte 0
-						dst_data[dst_cursor + 3] = src_data_lock[src_cursor + 3]; // Alpha is byte 3
-					}
-					_ => panic!("Unsupported format {:?}", self.format),
-				}
-				src_cursor += 4;
-				dst_cursor += 4;
-			}
-			src_cursor += self.stride - (self.size.x * 4);
-		}
-
-		let image = Image::new(
-			Extent3d {
-				width: self.size().x as u32,
-				height: self.size().y as u32,
-				depth_or_array_layers: 1,
-			},
-			TextureDimension::D2,
-			dst_data,
-			TextureFormat::Rgba8UnormSrgb,
-			RenderAssetUsages::RENDER_WORLD,
-		);
-
-		Some(images.add(image))
 	}
 
 	pub fn is_transparent(&self) -> bool {
-		match self.format {
+		match self.wl_format {
 			Format::Xrgb8888 => false,
 			Format::Argb8888 => true,
 			_ => true,

src/wayland/core/surface.rs

@@ -23,6 +23,7 @@ use bevy_dmabuf::import::ImportedDmatexs;
 use mint::Vector2;
 use parking_lot::Mutex;
 use std::sync::{Arc, OnceLock, Weak};
+use tokio::task::LocalSet;
 use waynest::{
 	server::{
 		Client, Dispatcher, Result,
@@ -370,6 +371,19 @@ impl WlSurface for Surface {
 	/// https://wayland.app/protocols/wayland#wl_surface:request:commit
 	#[tracing::instrument(level = "debug", skip_all)]
 	async fn commit(&self, _client: &mut Client, _sender_id: ObjectId) -> Result<()> {
+		// we want the upload to complete before we give the image to bevy
+		let buffer_option = self
+			.state
+			.lock()
+			.pending
+			.buffer
+			.as_ref()
+			.map(|b| b.buffer.clone());
+		if let Some(buffer) = buffer_option {
+			tokio::task::spawn_blocking(move || buffer.on_commit())
+				.await
+				.unwrap();
+		}
 		self.state.lock().apply();
 		self.state.lock().pending.frame_callbacks.clear();
 		let current_state = self.current_state();