Galactica/crates/render/src/gpustate/mod.rs

use anyhow::Result;
use bytemuck;
use cgmath::Point2;
use galactica_constants;

use rand::seq::SliceRandom;
use std::{iter, rc::Rc};
use wgpu::{self, BufferAddress};
use winit::{self, window::Window};

use crate::{
	content,
	globaluniform::{GlobalDataContent, GlobalUniform},
	pipeline::PipelineBuilder,
	starfield::Starfield,
	texturearray::TextureArray,
	vertexbuffer::{
		consts::{SPRITE_INDICES, SPRITE_VERTICES},
		types::{
			ObjectInstance, ParticleInstance, RadialBarInstance, StarfieldInstance, TexturedVertex,
			UiInstance,
		},
		BufferObject, VertexBuffer,
	},
	RenderState,
};

// Additional implementaitons for GPUState
mod hud;
mod world;

/// A high-level GPU wrapper. Consumes game state,
/// produces pretty pictures.
pub struct GPUState {
	/// The window to we draw on
	pub window: Window,

	/// The size of the window we draw on
	pub window_size: winit::dpi::PhysicalSize<u32>,

	device: wgpu::Device,
	config: wgpu::SurfaceConfiguration,
	surface: wgpu::Surface,
	queue: wgpu::Queue,

	window_aspect: f32,

	object_pipeline: wgpu::RenderPipeline,
	starfield_pipeline: wgpu::RenderPipeline,
	particle_pipeline: wgpu::RenderPipeline,
	ui_pipeline: wgpu::RenderPipeline,
	radialbar_pipeline: wgpu::RenderPipeline,

	starfield: Starfield,
	texture_array: TextureArray,
	global_uniform: GlobalUniform,
	vertex_buffers: VertexBuffers,
}

struct VertexBuffers {
	// Keeps track of length of each buffer
	pub object_counter: BufferAddress,
	pub ui_counter: BufferAddress,
	pub particle_counter: BufferAddress,
	pub radialbar_counter: BufferAddress,

	object: Rc<VertexBuffer>,
	starfield: Rc<VertexBuffer>,
	ui: Rc<VertexBuffer>,

	particle: Rc<VertexBuffer>,
	radialbar: Rc<VertexBuffer>,
}

/// Basic wgsl preprocesser
fn preprocess_shader(
	shader: &str,
	global_uniform: &GlobalUniform,
	global_uniform_group: u32,
) -> String {
	// Insert dynamically-generated global definitions
	let shader = shader.replace(
		"// INCLUDE: global uniform header",
		&global_uniform.shader_header(global_uniform_group),
	);

	// Insert common functions
	let shader = shader.replace(
		"// INCLUDE: animate.wgsl",
		&include_str!(concat!(
			env!("CARGO_MANIFEST_DIR"),
			"/shaders/include/",
			"animate.wgsl"
		)),
	);

	let shader = shader.replace(
		"// INCLUDE: anchor.wgsl",
		&include_str!(concat!(
			env!("CARGO_MANIFEST_DIR"),
			"/shaders/include/",
			"anchor.wgsl"
		)),
	);

	return shader;
}

impl GPUState {
	/// Make a new GPUState that draws on `window`
	pub async fn new(window: Window, ct: &content::Content) -> Result<Self> {
		let window_size = window.inner_size();
		let window_aspect = window_size.width as f32 / window_size.height as f32;

		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: window_size.width,
				height: window_size.height,
				present_mode: surface_caps.present_modes[0],
				alpha_mode: surface_caps.alpha_modes[0],
				view_formats: vec![],
			};

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

		let vertex_buffers = VertexBuffers {
			object_counter: 0,
			ui_counter: 0,
			particle_counter: 0,
			radialbar_counter: 0,

			object: Rc::new(VertexBuffer::new::<TexturedVertex, ObjectInstance>(
				"object",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				galactica_constants::OBJECT_SPRITE_INSTANCE_LIMIT,
			)),

			starfield: Rc::new(VertexBuffer::new::<TexturedVertex, StarfieldInstance>(
				"starfield",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				galactica_constants::STARFIELD_SPRITE_INSTANCE_LIMIT,
			)),

			ui: Rc::new(VertexBuffer::new::<TexturedVertex, UiInstance>(
				"ui",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				galactica_constants::UI_SPRITE_INSTANCE_LIMIT,
			)),

			particle: Rc::new(VertexBuffer::new::<TexturedVertex, ParticleInstance>(
				"particle",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				galactica_constants::PARTICLE_SPRITE_INSTANCE_LIMIT,
			)),

			radialbar: Rc::new(VertexBuffer::new::<TexturedVertex, RadialBarInstance>(
				"radial bar",
				&device,
				Some(SPRITE_VERTICES),
				Some(SPRITE_INDICES),
				10,
			)),
		};

