From 81a741ad36c26a0f9a5cc24ffce774a9a4eda997 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@gmail.com>
Date: Tue, 15 Jul 2025 02:26:22 -0700
Subject: [PATCH] fix(wayland/dmabuf): more robust format handling

---
 src/wayland/dmabuf/feedback.rs | 15 ++++-----------
 src/wayland/dmabuf/mod.rs      | 21 +++++++++++----------
 2 files changed, 15 insertions(+), 21 deletions(-)

diff --git a/src/wayland/dmabuf/feedback.rs b/src/wayland/dmabuf/feedback.rs
index c1f99cb..7b0add7 100644
--- a/src/wayland/dmabuf/feedback.rs
+++ b/src/wayland/dmabuf/feedback.rs
@@ -1,6 +1,5 @@
 use super::Dmabuf;
 use crate::wayland::vulkano_data::VULKANO_CONTEXT;
-use drm_fourcc::DrmFourcc;
 use memfd::MemfdOptions;
 use std::{
 	io::Write,
@@ -23,8 +22,7 @@ impl DmabufFeedback {
 	pub async fn send_params(&self, client: &mut Client, sender_id: ObjectId) -> Result<()> {
 		let num_formats = self.0.formats.len();
 		// Send format table first
-		self.send_format_table(client, sender_id, &self.0.formats)
-			.await?;
+		self.send_format_table(client, sender_id).await?;
 
 		// Get the device information from Vulkan properties
 		let props = VULKANO_CONTEXT.get().unwrap().phys_dev.properties();
@@ -62,17 +60,12 @@ impl DmabufFeedback {
 		Ok(())
 	}
 
-	pub async fn send_format_table(
-		&self,
-		client: &mut Client,
-		sender_id: ObjectId,
-		formats: &[(DrmFourcc, u64)],
-	) -> Result<()> {
+	pub async fn send_format_table(&self, client: &mut Client, sender_id: ObjectId) -> Result<()> {
 		// Format + modifier pair (16 bytes):
 		// - format: u32
 		// - padding: 4 bytes
 		// - modifier: u64
-		let size = formats.len() as u32 * 16u32;
+		let size = self.0.formats.len() as u32 * 16u32;
 		// Create a temporary file for the format table
 		let mfd = MemfdOptions::default()
 			.create("stardustxr-format-table")
@@ -80,7 +73,7 @@ impl DmabufFeedback {
 
 		mfd.as_file().set_len(size as u64)?;
 
-		for (format, modifier) in formats {
+		for (format, modifier) in self.0.formats.iter() {
 			let format = format.clone() as u32;
 			// Write the format+modifier pair
 			mfd.as_file().write_all(&format.to_ne_bytes())?;
diff --git a/src/wayland/dmabuf/mod.rs b/src/wayland/dmabuf/mod.rs
index 4af6d55..f9ef324 100644
--- a/src/wayland/dmabuf/mod.rs
+++ b/src/wayland/dmabuf/mod.rs
@@ -2,6 +2,11 @@ pub mod buffer_backing;
 pub mod buffer_params;
 pub mod feedback;
 
+use super::{
+	util::ClientExt,
+	vulkano_data::{DMA_CAPABLE_FORMATS, VULKANO_CONTEXT},
+};
+use crate::core::registry::Registry;
 use bevy_dmabuf::{
 	format_mapping::{drm_fourcc_to_vk_format, vk_format_to_drm_fourcc},
 	wgpu_init::vulkan_to_wgpu,
@@ -9,6 +14,7 @@ use bevy_dmabuf::{
 use buffer_params::BufferParams;
 use drm_fourcc::DrmFourcc;
 use feedback::DmabufFeedback;
+use rustc_hash::FxHashSet;
 use waynest::{
 	server::{
 		Client, Dispatcher, Error, Result,
@@ -20,13 +26,6 @@ use waynest::{
 	wire::ObjectId,
 };
 
-use crate::core::registry::Registry;
-
-use super::{
-	util::ClientExt,
-	vulkano_data::{DMA_CAPABLE_FORMATS, VULKANO_CONTEXT},
-};
-
 /// Main DMA-BUF interface implementation
 ///
 /// This interface allows clients to create wl_buffers from DMA-BUFs.
@@ -46,14 +45,14 @@ pub struct Dmabuf {
 	// Track active buffer parameters objects by their ID
 	active_params: Registry<BufferParams>,
 	pub(self) version: u32,
-	pub(self) formats: Vec<(DrmFourcc, u64)>,
+	pub(self) formats: FxHashSet<(DrmFourcc, u64)>,
 }
 
 impl Dmabuf {
 	/// Create a new DMA-BUF interface instance
 	pub async fn new(client: &mut Client, id: ObjectId, version: u32) -> Result<Self> {
 		let vk = VULKANO_CONTEXT.wait();
-		let formats = DMA_CAPABLE_FORMATS
+		let formats: FxHashSet<(DrmFourcc, u64)> = DMA_CAPABLE_FORMATS
 			.iter()
 			.filter(|f| {
 				vk_format_to_drm_fourcc((**f).into())
@@ -76,6 +75,8 @@ impl Dmabuf {
 			.flat_map(|(f, mods)| mods.into_iter().map(move |modifier| (f, modifier)))
 			.collect();
 
+		dbg!(&formats);
+
 		let dmabuf = Self {
 			// formats: Mutex::new(formats),
 			active_params: Registry::new(),
@@ -83,7 +84,7 @@ impl Dmabuf {
 			formats,
 		};
 
-		if version > 3 {
+		if version < 3 {
 			for (format, _) in &dmabuf.formats {
 				dmabuf.format(client, id, *format as u32).await?;
 			}