Galactica/src/render/gpustate.rs

use anyhow::Result;
use bytemuck;
use cgmath::{EuclideanSpace, Point2};
use std::iter;
use wgpu::{self, util::DeviceExt};
use winit::{self, window::Window};

use crate::{Camera, Sprite};

use super::{
	bufferdata::{PlainVertex, SpriteInstance, StarInstance, TexturedVertex},
	pipeline::PipelineBuilder,
	texturearray::TextureArray,
	util::Transform,
	SPRITE_INDICES, SPRITE_VERTICES, STARFIELD_VERTICES,
};

pub struct GPUState {
	device: wgpu::Device,
	config: wgpu::SurfaceConfiguration,
	surface: wgpu::Surface,
	queue: wgpu::Queue,

	pub window: Window,
	pub size: winit::dpi::PhysicalSize<u32>,

	sprite_pipeline: wgpu::RenderPipeline,
	starfield_pipeline: wgpu::RenderPipeline,

	star_vertex_buffer: wgpu::Buffer,
	sprite_vertex_buffer: wgpu::Buffer,
	sprite_index_buffer: wgpu::Buffer,
	texture_array: TextureArray,

	sprite_instance_buffer: wgpu::Buffer,
	star_instance_buffer: wgpu::Buffer,
}

impl GPUState {
	// We can draw at most this many sprites on the screen.
	// TODO: compile-time option
	pub const SPRITE_LIMIT: u64 = 100;
	pub const STAR_LIMIT: u64 = 100;