		// Load uniforms
		let global_uniform = GlobalUniform::new(&device);
		let texture_array = TextureArray::new(&device, &queue, ct)?;

		// Make sure these match the indices in each shader
		let bind_group_layouts = &[
			&texture_array.bind_group_layout,
			&global_uniform.bind_group_layout,
		];

		// Create render pipelines
		let object_pipeline = PipelineBuilder::new("object", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"object.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.object)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let starfield_pipeline = PipelineBuilder::new("starfield", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"starfield.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.starfield)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let ui_pipeline = PipelineBuilder::new("ui", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/shaders/", "ui.wgsl")),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.ui)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let particle_pipeline = PipelineBuilder::new("particle", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"particle.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.particle)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let radialbar_pipeline = PipelineBuilder::new("radialbar", &device)
			.set_shader(&preprocess_shader(
				&include_str!(concat!(
					env!("CARGO_MANIFEST_DIR"),
					"/shaders/",
					"radialbar.wgsl"
				)),
				&global_uniform,
				1,
			))
			.set_format(config.format)
			.set_triangle(true)
			.set_vertex_buffer(&vertex_buffers.radialbar)
			.set_bind_group_layouts(bind_group_layouts)
			.build();

		let mut starfield = Starfield::new();
		starfield.regenerate();

