wrgpgpu 0.1.5

#![doc = include_str!("../readme.md")]

use std::{marker::PhantomData, sync::Arc};

/// Bindings
pub mod bindings;

use bindings::{BindGroupData, BindGroups};

#[doc(inline)]
pub use bindings::{
	buffer::{StorageBufferBind, UniformBufferBind},
	texture::{RgbaStorageTextureBind, StorageTextureBind, TextureBind},
	Bind, BindGroup,
};
pub use image::RgbaImage;
pub use wgpu::include_wgsl;

/// A hardware GPU acceleration device
pub struct Device {
	pub device: Arc<wgpu::Device>,
	pub queue: Arc<wgpu::Queue>,
}

/// Arguments used to create shaders
pub struct ShaderArgs<'a> {
	pub label: &'static str,
	pub shader: wgpu::ShaderModuleDescriptor<'a>,
	pub entrypoint: &'a str,
}

/// A compute shader which has a specific set of bind groups
pub struct ComputeShader<B: BindGroups> {
	pub bind_group_layouts: Vec<wgpu::BindGroupLayout>,
	pub pipeline: wgpu::ComputePipeline,
	_bind_group: PhantomData<B>,
}

impl Device {
	/// Create a device from the "best" choice for high performance usage
	pub async fn auto_high_performance() -> Self {
		let instance = Arc::new(wgpu::Instance::new(&wgpu::InstanceDescriptor {
			backends: wgpu::Backends::VULKAN,
			#[cfg(debug_assertions)]
			flags: wgpu::InstanceFlags::debugging(),
			..Default::default()
		}));
		log::info!("uwu: {instance:?}");
		let adapter = Arc::new(
			instance
				.request_adapter(&wgpu::RequestAdapterOptions {
					power_preference: wgpu::PowerPreference::HighPerformance,
					compatible_surface: None,
					force_fallback_adapter: false,
				})
				.await
				.unwrap(),
		);
		let (device, queue) = adapter
			.request_device(
				&wgpu::DeviceDescriptor {
					required_features: wgpu::Features::empty(),
					required_limits: wgpu::Limits::default(),
					label: Some("wrgpgpu_device"),
					memory_hints: Default::default(),
				},
				None,
			)
			.await
			.unwrap();
		Self {
			device: Arc::new(device),
			queue: Arc::new(queue),
		}
	}
	/// Create a device from existing wgpu types
	pub fn from_wgpu(device: Arc<wgpu::Device>, queue: Arc<wgpu::Queue>) -> Self {
		Self {
			device,
			queue,
		}
	}

	/// Create a shader on the device
	pub fn create_shader<B: BindGroups>(&self, args: ShaderArgs<'_>) -> ComputeShader<B> {
		let bind_group_layouts = B::get_layouts(self);

		let layout = self
			.device
			.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
				label: Some(args.label),
				bind_group_layouts: &bind_group_layouts.iter().collect::<Vec<_>>(),
				push_constant_ranges: &[],
			});
		let pipeline = self
			.device
			.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
				label: Some(args.label),
				layout: Some(&layout),
				module: &self.device.create_shader_module(args.shader),
				entry_point: Some(args.entrypoint),
				cache: None,
				compilation_options: wgpu::PipelineCompilationOptions::default(),
			});

		ComputeShader {
			bind_group_layouts,
			pipeline,
			_bind_group: PhantomData,
		}
	}

	/// Create a bind group for use with a shader
	pub fn bind<B: BindGroupData>(&self, data: &B) -> BindGroup<B> {
		// TODO: cache bind group layouts, really important
		BindGroup::new(self, &BindGroup::<B>::get_layout(self), data)
	}

	/// Dispatch a compute shader pass
	pub fn dispatch<B: BindGroups>(
		&self,
		shader: &ComputeShader<B>,
		bindings: &B,
		workgroups: (u32, u32, u32),
	) {
		let mut encoder = self
			.device
			.create_command_encoder(&wgpu::CommandEncoderDescriptor {
				label: Some("Compute Shader Command Encoder"),
			});
		{
			let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
				label: Some("Compute Shader Pass"),
				timestamp_writes: None,
			});
			pass.set_pipeline(&shader.pipeline);
			for n in 0..B::LEN {
				pass.set_bind_group(n as u32, bindings.get_bind_group_n(n), &[]);
			}
			pass.dispatch_workgroups(workgroups.0, workgroups.1, workgroups.2);
		}
		self.queue.submit(std::iter::once(encoder.finish()));
	}

	pub fn is_complete(&self) -> bool {
		self.device.poll(wgpu::Maintain::Poll).is_queue_empty()
	}
}