	pub async fn new(window: Window) -> Result<Self> {
		let size = window.inner_size();

		let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
			backends: wgpu::Backends::all(),
			..Default::default()
		});

		let surface = unsafe { instance.create_surface(&window) }.unwrap();

		// Basic setup
		let device;
		let queue;
		let config;

		{
			let adapter = instance
				.request_adapter(&wgpu::RequestAdapterOptions {
					power_preference: wgpu::PowerPreference::default(),
					compatible_surface: Some(&surface),
					force_fallback_adapter: false,
				})
				.await
				.unwrap();

			(device, queue) = adapter
				.request_device(
					&wgpu::DeviceDescriptor {
						features: wgpu::Features::TEXTURE_BINDING_ARRAY | wgpu::Features::SAMPLED_TEXTURE_AND_STORAGE_BUFFER_ARRAY_NON_UNIFORM_INDEXING,
						// We may need limits if we compile for wasm
						limits: wgpu::Limits::default(),
						label: Some("gpu device"),
					},
					None,
				)
				.await
				.unwrap();

			// Assume sRGB
			let surface_caps = surface.get_capabilities(&adapter);
			let surface_format = surface_caps
				.formats
				.iter()
				.copied()
				.filter(|f| f.is_srgb())
				.filter(|f| f.has_stencil_aspect())
				.next()
				.unwrap_or(surface_caps.formats[0]);

			config = wgpu::SurfaceConfiguration {
				usage: wgpu::TextureUsages::RENDER_ATTACHMENT,
				format: surface_format,
				width: size.width,
				height: size.height,
				present_mode: surface_caps.present_modes[0],
				alpha_mode: surface_caps.alpha_modes[0],
				view_formats: vec![],
			};

			surface.configure(&device, &config);
		}

		// Load textures
		let texture_array = TextureArray::new(&device, &queue)?;

		// Render pipelines
		let sprite_pipeline = PipelineBuilder::new("sprite", &device)
			.set_shader(include_str!(concat!(
				env!("CARGO_MANIFEST_DIR"),
				"/src/render/shaders/",
				"sprite.wgsl"
			)))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffers(&[TexturedVertex::desc(), SpriteInstance::desc()])
			.set_bind_group_layouts(&[&texture_array.bind_group_layout])
			.build();

		let starfield_pipeline = PipelineBuilder::new("starfield", &device)
			.set_shader(include_str!(concat!(
				env!("CARGO_MANIFEST_DIR"),
				"/src/render/shaders/",
				"starfield.wgsl"
			)))
			.set_format(config.format)
			.set_triangle(false)
			.set_vertex_buffers(&[PlainVertex::desc(), StarInstance::desc()])
			.build();

		let sprite_vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: Some("vertex buffer"),
			contents: bytemuck::cast_slice(SPRITE_VERTICES),
			usage: wgpu::BufferUsages::VERTEX,
		});

		let sprite_index_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: Some("vertex index buffer"),
			contents: bytemuck::cast_slice(SPRITE_INDICES),
			usage: wgpu::BufferUsages::INDEX,
		});

		let star_vertex_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
			label: Some("vertex buffer"),
			contents: bytemuck::cast_slice(STARFIELD_VERTICES),
			usage: wgpu::BufferUsages::VERTEX,
		});

		let sprite_instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
			label: Some("sprite instance buffer"),
			usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
			size: SpriteInstance::SIZE * Self::SPRITE_LIMIT,
			mapped_at_creation: false,
		});

		let star_instance_buffer = device.create_buffer(&wgpu::BufferDescriptor {
			label: Some("star instance buffer"),
			usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
			size: StarInstance::SIZE * Self::STAR_LIMIT,
			mapped_at_creation: false,
		});

		return Ok(Self {
			surface,
			device,
			queue,
			config,
			size,
			window,
			sprite_pipeline,
			starfield_pipeline,
			sprite_vertex_buffer,
			sprite_index_buffer,
			star_vertex_buffer,
			sprite_instance_buffer,
			star_instance_buffer,
			texture_array,
		});
	}

	pub fn window(&self) -> &Window {
		&self.window
	}

	pub fn resize(&mut self, new_size: winit::dpi::PhysicalSize<u32>) {
		if new_size.width > 0 && new_size.height > 0 {
			self.size = new_size;
			self.config.width = new_size.width;
			self.config.height = new_size.height;
			self.surface.configure(&self.device, &self.config);
		}
	}

	pub fn update(&mut self) {}

	pub fn render(
		&mut self,
		sprites: &Vec<Sprite>,
		camera: Camera,
	) -> Result<(), wgpu::SurfaceError> {
		let output = self.surface.get_current_texture()?;
		let view = output
			.texture
			.create_view(&wgpu::TextureViewDescriptor::default());

		let mut encoder = self
			.device
			.create_command_encoder(&wgpu::CommandEncoderDescriptor {
				label: Some("sprite render encoder"),
			});

		let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			label: Some("sprite render pass"),

			color_attachments: &[Some(wgpu::RenderPassColorAttachment {
				view: &view,
				resolve_target: None,
				ops: wgpu::Operations {
					load: wgpu::LoadOp::Clear(wgpu::Color {
						r: 0.0,
						g: 0.0,
						b: 0.0,
						a: 1.0,
					}),
					store: wgpu::StoreOp::Store,
				},
			})],
			depth_stencil_attachment: None,
			occlusion_query_set: None,
			timestamp_writes: None,
		});

		// Correct for screen aspect ratio
		// (it may not be square!)
		let screen_aspect = self.size.width as f32 / self.size.height as f32;

		// Game coordinates (relative to camera) of ne and sw corners of screen.
		// Used to skip off-screen sprites.
		let clip_ne = Point2::from((-1.0, 1.0)) * camera.zoom;
		let clip_sw = Point2::from((1.0, -1.0)) * camera.zoom;

		let mut instances: Vec<SpriteInstance> = Vec::new();

		for s in sprites {
			let pos = s.post_parallax_position(camera) - camera.pos.to_vec();
			let texture = self.texture_array.get_texture(&s.name[..]);

			// Game dimensions of this sprite post-scale.
			//
			// We only need height / 2 to check if we're on the screen,
			// but we omit the division.
			// This gives us a small margin, and lets us re-use the value
			// without an extra multiply.
			let height = s.height * s.scale;
			let width = height * texture.aspect;

			// Don't draw (or compute matrices for)
			// sprites that are off the screen
			if pos.x < clip_ne.x - width
				|| pos.y > clip_ne.y + height
				|| pos.x > clip_sw.x + width
				|| pos.y < clip_sw.y - height
			{
				continue;
			}

			// Compute the values we need to draw
			// scale: combines texture scale and zoom scale.
			// screen_pos: position of this sprite in screen coordinates
			//
			// We can't use screen_pos to exclude off-screen sprites because
			// it can't account for height and width.
			let scale = height / camera.zoom;
			let screen_pos: Point2<f32> = pos / camera.zoom;

			instances.push(SpriteInstance {
				transform: Transform {
					pos: screen_pos,
					aspect: texture.aspect,
					screen_aspect,
					rotate: s.angle,
					scale,
				}
				.to_matrix()
				.into(),
				texture_index: texture.index,
			})
		}

		// Enforce sprite limit
		if sprites.len() as u64 > Self::SPRITE_LIMIT {
			// TODO: no panic, handle this better.
			panic!("Sprite limit exceeded!")
		}

		// Write new sprite data to buffer
		self.queue.write_buffer(
			&self.sprite_instance_buffer,
			0,
			bytemuck::cast_slice(&instances),
		);

		render_pass.set_bind_group(0, &self.texture_array.bind_group, &[]);

		let nstances: Vec<StarInstance> = (-10..10)
			.map(|x| StarInstance {
				transform: cgmath::Matrix4::from_translation(cgmath::Vector3 {
					x: x as f32 / 10.0,
					y: 0.0,
					z: 0.0,
				})
				.into(),
			})
			.collect();

		self.queue.write_buffer(
			&self.star_instance_buffer,
			0,
			bytemuck::cast_slice(&nstances),
		);

		// Starfield pipeline
		render_pass.set_vertex_buffer(0, self.star_vertex_buffer.slice(..));
		render_pass.set_vertex_buffer(1, self.star_instance_buffer.slice(..));
		render_pass.set_pipeline(&self.starfield_pipeline);
		render_pass.draw(0..STARFIELD_VERTICES.len() as _, 0..nstances.len() as _);

		// Sprite pipeline
		render_pass.set_vertex_buffer(0, self.sprite_vertex_buffer.slice(..));
		render_pass.set_vertex_buffer(1, self.sprite_instance_buffer.slice(..));
		render_pass.set_index_buffer(
			self.sprite_index_buffer.slice(..),
			wgpu::IndexFormat::Uint16,
		);
		render_pass.set_pipeline(&self.sprite_pipeline);
		render_pass.draw_indexed(0..SPRITE_INDICES.len() as u32, 0, 0..instances.len() as _);

		// begin_render_pass borrows encoder mutably, so we can't call finish()
		// without dropping this variable.
		drop(render_pass);

		self.queue.submit(iter::once(encoder.finish()));
		output.present();

		Ok(())
	}
}