		return Ok(Self {
			device,
			config,
			surface,
			queue,

			window,
			window_size,
			window_aspect,

			object_pipeline,
			starfield_pipeline,
			ui_pipeline,
			particle_pipeline,
			radialbar_pipeline,

			starfield,
			texture_array,
			global_uniform,
			vertex_buffers,
		});
	}

	/// Get the window this GPUState is attached to
	pub fn window(&self) -> &Window {
		&self.window
	}

	/// Update window size.
	/// This should be called whenever our window is resized.
	pub fn resize(&mut self) {
		let new_size = self.window.inner_size();
		if new_size.width > 0 && new_size.height > 0 {
			self.window_size = new_size;
			self.window_aspect = new_size.width as f32 / new_size.height as f32;
			self.config.width = new_size.width;
			self.config.height = new_size.height;
			self.surface.configure(&self.device, &self.config);
		}
		self.update_starfield_buffer()
	}

	/// Entrypoint for all vertex buffer builders
	pub(super) fn update_all_buffers(&mut self, state: &RenderState) {
		// Game coordinates (relative to camera) of ne and sw corners of screen.
		// Used to skip off-screen sprites.
		let clip_ne = Point2::from((-self.window_aspect, 1.0)) * state.camera_zoom;
		let clip_sw = Point2::from((self.window_aspect, -1.0)) * state.camera_zoom;

		// TODO: sorting. We don't need to sort ships, but we do need to sort system objects by z-level
		// (which we don't yet draw)
		// that should probably be done in iter_system().

		// Order matters, it determines what is drawn on top.
		// The order inside ships and projectiles doesn't matter,
		// but ships should always be under projectiles.
		self.world_push_system(state, (clip_ne, clip_sw));
		self.world_push_ship(state, (clip_ne, clip_sw));
		self.world_push_projectile(state, (clip_ne, clip_sw));

		self.hud_add_radar(state);
		self.hud_add_status(state);
	}

	/// Make a StarfieldInstance for each star that needs to be drawn.
	/// Will panic if STARFIELD_INSTANCE_LIMIT is exceeded.
	///
	/// Starfield data rarely changes, so this is called only when it's needed.
	pub fn update_starfield_buffer(&mut self) {
		self.queue.write_buffer(
			&self.vertex_buffers.starfield.instances,
			0,
			bytemuck::cast_slice(&self.starfield.make_instances(self.window_aspect)),
		);
	}

	/// Initialize the rendering engine
	pub fn init(&mut self) {
		// Update global values
		self.queue.write_buffer(
			&self.global_uniform.atlas_buffer,
			0,
			bytemuck::cast_slice(&[self.texture_array.image_locations]),
		);
		self.queue.write_buffer(
			&self.global_uniform.sprite_buffer,
			0,
			bytemuck::cast_slice(&[self.texture_array.sprite_data]),
		);

		self.update_starfield_buffer();
	}

	/// Main render function. Draws sprites on a window.
	pub fn render(&mut self, state: RenderState) -> Result<(), wgpu::SurfaceError> {
		let output = self.surface.get_current_texture()?;
		let view = output.texture.create_view(&Default::default());

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

		let mut render_pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
			label: Some("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,
		});

		self.vertex_buffers.object_counter = 0;
		self.vertex_buffers.ui_counter = 0;
		self.vertex_buffers.radialbar_counter = 0;
		// Don't reset particle counter, it's special

		let s = state.content.get_starfield_handle();

		// Update global values
		self.queue.write_buffer(
			&self.global_uniform.data_buffer,
			0,
			bytemuck::cast_slice(&[GlobalDataContent {
				camera_position: state.camera_pos.into(),
				camera_zoom: [state.camera_zoom, 0.0],
				camera_zoom_limits: [galactica_constants::ZOOM_MIN, galactica_constants::ZOOM_MAX],
				window_size: [
					self.window_size.width as f32,
					self.window_size.height as f32,
				],
				window_scale: [self.window.scale_factor() as f32, 0.0],
				window_aspect: [self.window_aspect, 0.0],
				starfield_sprite: [s.get_index(), 0],
				starfield_tile_size: [galactica_constants::STARFIELD_SIZE as f32, 0.0],
				starfield_size_limits: [
					galactica_constants::STARFIELD_SIZE_MIN,
					galactica_constants::STARFIELD_SIZE_MAX,
				],
				current_time: [state.current_time, 0.0],
			}]),
		);

		// Write all new particles to GPU buffer
		state.particles.shuffle(&mut rand::thread_rng());
		for i in state.particles.iter() {
			self.queue.write_buffer(
				&self.vertex_buffers.particle.instances,
				ParticleInstance::SIZE * self.vertex_buffers.particle_counter,
				bytemuck::cast_slice(&[ParticleInstance {
					position: [i.pos.x, i.pos.y],
					velocity: i.velocity.into(),
					angle: i.angle.0,
					angvel: i.angvel.0,
					size: i.size,
					sprite_index: i.sprite.get_index(),
					created: state.current_time,
					expires: state.current_time + i.lifetime,
					fade: i.fade,
				}]),
			);
			self.vertex_buffers.particle_counter += 1;
			if self.vertex_buffers.particle_counter
				== galactica_constants::PARTICLE_SPRITE_INSTANCE_LIMIT
			{
				self.vertex_buffers.particle_counter = 0;
			}
		}
		state.particles.clear();

		// Create sprite instances
		self.update_all_buffers(&state);

		// These should match the indices in each shader,
		// and should each have a corresponding bind group layout.
		render_pass.set_bind_group(0, &self.texture_array.bind_group, &[]);
		render_pass.set_bind_group(1, &self.global_uniform.bind_group, &[]);

		// Starfield pipeline
		self.vertex_buffers.starfield.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.starfield_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.starfield.instance_count,
		);

		// Sprite pipeline
		self.vertex_buffers.object.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.object_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.vertex_buffers.object_counter as _,
		);

		// Particle pipeline
		self.vertex_buffers.particle.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.particle_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..galactica_constants::PARTICLE_SPRITE_INSTANCE_LIMIT as _,
		);

		// Ui pipeline
		self.vertex_buffers.ui.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.ui_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.vertex_buffers.ui_counter as _,
		);

		// Radial progress bars
		// TODO: do we need to do this every time?
		self.vertex_buffers.radialbar.set_in_pass(&mut render_pass);
		render_pass.set_pipeline(&self.radialbar_pipeline);
		render_pass.draw_indexed(
			0..SPRITE_INDICES.len() as u32,
			0,
			0..self.vertex_buffers.radialbar_counter 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(())
	}
}